Length-Bucketed DataLoader

Build a DataLoader for sequences of different lengths that wastes as little compute on padding as possible.

A batch of sequences must be padded to its longest member. If a batch mixes a length 3 sequence with a length 60 one, most of the tensor is padding and the model spends most of its time on zeros. Bucketing groups sequences of similar length into the same batch:

order   = indices sorted by sequence length
batches = consecutive chunks of batch_size from order
each epoch: shuffle the order of the batches (not their contents)

You need three pieces of torch.utils.data:

• a Dataset whose item i is (tokens as a long tensor, label),
• a batch sampler, an iterable of lists of indices, one list per batch,
• a collate_fn that pads a list of items into one batch.

Each batch is a tuple (padded, lengths, mask, labels):

padded   (B, T) int64, padded with 0, T is the longest in this batch
lengths  (B,)   int64
mask     (B, T) bool, True on real tokens
labels   (B,)   int64

Function signature

def make_loader(sequences, labels, batch_size, seed=0):

Example

Input

sequences = [[5, 6, 7], [1], [8, 9], [2, 3, 4, 5]]
labels = [0, 1, 0, 1]
loader = make_loader(sequences, labels, batch_size=2)

Output

two batches, in a seeded random order:
padded [[1, 0], [8, 9]]                      lengths [1, 2]   labels [1, 0]
padded [[5, 6, 7, 0], [2, 3, 4, 5]]          lengths [3, 4]   labels [0, 1]

Reasoning

Sorted by length the indices are 1, 2, 0, 3. The chunks [1, 2] and [0, 3] need 1 and 1 padding tokens. A random split could pair lengths 1 and 4 and need 3.

Constraints and notes

• Return a torch.utils.data.DataLoader. Pass your sampler as batch_sampler=, not sampler=.
• Every example appears exactly once per epoch, and no batch is larger than batch_size.
• The batch order must depend on seed: the same seed gives the same order, and a different seed gives a different one.
• Graded on the result too: total padding must be within 10% of the minimum possible for that batch size.