Training

Training#

Use this page to quickly reference how to calculate a loss, configure an optimizer, update model parameters, inspect gradients, train across mini-batches, and record training progress.

Loss Functions#

Choosing a Loss#

Choose the loss function according to the prediction task and the model output.

Task	Model output	Targets	Typical loss
Multiclass classification	logits `(N, K)`	class indices `(N,)`	`CrossEntropyLoss()`
Binary classification	logits `(N,)`	binary values `(N,)`	`BCEWithLogitsLoss()`
Scalar regression	predictions `(N,)`	real values `(N,)`	`MSELoss()`
Multi-output regression	predictions `(N, T)`	real values `(N, T)`	`MSELoss()`

Multiclass targets are integer tensors. Binary-classification and regression targets are floating-point tensors.

`CrossEntropyLoss`#

Use cross-entropy loss for multiclass classification with raw logits and integer class targets.

loss_fn = torch.nn.CrossEntropyLoss()

logits = model(inputs)
loss = loss_fn(logits, targets)

The loss function expects the following tensors, where N is the batch size and K is the number of classes.

logits - Floating-point tensor of shape (N, K) containing raw, unnormalized scores for each class.
targets - Integer tensor of shape (N,) containing integer class indices in the range [0, K-1].

Warning

Do not apply softmax before nn.CrossEntropyLoss(). The loss expects raw logits.

To obtain probabilities or class predictions after the model output has been computed:

probabilities = torch.softmax(logits, dim=1)

predictions = logits.argmax(dim=1)

`BCEWithLogitsLoss`#

Use binary cross-entropy with logits for binary classification.

loss_fn = torch.nn.BCEWithLogitsLoss()

logits = model(inputs)
loss = loss_fn(logits, targets)

The loss function expects the following tensors, where N is the batch size.

logits - Floating-point tensor of shape (N,) or (N, 1) containing raw scores for the positive class.
targets - Floating-point tensor of the same shape as logits, with binary values for the target class.

Adopt a consistent contract for the shapes of logits and targets. For example, use (N,) for both tensors. In this case, the model should flatten the output before returning it.

Warning

Do not apply sigmoid before nn.BCEWithLogitsLoss(). The loss expects raw logits.

To obtain probabilities or class predictions after the model output has been computed:

probabilities = torch.sigmoid(logits)

predictions = (logits >= 0).to(torch.int64)

A logit threshold of 0 corresponds to a probability threshold of 0.5.

`MSELoss`#

Use mean squared error for regression when the model predicts continuous values.

loss_fn = torch.nn.MSELoss()

predictions = model(inputs)
loss = loss_fn(predictions, targets)

Predictions and targets should have matching shapes and floating-point dtypes. For scalar regression, a consistent contract is to use the shape (N,) for both predictions and targets. Avoid relying on broadcasting to reconcile unintended shape differences.

Loss Reduction#

Most PyTorch loss functions support a reduction argument that controls how losses from multiple examples are combined.

Reduction	Result
`"mean"`	Mean over the losses
`"sum"`	Sum of the losses
`"none"`	No reduction; return a loss for each target

The epoch-loss aggregation pattern assumes that the loss function returns a mean loss for each batch.

Optimization#

Creating an Optimizer#

An optimizer updates the parameters passed to it.

optimizer = torch.optim.Adam(model.parameters())

Only parameters passed to the optimizer are eligible to be updated by optimizer.step().

Learning Rate#

Set the learning rate when creating the optimizer.

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

The learning rate controls the scale of parameter updates.

For diagnosing learning-rate problems, see Diagnostics.

Weight Decay#

Configure weight decay through the optimizer.

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)

Weight decay discourages unnecessarily large parameter values during optimization. Whether it improves validation performance must be determined experimentally.

Using an Optimizer#

A parameter update involves three separate operations: clearing old gradients, computing new gradients, and using those gradients to update the parameters.

loss.backward() computes the gradient of the loss with respect to every trainable parameter that contributed to the computation. These gradients are stored in each parameter’s .grad attribute.
optimizer.step() uses the stored gradients to update the parameters that were passed to the optimizer. The stored gradients are not cleared automatically after this call.
optimizer.zero_grad() clears the stored gradients for all parameters that were passed to the optimizer. This is normally done before computing gradients for a new batch. Otherwise, the gradients from the previous batch would be accumulated with the new gradients.

Inspecting Gradients#

Inspect gradients after loss.backward() and before clearing them.
```
for name, parameter in model.named_parameters():
    print(name, parameter.grad)
```
A parameter with requires_grad=False does not accumulate a gradient.
Check that an expected gradient exists with:
```
assert parameter.grad is not None
```
A None means that no gradient was accumulated for the parameter during the last backward pass.

Check that gradient values are finite:

for parameter in model.parameters():
    if parameter.grad is not None:
        assert torch.isfinite(parameter.grad).all()

Non-finite gradients indicate a numerical or optimization problem that should be investigated promptly.

Training Across Batches#

Training Mode#

Place the model in training mode before parameter updates.

model.train()

