Iterator & Generator Protocol¶

Overview¶

Iterators and generators are fundamental to ML pipelines:

DataLoaders use iterators to stream data
Generators are memory-efficient (don't load entire dataset)
Yield enables lazy evaluation for large datasets
JAX transformations leverage generator composition

This is why PyTorch DataLoaders don't load all data into memory upfront.

The Iterator Protocol¶

Core Concept¶

An iterator is an object that implements two methods:

class Iterator:
 def __iter__(self):
 """Return self (iterator object)."""
 return self

 def __next__(self):
 """Return next value, raise StopIteration when done."""
 if self.has_next():
 return self.get_next()
 raise StopIteration

Simple Example¶

# Built-in iterator
data = [1, 2, 3]
iterator = iter(data) # Calls data.__iter__()

print(next(iterator)) # 1, calls iterator.__next__()
print(next(iterator)) # 2
print(next(iterator)) # 3
print(next(iterator)) # StopIteration exception

# For loop uses iterator protocol
for value in data: # Implicitly calls iter() and __next__()
 print(value)

Custom Iterator¶

class RangeIterator:
 """Custom iterator (like range)."""
 def __init__(self, start, end):
 self.current = start
 self.end = end

 def __iter__(self):
 return self

 def __next__(self):
 if self.current < self.end:
 value = self.current
 self.current += 1
 return value
 raise StopIteration

# Usage
for value in RangeIterator(0, 5):
 print(value) # 0, 1, 2, 3, 4

Generators: Simplified Iterators¶

Core Concept¶

A generator is a function that yields values instead of returning:

# Generator function (uses yield)
def count_to_n(n):
 """Generator that yields 0 to n-1."""
 i = 0
 while i < n:
 yield i # Pause here, return value, resume when next() called
 i += 1

# Generator object (not executed yet)
gen = count_to_n(3)
print(type(gen)) # <class 'generator'>

# Iteration
print(next(gen)) # 0 (function resumes after first yield)
print(next(gen)) # 1
print(next(gen)) # 2
print(next(gen)) # StopIteration

Why Generators Matter for ML¶

# WITHOUT generators
def load_all_data():
 data = []
 for line in open("huge_dataset.txt"):
 data.append(process(line)) # Stores everything in RAM
 return data

all_data = load_all_data() # Uses gigabytes of RAM!

# WITH generators
def load_data_lazy():
 for line in open("huge_dataset.txt"):
 yield process(line) # Only one sample in RAM at a time

for sample in load_data_lazy():
 # Process one sample at a time
 train_model(sample)

Memory Comparison:

All-at-once: Load 1B samples × 1KB each = 1TB RAM (impossible!)
Generator: Load 1 sample at a time = 1KB RAM (feasible!)

Generator Expressions¶

Lazy versions of list comprehensions:

# List comprehension
squares_list = [x**2 for x in range(1000000)] # Allocates huge list

# Generator expression
squares_gen = (x**2 for x in range(1000000)) # Tiny object

# Iterate through generator
for square in squares_gen:
 print(square)

# Or convert to list if needed (but then defeats purpose)
squares_list = list(squares_gen) # Now it's big

Chaining Generators¶

# Multiple generators compose efficiently
data = range(1000000)

# Don't do this (allocates intermediate lists)
filtered = [x for x in data if x % 2 == 0]
squared = [x**2 for x in filtered]
top_100 = squared[:100]

# Do this (composable, lazy)
filtered = (x for x in data if x % 2 == 0)
squared = (x**2 for x in filtered)
top_100 = list(itertools.islice(squared, 100))

PyTorch DataLoader: Iterator Pattern¶

How DataLoader Uses Iterators¶

import torch
from torch.utils.data import DataLoader, Dataset

class MyDataset(Dataset):
 """Custom dataset."""
 def __len__(self):
 return 1000

 def __getitem__(self, idx):
 # Called by DataLoader to fetch individual samples
 return torch.randn(3, 224, 224), idx % 10

# Create iterator
dataset = MyDataset()
loader = DataLoader(
 dataset,
 batch_size=32,
 shuffle=True,
 num_workers=4 # Load batches in parallel
)

# DataLoader is iterable
for batch_images, batch_labels in loader: # Uses iterator protocol
 print(batch_images.shape) # (32, 3, 224, 224)
 # Process batch

Multiple Epochs¶

model = torch.nn.Linear(10, 5)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for epoch in range(3):
 # Each epoch creates new iterator over same dataset
 for batch_x, batch_y in loader:
 # Forward pass
 pred = model(batch_x)
 loss = torch.nn.functional.mse_loss(pred, batch_y)

 # Backward pass
 optimizer.zero_grad()
 loss.backward()
 optimizer.step()

 print(f"Epoch {epoch} done")

Key insight: DataLoader creates a fresh iterator each epoch, respecting shuffle order.

Advanced: Stateful Generators¶

Generators can maintain state:

# Generator with state
def batch_generator(data, batch_size):
 """Yield batches of data."""
 for i in range(0, len(data), batch_size):
 batch = data[i:i+batch_size]
 yield batch
 # State (i) is preserved between yields

# Usage
batches = batch_generator([1, 2, 3, 4, 5], batch_size=2)
print(next(batches)) # [1, 2]
print(next(batches)) # [3, 4]
print(next(batches)) # [5]
print(next(batches)) # StopIteration

Generator with Initialization¶

def infinite_loader(data):
 """Cycle through data infinitely (for validation)."""
 while True:
 for sample in data:
 yield sample

# Usage in validation loop
gen = infinite_loader(val_dataset)
for _ in range(num_validation_batches):
 batch = next(gen)
 validate(model, batch)

Important Generator Features¶

send() Method: Bidirectional Communication¶

def echo_generator():
 """Receive values via send()."""
 while True:
 value = yield # Wait for value to be sent
 print(f"Received: {value}")

gen = echo_generator()
next(gen) # Prime the generator (move to first yield)
gen.send("Hello") # Send value into generator
gen.send("World")
# Output:
# Received
# Received

ML Application: Feedback-driven data generation

def curriculum_generator(all_data):
 """Adjust difficulty based on model performance."""
 easy_data = [d for d in all_data if d['difficulty'] < 5]
 hard_data = [d for d in all_data if d['difficulty'] >= 5]

 current = easy_data

 while True:
 # Yield a batch
 batch = yield current

 # Receive feedback on model performance
 if batch.get('avg_loss') < 0.1:
 # Model is doing well, increase difficulty
 current = hard_data
 else:
 # Keep easy examples
 current = easy_data

Nested Generators with yield from¶

# yield from (Python 3.3+)
def flatten(nested_list):
 """Flatten nested lists using yield from."""
 for item in nested_list:
 if isinstance(item, list):
 yield from flatten(item) # Delegate to nested generator
 else:
 yield item

# Usage
data = [1, [2, [3, 4]], 5]
for value in flatten(data):
 print(value) # 1, 2, 3, 4, 5

# Practical example
def iter_dataset(root_dir):
 """Recursively iterate all samples."""
 for item in os.listdir(root_dir):
 path = os.path.join(root_dir, item)
 if os.path.isdir(path):
 yield from iter_dataset(path) # Recurse
 else:
 yield path # Leaf node

Itertools: Functional Iterator Tools¶

Common Patterns¶

import itertools

# chain
data1 = [1, 2, 3]
data2 = [4, 5, 6]
combined = itertools.chain(data1, data2)

# islice
first_100 = itertools.islice(all_data, 100)

# cycle
repeated = itertools.cycle([1, 2, 3])

# zip_longest
a = [1, 2]
b = [3, 4, 5]
zipped = itertools.zip_longest(a, b, fillvalue=0) # [(1,3), (2,4), (0,5)]

# groupby
sorted_data = [1, 1, 1, 2, 2, 3, 3, 3, 3]
for key, group in itertools.groupby(sorted_data):
 print(f"{key}: {list(group)}")

Practical DataLoader Example¶

from itertools import chain, islice

# Combine train and validation iterators
train_loader = DataLoader(train_dataset, batch_size=32)
val_loader = DataLoader(val_dataset, batch_size=32)

# First train 100 batches, then validate
combined = chain(islice(train_loader, 100), val_loader)

for batch in combined:
 if batch in train_loader:
 train_step(batch)
 else:
 val_step(batch)

Performance: Lists vs Generators¶

import sys

# List
list_comp = [x**2 for x in range(1000000)]
print(sys.getsizeof(list_comp)) # ~8MB

# Generator
gen_expr = (x**2 for x in range(1000000))
print(sys.getsizeof(gen_expr)) # ~136 bytes

# Time to create
import time

start = time.time()
list_comp = [x**2 for x in range(10000000)]
print(f"List: {time.time() - start:.3f}s") # ~0.5s

start = time.time()
gen_expr = (x**2 for x in range(10000000))
print(f"Generator: {time.time() - start:.3f}s") # ~0.0s (just creates object)

For ML: Use generators for datasets, lists only when needed.

The Iterable vs Iterator Distinction¶

# Iterable
data = [1, 2, 3]
iterator1 = iter(data)
iterator2 = iter(data)
# iterator1 and iterator2 are different objects!
# Each starts from beginning

# Iterator
iterator = iter([1, 2, 3])
print(iter(iterator) is iterator) # True! Iterator's __iter__ returns self
print(next(iterator)) # 1
print(next(iterator)) # 2

Common Pitfalls¶

Pitfall 1: Exhausted Generators¶

gen = (x**2 for x in range(5))

# First iteration
for x in gen:
 print(x)

# Second iteration doesn't work!
for x in gen:
 print(x) # Prints nothing - generator is exhausted

# Solution
gen = (x**2 for x in range(5))

Pitfall 2: Closing Generator¶

def my_gen():
 try:
 yield 1
 yield 2
 finally:
 print("Cleaning up...")

gen = my_gen()
next(gen) # 1
gen.close() # Raises GeneratorExit, triggers finally block
# Output

# Practical
def file_line_generator(filename):
 with open(filename) as f:
 try:
 for line in f:
 yield line.strip()
 finally:
 print("File closed") # Always runs

Summary¶

Iterators = Objects with __iter__() and __next__()
Generators = Functions with yield (simplified iterators)
Generator expressions = Lazy list comprehensions
DataLoader = Iterator pattern for batching
Memory efficiency = Process one sample at a time
Composition = Chain multiple generators efficiently

-

02 Data Structures Essentials - Lists and batching
[04 Context Managers & Resource Management](/05-py3/01-fundamentals/(04-context-managers-resource-management/) - Resource management with generators
00 Readme - Functional composition
[04 Profiling & Performance Analysis](/05-py3/09-bytecode-and-execution/(04-profiling-performance-analysis/) - Performance analysis

Iterator & Generator Protocol¶

Overview¶

The Iterator Protocol¶

Core Concept¶

Simple Example¶

Custom Iterator¶

Generators: Simplified Iterators¶

Core Concept¶

Why Generators Matter for ML¶

Generator Expressions¶

Chaining Generators¶

PyTorch DataLoader: Iterator Pattern¶

How DataLoader Uses Iterators¶

Multiple Epochs¶

Advanced: Stateful Generators¶

Generator with Initialization¶

Important Generator Features¶

send() Method: Bidirectional Communication¶

Nested Generators with yield from¶

Itertools: Functional Iterator Tools¶

Common Patterns¶

Practical DataLoader Example¶

Performance: Lists vs Generators¶

The Iterable vs Iterator Distinction¶

Common Pitfalls¶

Pitfall 1: Exhausted Generators¶

Pitfall 2: Closing Generator¶

Summary¶

Related Topics¶