🎓 A1Learner
🔍

PyTorch Iterable Datasets

An IterableDataset streams rows one after another instead of serving them by index. The data arrives as a flow, like readings from a sensor.

Pulling rows by hand

iter() starts the stream. Each next() takes the following row.

stream = iter(CounterStream())
print(next(stream))
print(next(stream))
print(next(stream))

A stream has no index

Asking for row 0 by index fails. A stream can only move forward.

CounterStream()[0]

Stopping the stream

A streamed dataset never holds all rows at once. It makes each row as the reader asks, and the reader stops it.

for step, batch in enumerate(loader, start=1):
    print(f"Batch {step}:", batch.tolist())
    if step == 5:
        break

Example

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


class CounterStream(IterableDataset):
    def __iter__(self):
        value = 0
        while True:
            yield torch.tensor([value, value * value])
            value += 1


loader = DataLoader(CounterStream(), batch_size=4)

for step, batch in enumerate(loader, start=1):
    print(f"Batch {step}:", batch.tolist())
    if step == 5:
        break

print("Stopped after 5 batches")

Key points

  • IterableDataset: a dataset that yields rows in order from __iter__
  • Use streaming for endless data or data too large to store
  • Stop an endless stream with break, or it runs forever

In this example, CounterStream is an IterableDataset with an endless __iter__. It yields a pair of a counter and its square, forever. A DataLoader with batch_size 4 groups the stream into batches. The loop prints each batch and counts the steps. After batch 5, break stops the loop, and the program ends. The stream itself never ends. The break is what makes this program terminate.

Note

A DataLoader cannot shuffle an IterableDataset, because shuffling needs the index access a stream does not have.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Which dataset fits a sensor feed that never ends?