Training mode enables the training behaviour of modules such as dropout and batch normalization.

The Training Step#

A canonical training step follows this order:

optimizer.zero_grad()

outputs = model(inputs)
loss = loss_fn(outputs, targets)

loss.backward()
optimizer.step()

The model parameters persist across training steps. Each new batch computes a new loss and gradient, then updates the same model parameters.

Training One Epoch#

One epoch processes every batch from the training loader once.

model.train()

for inputs, targets in train_loader:
    inputs = inputs.to(device)
    targets = targets.to(device)

    optimizer.zero_grad()

    outputs = model(inputs)
    loss = loss_fn(outputs, targets)

    loss.backward()
    optimizer.step()

If the loader contains B batches, the loop performs B parameter updates during the epoch.

Monitoring Training#

Aggregating Training Loss#

Loss functions with reduction="mean" return a mean for the values represented by the current batch. To combine batch losses into one epoch-level mean, they should be weighted by the number of target values that contributed to it. If batch \(b\) reports mean loss \(L_b\) over \(M_b\) target values, the aggregated training loss is

\[ L_{\text{epoch}} = \frac{\sum_b M_b L_b}{\sum_b M_b}. \]

A general accumulation pattern for the loss functions mentioned on this page is the following.

total_loss = 0.0
total_weight = 0

for inputs, targets in train_loader:
    ...

    batch_weight = targets.numel()
    total_loss += loss.item() * batch_weight
    total_weight += batch_weight

epoch_loss = total_loss / total_weight

Weighting the batch loss by targets.numel() takes into account the number of values that contributed to the mean loss. For example, nn.CrossEntropyLoss() computes a mean over N target class indices, while nn.MSELoss() computes a mean over N*T target values for multi-output regression.

The aggregation is used only to report training performance. Parameter updates have already been performed separately for each batch and are unaffected by how batch losses are aggregated afterward.

Note

During an epoch, the model is updated after every batch. The aggregated training loss does not represent the loss of a single model. It only summarizes the losses observed throughout the epoch.

Aggregating Training Metrics#

Training metrics can be accumulated during the same batch loop used for optimization. Instead of calculating a metric independently for each batch and then averaging those batch metrics, accumulate the quantities needed to compute the metric over the complete epoch.

Metric	Accumulate during the epoch	Compute after the epoch
Multiclass accuracy	Correct predictions – Number of examples	`correct / examples`
Binary accuracy	Correct predictions – Number of target values	`correct / values`
Mean absolute error	Absolute error – Number of target values	`absolute_error / values`
Mean squared error	Squared error – Number of target values	`squared_error / values`

Multiclass classification

A multiclass classifier produces one logit per class for each example. The predicted class is the one with the largest logit. Accuracy then measures the fraction of examples for which this predicted class matches the target class.

total_correct = 0
total_examples = 0

for inputs, targets in train_loader:
    ...

    predictions = outputs.argmax(dim=1)
    total_correct += (predictions == targets).sum().item()
    total_examples += targets.shape[0]

epoch_accuracy = total_correct / total_examples

Binary classification

A binary classifier produces one logit per example. The predicted class is positive if the logit is greater than or equal to zero, which corresponds to a probability greater than or equal to 0.5. Accuracy measures the fraction of examples where the predicted class matches the target class.

total_correct = 0
total_values = 0

for inputs, targets in train_loader:
    ...

    predictions = outputs >= 0
    binary_targets = targets >= 0.5
    assert predictions.shape == binary_targets.shape

    total_correct += (predictions == binary_targets).sum().item()
    total_values += targets.numel()

epoch_accuracy = total_correct / total_values

Regression

A regression model predicts continuous values. Mean absolute error (MAE) calculates the average absolute difference between prediction and target.

total_absolute_error = 0.0
total_values = 0

for inputs, targets in train_loader:
    ...

    total_absolute_error += (outputs - targets).abs().sum().item()
    total_values += targets.numel()

epoch_mae = total_absolute_error / total_values

The same principle applies to all metrics: accumulate the quantities required by the metric across batches, then compute the final value after the epoch. Avoid taking an unweighted mean of independently computed batch metrics when batches may represent different numbers of examples or target values.

Note

Training metrics aggregate values measured while the model is being updated. They are useful for monitoring optimization, but they differ from evaluation metrics measured on a fixed model.

For task-specific prediction rules and evaluation metrics, see Evaluation.

Learning Curves#

Store epoch-level measurements in a history structure.

history = {
    "train_loss": [],
    "train_accuracy": [],
}

Append the aggregate loss and metrics after each epoch.

history["train_loss"].append(epoch_loss)
history["train_accuracy"].append(epoch_accuracy)

After training, plot recorded measurements across epochs.

import matplotlib.pyplot as plt

plt.plot(range(1, epochs + 1), history["train_loss"])
plt.xlabel("Epoch")
plt.ylabel("Training loss")
plt.show()

Learning curves preserve how training measurements evolve over time. They should be recorded for both training and evaluation metrics to help diagnose optimization problems. Their interpretation is covered in Diagnostics.