KAHIBARO
5.3. Collate Functions and Batching
Table of Contents
What Batching Really Means in PyTorch
A DataLoader produces batches so your model can process multiple examples at once. Conceptually, batching is just grouping several samples from a Dataset into a single “mini batch” object. In practice, batching is also where many real world details show up, like variable length inputs, optional fields, and non tensor metadata.
In PyTorch, the DataLoader creates a batch by repeatedly calling the Dataset’s __getitem__ to get individual samples, collecting those samples into a Python list, then passing that list to a function called collate_fn. The collate function turns a list of samples into the batch format your training loop expects.
The Default Collate Function and Its Assumptions
If you do not specify collate_fn, PyTorch uses a default collator. It works well when each sample has the same structure and compatible shapes. Typical examples are image classification samples like (image_tensor, label_tensor) where all images have the same height and width.
The default collator usually stacks tensors along a new first dimension. If a single sample feature tensor has shape (C, H, W), a batch of size B becomes (B, C, H, W). If labels are scalars, a batch becomes a 1D tensor of shape (B,).
The default collate function requires that tensors it stacks have matching shapes. If your samples contain variable length tensors, like different sequence lengths or different image sizes, the default collator will fail.
How Samples Should Be Structured for Easy Batching
A good rule is to return a consistent structure from your Dataset. Common patterns are a tuple like (x, y) or a dictionary like {"x": x, "y": y, "id": id}. The collate function will batch each field across samples.
If you return dictionaries, your batch is often a dictionary where each key maps to a batched tensor. This can make training code clearer because you can access batch["x"] and batch["y"] directly.
Whatever structure your Dataset returns for one sample, your collate function should convert a list of those samples into the same kind of structure, but batched. Your training loop should not have to guess what it received.
Writing a Custom `collate_fn`
A custom collate function is just a Python callable that accepts batch, which is a list of samples, and returns a single batched object. The simplest custom collate separates fields, stacks what can be stacked, and keeps what should stay as a list.
Here is a common template for a Dataset that returns (x, y):
import torch
def collate_xy(batch):
xs, ys = zip(*batch) # tuples of length B
x = torch.stack(xs, dim=0) # (B, ...)
y = torch.tensor(ys) # (B,) if ys are numbers
return x, yYou then pass it to the DataLoader:
from torch.utils.data import DataLoader
loader = DataLoader(dataset, batch_size=32, collate_fn=collate_xy)Variable Length Data: Padding in the Collate Function
Many beginner projects involve variable length sequences, like text tokens, time series, or lists of detections. In these cases, batching often means padding shorter sequences so that all sequences in the batch share a common length.
A typical approach is to return both the padded batch and a mask or the original lengths. The collate function is a good place to do this because it can see the whole batch and compute the maximum length.
If your Dataset returns {"tokens": tokens_1d, "label": label} where tokens_1d has variable length, a collate function can pad:
import torch
from torch.nn.utils.rnn import pad_sequence
def collate_pad_tokens(batch, pad_value=0):
tokens_list = [item["tokens"] for item in batch]
labels = torch.tensor([item["label"] for item in batch], dtype=torch.long)
padded = pad_sequence(tokens_list, batch_first=True, padding_value=pad_value)
lengths = torch.tensor([t.numel() for t in tokens_list], dtype=torch.long)
return {"tokens": padded, "lengths": lengths, "label": labels}
Now tokens has shape (B, L_max) where $L_{max}$ is the maximum sequence length in the batch. Many models also need an attention mask. You can build it in collate as well:
def make_padding_mask(lengths, max_len):
idx = torch.arange(max_len).unsqueeze(0) # (1, max_len)
return idx < lengths.unsqueeze(1) # (B, max_len), booleanIf you pad sequences, you must ensure your model or loss ignores padding positions. That usually means using a mask or a loss setting like an ignore index for padded labels.
Batching “Ragged” or Heterogeneous Data
Sometimes you cannot or do not want to pad, for example when samples contain a variable number of objects, boxes, or graph nodes. In those cases a valid batch can be a list of per sample tensors rather than a single stacked tensor.
A collate function can intentionally return lists:
def collate_keep_lists(batch):
images = torch.stack([b["image"] for b in batch]) # images are fixed size
boxes = [b["boxes"] for b in batch] # variable number of boxes
labels = [b["labels"] for b in batch]
return {"image": images, "boxes": boxes, "labels": labels}This is common in detection style tasks. The training step then iterates over list fields or uses model code designed for lists.
Batching and Dtypes: Getting `torch.tensor` Right
When creating tensors inside collate_fn, explicitly set dtype when needed. Labels for classification are typically torch.long. Regression targets are often torch.float32. If you do torch.tensor(ys) without specifying dtype, PyTorch infers it from the Python values, which can be surprising.
Classification targets used with common classification losses usually must be integer class indices, typically dtype torch.long. Create them as torch.tensor(labels, dtype=torch.long) to avoid silent mistakes.
Where Device Placement Fits In
Even though it is technically possible to move tensors to GPU inside collate_fn, it is usually better to keep collation on CPU and move the resulting batch to the device in your training loop. This keeps the data pipeline simpler and avoids tricky interactions with DataLoader workers.
So collate should focus on structure, padding, stacking, masking, and dtype, not GPU transfers.
Debugging Collation Problems Quickly
Most collation issues show up as shape errors or type errors when stacking. A practical way to debug is to fetch a single batch and print its structure and shapes:
batch = next(iter(loader))
print(type(batch))
If it is a tuple or dict, inspect each field’s shape, dtype, and whether it is a list.
If you see an error like “stack expects each tensor to be equal size”, the fix is not in the model. It is almost always a batching issue, either pad variable length data or return lists instead of stacking.
Summary: When to Use a Custom Collate Function
You need a custom collate_fn when any of these are true. Your samples have variable length components and you want padding and masks. Your samples contain optional or variable count fields that should remain as lists. Your batch needs special dtype handling or extra computed fields like lengths or masks. When samples are uniform tensors, the default collator is enough and usually preferable.
Views: 97