The Training Loop
Networks learn by repeating: predict, measure loss, adjust weights via gradient descent.
Training a network is a loop repeated over the data many times. Each full pass is an epoch.
One step
- Forward pass — compute predictions.
- Loss — measure how wrong they are.
- Backward pass — compute gradients (which way each weight should move).
- Update — the optimiser nudges weights by the learning rate.
Frameworks do the gradients for you (autograd), but the concept is a simple loop.
Example
# Gradient descent on one weight (concept)
w = 0.0
learning_rate = 0.1
X = [1, 2, 3]; y = [2, 4, 6] # true relation y = 2x
for epoch in range(50):
grad = 0.0
for xi, yi in zip(X, y):
pred = w * xi
grad += 2 * (pred - yi) * xi # d(loss)/dw
w -= learning_rate * grad / len(X)
print(round(w, 3)) # ~2.0 learned the slopeWhen to use it
- A researcher implements a custom PyTorch training loop to add a per-batch learning-rate warm-up schedule not available in Keras.
- An engineer adds gradient clipping inside the training loop to stabilise training of a recurrent network that experiences exploding gradients.
- A team logs training loss and validation accuracy at the end of each epoch to a file so they can plot learning curves after training finishes.
More examples
Minimal PyTorch training loop
Implements the core training loop: forward pass, loss computation, zero gradients, backprop, and weight update — the four steps of every gradient descent iteration.
import torch
import torch.nn as nn
model = nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()
X = torch.randn(100, 2)
y = X[:, 0] * 2 - X[:, 1] + 0.5
for epoch in range(50):
pred = model(X).squeeze()
loss = loss_fn(pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Final loss:', loss.item())Training loop with validation
Extends the training loop with a validation phase using model.eval() and torch.no_grad() to compute val loss without affecting gradients.
import torch
import torch.nn as nn
model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
X_tr = torch.randn(200, 4); y_tr = torch.randint(0, 3, (200,))
X_val = torch.randn(50, 4); y_val = torch.randint(0, 3, (50,))
for epoch in range(5):
model.train()
loss = loss_fn(model(X_tr), y_tr)
optimizer.zero_grad(); loss.backward(); optimizer.step()
model.eval()
with torch.no_grad():
val_loss = loss_fn(model(X_val), y_val)
print(f'Epoch {epoch+1}: train={loss:.3f} val={val_loss:.3f}')Gradient clipping in the loop
Adds nn.utils.clip_grad_norm_ before optimizer.step() to cap the gradient magnitude at 1.0, preventing exploding gradients in deep or recurrent networks.
import torch
import torch.nn as nn
model = nn.Sequential(nn.Linear(10, 64), nn.Tanh(), nn.Linear(64, 1))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()
X = torch.randn(100, 10)
y = torch.randn(100, 1)
for epoch in range(20):
pred = model(X)
loss = loss_fn(pred, y)
optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
print('Done, final loss:', loss.item())
Discussion