🎓 A1Learner
🔍

PyTorch Collate Functions

A collate function turns the list of rows a Dataset returns into one batch tensor. The DataLoader calls it on every batch.

The default collate gives up

Rows of different lengths cannot stack into one tensor. The default collate stops with an error.

loader = DataLoader(Readings(), batch_size=4)
next(iter(loader))

Padding by hand

Padding fills the short row with zeros up to the longest row. Every row then has the same length.

batch = [[1.0, 2.0, 3.0], [4.0, 5.0]]
longest = max(len(r) for r in batch)
padded = torch.zeros(len(batch), longest)
for i, row in enumerate(batch):
    padded[i, : len(row)] = torch.tensor(row)
print(padded)
print("Shape:", tuple(padded.shape))

Collate receives a list

The function gets a plain list of rows, one per dataset row in the batch. What it returns becomes the batch.

def show_collate(batch):
    print("Collate received:", type(batch).__name__, "of", len(batch), "rows")
    return batch

loader = DataLoader(Readings(), batch_size=3, collate_fn=show_collate)
next(iter(loader))

Example

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

torch.manual_seed(0)

# Each row is a short reading: a list of numbers with its own length.
rows = [
    [1.0, 2.0, 3.0],
    [4.0, 5.0],
    [6.0, 7.0, 8.0, 9.0, 10.0],
    [11.0],
    [12.0, 13.0, 14.0, 15.0],
    [16.0, 17.0],
]


class Readings(Dataset):
    def __len__(self):
        return len(rows)

    def __getitem__(self, index):
        return rows[index]


def pad_collate(batch):
    lengths = [len(row) for row in batch]
    longest = max(lengths)
    padded = torch.zeros(len(batch), longest)
    for i, row in enumerate(batch):
        padded[i, : len(row)] = torch.tensor(row)
    return padded, torch.tensor(lengths)


loader = DataLoader(Readings(), batch_size=4, collate_fn=pad_collate)

for padded, lengths in loader:
    print("Batch shape:", tuple(padded.shape))
    print("Lengths:", lengths.tolist())
    print(padded)

Key points

  • collate_fn: the function a DataLoader uses to build one batch from a list of rows
  • Use a custom collate_fn when rows have different lengths or nested parts
  • Padding: filling short rows with zeros so all rows share one length
  • Returning the true lengths beside the padded batch lets later code ignore the padding

In this example, the Readings dataset holds six short readings, each a list of numbers with its own length. The pad_collate function measures every row in the batch. It builds a zeros tensor as wide as the longest row. Then it copies each row into place, and the empty cells stay zero. It returns the padded tensor and the true lengths. The DataLoader runs with batch_size 4 and this collate function. The first batch prints with shape (4, 5), padded to the longest reading. The second batch holds the last two rows and prints with shape (2, 4).

Note

The default collate only works when every row already has the same shape.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why does the default collate fail on rows of length 3, 2, and 5?