A collate function is a DataLoader callback that converts the list of individual Dataset outputs into one batch. Use it when the default collator cannot combine items directly because their shapes, lengths, or target structures differ. section 4.3.1 section 8.4.2
Common collation strategies include:
- For variable-length sequences, collect lengths and inputs, pad shorter tensors, and optionally pass the padded result with lengths to pack_padded_sequence. section 4.3.1
- For sequence-to-sequence data, pad source and target tensors independently to their batch maximums, then return the nested structure required by the training loop. section 11.2.1
- For MixUp, first apply a base collator, then sample a mixing coefficient, permute the batch, mix inputs, and return paired labels. section 14.3.2
- For Faster R-CNN, preserve variable-sized images and annotations by returning a list of tensors and a list of dictionaries instead of one padded tensor. section 8.4.2
Padding a translation batch
A pad_batch function finds max_x and max_y, right-pads each input and target with the PAD token, stacks both lists, and returns ((X, Y), Y). Configure it with DataLoader(..., collate_fn=pad_batch). The duplicated Y supplies both the model input structure and the training target expected by the loop. section 11.2.1
def pad_and_pack(batch):
input_tensors = [] #1
labels = []
lengths = []
for x, y in batch:
input_tensors.append(x)
labels.append(y)
lengths.append(x.shape[0]) #2
x_padded = torch.nn.utils.rnn.pad_sequence( #3
input_tensors, batch_first=False)
x_packed = #4
torch.nn.utils.rnn.pack_padded_sequence(x_padded, lengths,
batch_first=False, enforce_sorted=False)
y_batched = torch.as_tensor(labels, dtype=torch.long) #5
return x_packed, y_batched #6
#1 Organizes the batch input lengths, inputs, and outputs as separate lists
#2 Assumes shape is (T, *)
#3 Creates the padded version of the input
#4 Creates the packed version from the padded and lengths
#5 Converts the lengths into a tensor
#6 Returns a tuple of the packed inputs and heir labels
Inspect the step-by-step pad_and_pack implementation: it separates inputs, labels, and lengths, pads sequences, packs them, and converts labels into a tensor. section 4.3.1
