Day 29: Writing the training loop by hand
The five lines that train every neural network
Every training loop, from a toy MLP to GPT, is the same skeleton: for each batch — forward, compute loss, zero grads, backward, step. The roadmap insists you write this by hand every time this month rather than reaching for a framework, precisely so it becomes muscle memory. Below is the complete loop; every line is something you now understand from first principles.
model.train() # enable dropout/batchnorm training behavior
for epoch in range(num_epochs):
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad() # clear accumulated grads (Day 27)
outputs = model(images) # forward pass (Day 23)
loss = criterion(outputs, labels) # cross-entropy (Day 24)
loss.backward() # autograd backward (Day 27)
optimizer.step() # update params (Day 24)
print(f"epoch {epoch}: loss={loss.item():.4f}")model.train() vs model.eval() is not optional
Some layers behave differently in training vs inference — dropout is active only in training, batchnorm uses batch statistics in training but running averages in eval. Forgetting model.eval() before validation is a classic bug that produces mysteriously worse (and non-reproducible) validation numbers. Set the mode explicitly, always.
Key terms
- criterion
- The loss function object (e.g. nn.CrossEntropyLoss) applied to model outputs and true labels.
- optimizer.step()
- Applies one parameter update using the gradients currently stored in each parameter's .grad.
- model.train() / model.eval()
- Switches layers like dropout and batchnorm between their training and inference behaviors.
In the training loop, what is the correct order of these four calls?
