PyTorch — DataLoader, GPU, save/load

📅 Day 45. PyTorch's data abstraction is one of its cleanest design choices. Two classes, separation of concerns.

📦 Dataset. Defines 'how to get one example'. Two methods — __len__ returns the dataset size, __getitem__(i) returns the i-th example as a (X, y) tuple. That's it. Subclass torch.utils.data.Dataset, implement these two methods, you have a dataset.

The Dataset doesn't care about batching or shuffling. It just knows how to retrieve individual examples. This separation lets you focus on the data loading logic without worrying about batching plumbing.

🚚 DataLoader. Wraps a Dataset to add batching, shuffling, and parallel loading. DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4) gives you batches of 32 examples, randomly shuffled, loaded by 4 worker processes in the background.

The DataLoader is what you iterate in your training loop — for X, y in train_loader: gives you batches.

🔌 Why this design wins:

→ Drop in any data source. Files, S3, an HTTP API, a generator, a database query. Implement Dataset; DataLoader handles the rest.

→ Parallel loading. num_workers > 0 spawns worker processes that pre-fetch batches while your model is training. Removes data loading from the critical path.

→ Memory efficiency. The DataLoader pulls one batch at a time. The dataset can be too big for memory; only the active batches need to fit.

→ Reusability. Same DataLoader code works for image, text, audio, tabular. The Dataset is the swappable part.

💡 For quick experiments, torch.utils.data.TensorDataset wraps existing tensors directly. tensor_dataset = TensorDataset(X, y) gives you a Dataset without writing a class. Pair with DataLoader and you're training in 5 lines.

🚀 Master the pattern once; reuse it forever.

🖥 GPU usage in PyTorch is straightforward — until it isn't. The number-one bug — model on GPU, data on CPU. PyTorch raises 'Expected all tensors to be on the same device' and you wonder why.

📋 Three things must live on the same device:

1️⃣ The model. model = model.to(device).

2️⃣ The inputs. X = X.to(device) inside your training loop.

3️⃣ The targets. y = y.to(device) inside your training loop.

Mismatch any of them → error.

🎯 My default top-of-file device detection:

import torch

device = (
    'cuda' if torch.cuda.is_available() else
    'mps'  if torch.backends.mps.is_available() else
    'cpu'
)

This tries CUDA (NVIDIA GPUs) first, then MPS (Apple Silicon), then falls back to CPU. The same code runs on a Linux GPU server, a MacBook M-series, or a Windows laptop without changes.

📦 Then in code:

model = MLP(...).to(device)  # once at startup

for X, y in loader:
    X, y = X.to(device), y.to(device)  # every batch
    out = model(X)
    ...

⚡ Pro tips:

→ For very large datasets, pin_memory=True in DataLoader speeds up CPU→GPU transfer (~10-30% on heavy training).

→ For multiple GPUs, torch.nn.DataParallel or torch.nn.DistributedDataParallel splits the batch across GPUs. Modern projects use DDP; DataParallel is legacy.

→ For Apple Silicon, MPS support has improved a lot. Most operations now work; some don't. Check error messages and fall back to CPU for unsupported ops.

💡 Get device handling right once at the top of your codebase. Then forget about it. The 'same device' bug becomes invisible.

💾 PyTorch model serialisation has one rule that beginners constantly violate — save the state_dict, NOT the model object.

📦 The wrong way:

torch.save(model, 'model.pt')
model = torch.load('model.pt')

This pickles the entire model object including its class definition. Two problems:

→ Loading requires the SAME class definition to be importable. If you renamed a layer or refactored the file structure, loading fails with cryptic errors.

→ Tightly coupled to PyTorch internals. Changes to PyTorch versions can break loading.

✅ The right way — save state_dict, which is a Python dict mapping parameter names to tensors:

torch.save(model.state_dict(), 'model.pt')

# Later
model = MLP(in_d, hid, out_d).to(device)  # recreate the architecture
model.load_state_dict(torch.load('model.pt', map_location=device))
model.eval()  # set to inference mode

Notice the steps:

1️⃣ Recreate the model architecture (you need the class available).

2️⃣ Move it to the right device.

3️⃣ Load the state_dict (parameters only).

4️⃣ map_location=device handles the case where the model was saved on GPU but you're loading on CPU (or vice versa). Without it, you'd get device errors.

5️⃣ model.eval() switches BatchNorm and Dropout to inference mode. Critical step (covered tonight in detail).

💼 What this gives you:

→ Portable. Load on different machines, different PyTorch versions, different devices.

→ Future-proof. Your model architecture file can change names; as long as the layer structure is the same, the state_dict loads.

→ Smaller files. Just the weights, no Python pickling overhead.

💡 For production with frequent retraining, also save metadata — training date, evaluation metrics, hyperparameters used. JSON file alongside the .pt. Future-you will know what version is what.

💡 Two flips at inference time. Skip either and you have subtle bugs.

🎯 Flip 1 — model.eval()

This switches certain layers (BatchNorm, Dropout) into inference mode.

→ Dropout in train mode randomly zeros out a fraction of activations. In eval mode, it does nothing (passes inputs through unchanged).

→ BatchNorm in train mode uses the current batch's statistics for normalisation. In eval mode, it uses the running mean/variance accumulated during training.

Forget to call eval() and your model produces noisy, suboptimal outputs at inference. The bug is silent — outputs look reasonable, just not as good as they should be. Easy to miss; embarrassing in production.

🚫 Flip 2 — torch.no_grad()

Disables PyTorch's gradient tracking. Inference doesn't need gradients (we're not training); tracking them wastes memory and CPU.

with torch.no_grad():
    output = model(X)

Memory cost drops because intermediate activations don't need to be stored for backprop. Speed improves 30-50% because the autograd machinery is bypassed.

🎯 The combined inference pattern:

model.eval()
with torch.no_grad():
    for X in test_loader:
        X = X.to(device)
        output = model(X)
        # process output

💼 In production with PyTorch:

→ At model load, call model.eval() once. It stays in eval mode until you call model.train() to switch back.

→ Wrap inference code in torch.no_grad(). Or use the @torch.no_grad() decorator on inference functions for cleaner syntax.

→ For inference-only deployments (no training), consider torch.jit.trace or torch.compile for further speedups. We won't cover them in this sprint.

💡 Two lines. Skipping either is a top-5 PyTorch bug. Do them both, every time.

📅 End of Day 45.

✅ Recap:

📦 Dataset + DataLoader pattern. Dataset = single-example access. DataLoader = batching, shuffling, parallel loading. Reusable across data types.

🖥 Device handling. Detect at startup; .to(device) for model and every batch. Same device for everything or PyTorch yells.

💾 Save state_dict, not the model object. Recreate architecture at load time. map_location for cross-device loading. model.eval() after loading.

🎯 Inference always has model.eval() + torch.no_grad(). Otherwise dropout adds noise and gradient tracking wastes resources.

🧠 Reflection — PyTorch's API is small enough to fit in your head. Five concepts (Module, DataLoader, Loss, Optimiser, Device) plus the training loop covers 95% of what you'll write. Once you've internalised these, building any neural network architecture is composition.

🚀 Tomorrow, Day 46 — CNNs. Convolutions, pooling, why they still matter for vision. Tiny CNN for MNIST in 14 lines. ResNet's skip connections that unlocked deep networks. The 'always start from a pretrained model' rule.

💼 PyTorch fundamentals are done. Tomorrow we apply them to specific architectures.

👋 See you in the morning.