# PyTorch training loop
The minimal pieces needed to define and train a neural network in PyTorch, from the model class to mini-batched epochs.
For tooling around the loop (profiling, TensorBoard, Lightning), see [PyTorch Toolchain](PyTorch%20Toolchain.md).
## Model structure
`torch.nn.Module` is the basic building block and container of a neural network in PyTorch.
A model subclasses it, defines its layers in `__init__`, and computes the output in `forward()`:
```python
from torch import nn
class Model(nn.Module):
def __init__(self) -> None:
super().__init__()
...
def forward(self, x):
...
return outputs
```
For classification with a loss that expects logits, keep `forward()` focused on producing logits, and apply softmax or sigmoid outside the model, only when probabilities are actually needed.
`CrossEntropyLoss` combines log-softmax with the negative log-likelihood loss, and `BCEWithLogitsLoss` includes the sigmoid, so applying the activation in the model would apply it twice.
For regression, the outputs are the predictions themselves; `MSELoss` does not require a final activation.
## What you need to train
1. Tensors holding the training data, `X` and `y`.
2. An optimizer from `torch.optim`, such as `Adam`, given the model's `.parameters()`.
3. A loss function from `torch.nn`, such as `MSELoss` or `CrossEntropyLoss`.
```python
from torch import tensor, optim
from torch.nn import CrossEntropyLoss
X = tensor([...])
y = tensor([...])
model = Model()
optimizer = optim.Adam(model.parameters(), lr=0.01)
loss_fn = CrossEntropyLoss()
```
## The minimal loop
```python
for step in range(1, STEPS + 1):
optimizer.zero_grad()
output = model(X)
loss = loss_fn(output, y)
loss.backward()
optimizer.step()
if step % 100 == 0:
print(f"Loss at step {step}: {loss.item()}")
```
What each step does:
- `optimizer.zero_grad()` clears the gradients from the previous step. PyTorch accumulates gradients in each parameter's `p.grad`, so without this, every step would add to the old gradients.
- `model(X)` runs the forward pass and records the computation graph.
- `loss.backward()` triggers autograd to compute the gradient of the loss with respect to every parameter and store it in `p.grad`.
- `optimizer.step()` updates the weights and biases using `p.grad` according to the optimizer's rule (for Adam, see [Adam optimizer tips](../../../Science/Machine%20Learning/Deep%20Learning/Adam%20optimizer%20tips.md)).
Gradient accumulation exploits the same behavior on purpose: calling `backward()` on several mini-batches before one `step()` and `zero_grad()` simulates a larger batch without the extra memory.
With mean-reduced losses and $K$ equally sized mini-batches, divide each loss by $K$ before `backward()` to match the larger batch's average gradient; otherwise, the mini-batch gradients are summed:
```python
optimizer.zero_grad()
for i, (x_batch, y_batch) in enumerate(loader, start=1):
loss = loss_fn(model(x_batch), y_batch) / K
loss.backward()
if i % K == 0:
optimizer.step()
optimizer.zero_grad()
```
If gradients explode, add [Gradient Clipping](Gradient%20Clipping.md) between `backward()` and `step()`.
## Mini-batching
A `Dataset` returns a single sample, typically `(x, y)`, from `__getitem__(idx)`.
A `DataLoader` retrieves multiple samples, groups them into batches, and stacks their tensors along a new batch dimension.
`TensorDataset` wraps existing tensors as a `Dataset`:
```python
from torch.utils.data import DataLoader, TensorDataset
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=100, shuffle=True)
for epoch in range(10):
for x_batch, y_batch in loader:
optimizer.zero_grad()
loss = loss_fn(model(x_batch), y_batch)
loss.backward()
optimizer.step()
```
`shuffle=True` draws a new random order every epoch.
Batch size and learning rate interact; see [Relationship between learning rate and batch size](../../../Science/Machine%20Learning/Deep%20Learning/Relationship%20between%20learning%20rate%20and%20batch%20size.md) and [Optimize the learning rate](../../../Science/Machine%20Learning/Deep%20Learning/Optimize%20the%20learning%20rate.md).
Batch size also drives VRAM use.
Activations, input batch tensors, and some temporary buffers scale with batch size; model weights, gradients, and optimizer state stay roughly constant.
## Sources
- PyTorch tutorials: [Build the neural network](https://docs.pytorch.org/tutorials/beginner/basics/buildmodel_tutorial.html), [Automatic differentiation with torch.autograd](https://docs.pytorch.org/tutorials/beginner/basics/autogradqs_tutorial.html), [Optimizing model parameters](https://docs.pytorch.org/tutorials/beginner/basics/optimization_tutorial.html), [Datasets and DataLoaders](https://docs.pytorch.org/tutorials/beginner/basics/data_tutorial.html)
- PyTorch docs: [Optimizer.zero_grad](https://docs.pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html), [torch.utils.data](https://docs.pytorch.org/docs/stable/data.html), [CrossEntropyLoss](https://docs.pytorch.org/docs/main/generated/torch.nn.CrossEntropyLoss.html), [Gradient accumulation example](https://docs.pytorch.org/docs/main/notes/amp_examples.html#gradient-accumulation)