Custom Collate Functions and IterableDataset

When samples don't stack cleanly

The default collate function stacks every field of each sample into a tensor. That works when all your samples have the same shape — image batches, fixed-feature regression. It breaks when sequences have variable length, when samples carry per-item metadata, or when you want to batch in some custom way.

Custom collate

A collate function takes List[Sample] and returns a batch. Most common variant: pad variable-length sequences to the longest in the batch.

IterableDataset — streaming data

For datasets that are too big to fit on disk (giant text corpora, network-streamed data, online sensor feeds), implement IterableDataset instead of Dataset. You only define __iter__(self); PyTorch will iterate batches from your iterator.

The catch with multi-worker IterableDatasets: each worker receives the full iterator unless you split it manually. Use torch.utils.data.get_worker_info() inside __iter__ to slice the stream by worker id.

Code

Pad-to-longest collate for variable-length sequences·python

import torch
from torch.utils.data import DataLoader, Dataset
from torch.nn.utils.rnn import pad_sequence

class VarSeqDataset(Dataset):
    def __init__(self):
        self.seqs = [torch.randint(1, 100, (torch.randint(5, 25, (1,)).item(),)) for _ in range(64)]
        self.labels = torch.randint(0, 2, (64,))
    def __len__(self): return len(self.seqs)
    def __getitem__(self, i): return self.seqs[i], self.labels[i]

def collate_pad(batch):
    seqs, labels = zip(*batch)
    padded = pad_sequence(seqs, batch_first=True, padding_value=0)
    lengths = torch.tensor([len(s) for s in seqs])
    return padded, lengths, torch.tensor(labels)

loader = DataLoader(VarSeqDataset(), batch_size=8, collate_fn=collate_pad)
padded, lengths, labels = next(iter(loader))
print(padded.shape, lengths)   # torch.Size([8, max_len_in_batch]) tensor([...])

Returning per-sample metadata — dict-shaped batches·python

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

class TaggedDataset(Dataset):
    def __init__(self):
        self.X = torch.randn(32, 10)
        self.y = torch.randint(0, 2, (32,))
        self.tags = [f"sample_{i}" for i in range(32)]
    def __len__(self): return 32
    def __getitem__(self, i):
        return {'x': self.X[i], 'y': self.y[i], 'tag': self.tags[i]}

def collate_dict(batch):
    return {
        'x': torch.stack([b['x'] for b in batch]),
        'y': torch.stack([b['y'] for b in batch]),
        'tags': [b['tag'] for b in batch],   # keep as list
    }

loader = DataLoader(TaggedDataset(), batch_size=8, collate_fn=collate_dict)
batch = next(iter(loader))
print(batch['x'].shape, batch['y'].shape, batch['tags'])

IterableDataset with proper worker splitting·python

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

class StreamingDataset(IterableDataset):
    """Stream samples from a generator. Splits across workers correctly."""
    def __init__(self, n_total=10_000):
        self.n_total = n_total

    def __iter__(self):
        info = torch.utils.data.get_worker_info()
        if info is None:                    # single-process
            start, end = 0, self.n_total
        else:
            per = self.n_total // info.num_workers
            start = info.id * per
            end = self.n_total if info.id == info.num_workers - 1 else start + per

        for i in range(start, end):
            yield torch.randn(10), i % 2

loader = DataLoader(StreamingDataset(), batch_size=32, num_workers=4)
print(sum(b[0].size(0) for b in loader))   # ≈ 10000

Custom Collate Functions and IterableDataset

When samples don't stack cleanly

Custom collate

IterableDataset — streaming data

Code

External links

Exercise

Progress

Comments 0