# 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.

```python
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:

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

predictions = logits.argmax(dim=1)
```


### `BCEWithLogitsLoss`

Use binary cross-entropy with logits for binary classification.

```python
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:

```python
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.

```python
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.

```python
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.

```python
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](quickref-diagnosis.md).


### Weight Decay

Configure weight decay through the optimizer.

```python
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.

  ```python
  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:

  ```python
  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:

  ```python
  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.

```python
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:

```python
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.

```python
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.

```python
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`  |

:::{admonition} Multiclass classification
:class: dropdown, tip
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.

```python
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
```
:::

:::{admonition} Binary classification
:class: dropdown, tip
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.

```python
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
```
:::

:::{admonition} Regression
:class: dropdown, tip
A regression model predicts continuous values. Mean absolute error (MAE) calculates the average absolute difference between prediction and target.

```python
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](quickref-evaluation.md).

### Learning Curves

Store epoch-level measurements in a history structure.

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

Append the aggregate loss and metrics after each epoch.

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

After training, plot recorded measurements across epochs.

```python
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](quickref-diagnosis.md).