Code Lab 01 - Raw PyTorch Classifier
A small CNN image classifier trained end-to-end in raw PyTorch - no Trainer class, no wrapper library. This lab exists to make every mechanic from the Notes files (Dataset/DataLoader, the hand-written training loop, checkpointing, mixed precision) concrete in one runnable script.
← Back to Overview: PyTorch Fundamentals · Back to Concepts: Training Loop From Scratch · Checkpointing & Mixed Precision
- Run a complete raw-PyTorch training job - DataLoader, CNN, hand-written loop, validation, checkpoints, mixed precision
- Resume training from a checkpoint and confirm the run continues where it stopped
- Measure the effect of batch size and mixed precision on speed and memory
- Training Loop From Scratch and Checkpointing & Mixed Precision
- Python 3.10+; a GPU is optional (AMP is disabled automatically on CPU)
What's In This Lab
| Property | Detail |
|---|---|
| Task | Image classification (MNIST digits, via torchvision.datasets) |
| Model | A small hand-written CNN (SimpleCNN in model.py) |
| Training | Hand-written loop - custom Dataset usage via torchvision, DataLoader, forward/backward/step |
| Checkpointing | Saves model + optimizer state every N epochs to checkpoints/ |
| Mixed Precision | torch.amp.autocast("cuda") + GradScaler (auto-disabled gracefully on CPU-only machines) |
| Output | Final train/val accuracy printed at the end of the run |
| Complexity | Beginner-Intermediate |
| Files | 01-Raw-PyTorch-Classifier/{train.py, model.py, requirements.txt, README.mdx} |
Architecture
flowchart LR
D["📦 torchvision.datasets.MNIST"] --> DL["🔄 DataLoader\nbatch_size=64, shuffle=True"]
DL --> M["🧠 SimpleCNN\nconv → pool → conv → pool → fc"]
M --> L["📉 CrossEntropyLoss"]
L --> BW["⬅️ backward() + GradScaler"]
BW --> OPT["🔧 optimizer.step()"]
OPT -->|every N epochs| CKPT["💾 checkpoints/model_epochN.pt"]
OPT --> Next["🔁 next epoch"]
Next -->|after final epoch| Eval["📊 Final train/val accuracy"]
style D fill:#d8dfe8,stroke:#b0bac8
style DL fill:#e8e0d4,stroke:#c8b89a
style M fill:#dde4dc,stroke:#b0c4b0
style L fill:#dde4dc,stroke:#b0c4b0
style BW fill:#ddd8e4,stroke:#b8b0c8
style OPT fill:#ddd8e4,stroke:#b8b0c8
style CKPT fill:#d8dfe8,stroke:#b0bac8
style Eval fill:#e8e0d4,stroke:#c8b89a
Running
cd 02-Prog-Langs/PyTorch/CodeLabs/01-Raw-PyTorch-Classifier
pip install -r requirements.txt
python train.py
Full details, code walkthrough, and what each part demonstrates: see README.
Check Yourself
- The run is interrupted after epoch 3 and restarted with the checkpoint. Which state must the checkpoint contain for training to continue identically?
- Why is GradScaler created with
enabled=use_ampin train.py? - Validation accuracy is noisy between two evaluations of the same checkpoint. What is the first thing to check?
Exercises
Introduce each bug separately and record what happens: (1) remove optimizer.zero_grad(); (2) remove model.eval() from validation; (3) move only the inputs, not the labels, to the GPU; (4) divide accuracy by len(loader). Which fail loudly and which silently?
Solution
(3) fails loudly with a device-mismatch error. (1) trains erratically - accumulated gradients make effective steps grow - but runs. (2) and (4) are silent: (2) gives noisier, slightly lower validation accuracy (dropout on), (4) reports accuracy inflated by about the batch size. Silent bugs are why the notes stress checks and asserts.
Run 2 epochs with batch sizes 32, 128 and 512, with and without AMP, recording epoch time and peak GPU memory. Explain the pattern.
Solution
Larger batches improve throughput until the GPU is saturated or the data loader becomes the bottleneck; memory grows roughly linearly with batch size (activations). AMP lowers memory and speeds up convolutions on tensor-core GPUs; on a model this small the gain is modest because the GPU is underused - the lesson is to profile before optimising.
References
- PyTorch tutorials, Learn the Basics (2026)
- PyTorch documentation, Automatic Mixed Precision (2026)
- LeCun et al., Gradient-Based Learning Applied to Document Recognition (1998) - the MNIST paper
Last reviewed: 2026-09