Minibatches: Dataset & DataLoader - PyTorch Basics

Minibatches: Dataset & DataLoader

What you'll build: make_batches(X, y, batch_size, shuffle, drop_last, seed), which returns every minibatch of one epoch, like iterating a PyTorch DataLoader.

Background

Training on the whole dataset at every step is slow, and training on it in a fixed order biases each epoch. A data loader shuffles the sample indices once per epoch and cuts them into minibatches. In PyTorch this is

loader = DataLoader(TensorDataset(X, y), batch_size=32, shuffle=True, drop_last=False,
                    generator=torch.Generator().manual_seed(0))
for xb, yb in loader:
    ...

Algorithm

Return all minibatches of one epoch as a list:

1. Build an index order: np.arange(N) if shuffle=False, otherwise exactly np.random.default_rng(seed).permutation(N).
2. Cut the order into consecutive chunks of batch_size. The last chunk can be smaller.
3. If drop_last=True, drop the last chunk when it has fewer than batch_size samples.
4. For each chunk idx, emit the pair (X[idx], y[idx]), so rows of X stay paired with their labels.

Shuffling with a seeded generator keeps the run reproducible: the same seed gives the same batches.

N	batch_size	drop_last	number of batches	last batch size
10	4	False	3	2
10	4	True	2	4
12	4	either	3	4

Inputs

• X: array (N, ...); y: array (N,).
• batch_size: positive integer; shuffle, drop_last: booleans; seed: integer.

Expected Return Format

A list of tuples (X_batch, y_batch) with shapes (B, ...) and (B,), where B = batch_size except possibly for the last batch.

Constraints

• NumPy track: the shuffled order must be exactly np.random.default_rng(seed).permutation(N) (a hidden test checks it).
• PyTorch track: use TensorDataset and DataLoader with generator=torch.Generator().manual_seed(seed). The order then comes from torch's RNG, so the tests check coverage, pairing and determinism instead of an exact order.

Follow-up

• Why shuffle every epoch instead of once? What goes wrong if the data is sorted by label and you never shuffle?
• What do num_workers and pin_memory do in a real DataLoader?

Input: X = np.arange(20).reshape(10, 2), y = np.arange(10), batch_size=4, shuffle=False
Output: [(X[0:4], [0, 1, 2, 3]), (X[4:8], [4, 5, 6, 7]), (X[8:10], [8, 9])]

Input: same, drop_last=True
Output: [(X[0:4], [0, 1, 2, 3]), (X[4:8], [4, 5, 6, 7])]

import numpy as np

def make_batches(X, y, batch_size, shuffle=True, drop_last=False, seed=0) -> list:
    """
    Split (X, y) into minibatches for one epoch.

    Args:
        X: Inputs of shape (N, ...)
        y: Labels of shape (N,)
        batch_size: Samples per batch
        shuffle: Shuffle with np.random.default_rng(seed).permutation(N)
        drop_last: Drop the final batch if it is smaller than batch_size
        seed: Seed for the shuffle

    Returns:
        List of (X_batch, y_batch) tuples
    """
    # Your code here
    pass