Implementing PyTorch IterableDataset in Python

This article explains how to implement custom stream-based datasets in Python using PyTorch's torch.utils.data.IterableDataset. You will learn the fundamental structure of an iterable dataset, how to define the core iteration logic, and how to properly configure multi-process data loading to prevent data duplication across worker threads.

When to Use IterableDataset

PyTorch provides two types of datasets: map-style (Dataset) and iterable-style (IterableDataset). An IterableDataset is best suited for scenarios where random access via index (__getitem__) is either impossible or expensive. Common use cases include:

Basic Implementation

To create an iterable dataset, subclass torch.utils.data.IterableDataset and implement the __iter__() method, which must return a Python iterator or generator yielding individual samples.

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

class NumberStreamDataset(IterableDataset):
    def __init__(self, start: int, end: int):
        super().__init__()
        self.start = start
        self.end = end

    def __iter__(self):
        for value in range(self.start, self.end):
            yield torch.tensor(value, dtype=torch.float32)

# Usage
dataset = NumberStreamDataset(start=0, end=10)
loader = DataLoader(dataset, batch_size=2)

for batch in loader:
    print(batch)

Handling Multi-Process Data Loading

When using torch.utils.data.DataLoader with num_workers > 0, PyTorch creates duplicate worker processes. By default, each worker runs an identical copy of the dataset's __iter__() method, causing duplicate samples.

To split the workload, you must identify the worker process inside __iter__() using torch.utils.data.get_worker_info() or define a custom worker_init_fn.

Splitting Data Inside __iter__

The most common approach divides the stream according to the worker ID:

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

class MultiWorkerDataset(IterableDataset):
    def __init__(self, start: int, end: int):
        super().__init__()
        self.start = start
        self.end = end

    def __iter__(self):
        worker_info = get_worker_info()

        if worker_info is None:
            # Single-process data loading
            iter_start = self.start
            iter_end = self.end
        else:
            # Multi-process data loading: partition workload across workers
            total_items = self.end - self.start
            per_worker = int(math.ceil(total_items / float(worker_info.num_workers)))
            worker_id = worker_info.id
            
            iter_start = self.start + worker_id * per_worker
            iter_end = min(iter_start + per_worker, self.end)

        for value in range(iter_start, iter_end):
            yield torch.tensor(value, dtype=torch.float32)

# Usage with multiple workers
dataset = MultiWorkerDataset(start=0, end=20)
loader = DataLoader(dataset, batch_size=4, num_workers=2)

for batch in loader:
    print(batch)

Key Considerations