Training is the repeated adjustment of a model’s parameters so its outputs better match examples and their targets. In a neural network, the adjustable parameters are weights and biases. Rather than handcrafting every useful feature, training can also help the network discover representations in the data. The aim is not merely to fit the examples it has seen, but to produce useful outputs on unseen data. section 6.1.6 section 1.2
A typical PyTorch training loop repeats these steps:
- Get a batch of inputs and targets from a DataLoader, then run the model to produce outputs. section 8.4
- Use a loss function to measure the difference between outputs and targets. section 1.5
- Clear old gradients, call loss.backward() to calculate new ones, and call optimizer.step() to update the parameters. section 8.4 section 6.2.2
An epoch is one complete pass through the training dataset; an iteration usually processes one batch. A small dataset may fit in one batch, making an iteration an epoch. With larger datasets, many iterations make up an epoch. Shuffling examples and updating on minibatches gives stochastic gradient descent its randomness. Track loss as training proceeds, and check separate validation data: a low training loss alone does not show how well the model handles unseen examples. section 5.4.3 section 7.2.5
From a simple fit to an image classifier
For a one-input linear model, nn.Linear(1, 1) holds the weight and bias. Pass model.parameters() to optim.SGD, compare predictions with targets using nn.MSELoss(), then repeat backward and optimizer steps. The same loop structure trains an image classifier: a DataLoader supplies shuffled image batches, while a classification loss replaces mean squared error. In one bird-versus-airplane example, validation accuracy is checked after training. section 6.2.2 section 7.2.5
# In[31]:
import datetime #1
def training_loop(n_epochs, optimizer, model, loss_fn, train_loader):
for epoch in range(1, n_epochs + 1): #2
start_time = datetime.datetime.now()
loss_train = 0.0
for imgs, labels in train_loader: #3
outputs = model(imgs) #4
loss = loss_fn(outputs, labels) #5
optimizer.zero_grad() #6
loss.backward() #7
optimizer.step() #8
loss_train += loss.item() #9
end_time = datetime.datetime.now()
epoch_duration = (end_time - start_time).total_seconds()
if epoch == 1 or epoch % 10 == 0:
print('{} Epoch {}, Training loss {:.6f}, Time {:.2f}s'.format(
datetime.datetime.now(), epoch,
loss_train / len(train_loader), epoch_duration)) #10
The loop shows the forward pass, loss calculation, gradient reset, backward pass, and optimizer update inside an epoch-and-batch structure. section 8.4
If loss rises rapidly or becomes infinite, updates may be too large; try a smaller learning rate. Normalizing inputs can also make gradients easier to manage. The core loop carries over to other tasks, but the data and loss change: a diffusion model learns to predict added noise, while a segmentation model trains on images and masks and checks validation loss after each epoch. section 5.4.3 section 5.4.4 section 10.6.1 section 15.7.3
