Machine Learning course menu

Training a Neural Network in PyTorch: The Training Loop

Build and train a neural network in PyTorch: nn.Module and nn.Sequential, loss functions, optimisers like Adam, Dataset and DataLoader for mini-batches, the standard training loop, evaluation mode, validation and early stopping.

Advanced11 min readUpdated October 24, 2026

On this page

Defining a model
Data in mini-batches
The training loop
Overfitting and early stopping
Common mistakes
Practice
Key takeaway

The previous lesson updated weights by hand. Real PyTorch code uses three building blocks instead:

nn.Module: the model, holding layers and their weights;
a loss function, such as nn.CrossEntropyLoss;
an optimiser, such as torch.optim.Adam, which updates every weight.

Defining a model#

Python

import torch
from torch import nn

class Classifier(nn.Module):
    def __init__(self, n_features, n_classes):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(n_features, 32),
            nn.ReLU(),
            nn.Linear(32, n_classes),
        )

    def forward(self, x):
        return self.layers(x)

torch.manual_seed(0)
model = Classifier(n_features=30, n_classes=2)
print(model)
print("parameters:", sum(p.numel() for p in model.parameters()))
print("output shape for 5 rows:", model(torch.randn(5, 30)).shape)

Output

Classifier(
  (layers): Sequential(
    (0): Linear(in_features=30, out_features=32, bias=True)
    (1): ReLU()
    (2): Linear(in_features=32, out_features=2, bias=True)
  )
)
parameters: 1058
output shape for 5 rows: torch.Size([5, 2])

nn.Linear(30, 32) holds a 30×32 weight matrix and 32 biases (992 parameters). The model outputs raw scores called logits, one per class. nn.CrossEntropyLoss applies softmax internally, so do not add a softmax layer at the end.

Data in mini-batches#

Networks train on small batches of rows at a time. TensorDataset wraps tensors, and DataLoader shuffles and batches them:

Python

import torch
from torch.utils.data import DataLoader, TensorDataset

X = torch.randn(100, 30)
y = torch.randint(0, 2, (100,))
loader = DataLoader(TensorDataset(X, y), batch_size=32, shuffle=True)
for xb, yb in loader:
    print(xb.shape, yb.shape)

Output

torch.Size([32, 30]) torch.Size([32])
torch.Size([32, 30]) torch.Size([32])
torch.Size([32, 30]) torch.Size([32])
torch.Size([4, 30]) torch.Size([4])

100 rows in batches of 32 gives three full batches and one batch of 4. For your own data formats (images on disk, text files), you write a Dataset class with __len__ and __getitem__.

The training loop#

The breast cancer data, scaled, with a training/validation split:

Python

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

X, y = load_breast_cancer(return_X_y=True)
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.25, random_state=0, stratify=y)
scaler = StandardScaler().fit(X_train)
X_train = torch.tensor(scaler.transform(X_train), dtype=torch.float32)
X_val = torch.tensor(scaler.transform(X_val), dtype=torch.float32)
y_train, y_val = torch.tensor(y_train), torch.tensor(y_val)

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(30, 32), nn.ReLU(), nn.Linear(32, 2))
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)

for epoch in range(1, 31):
    model.train()
    for xb, yb in loader:
        optimizer.zero_grad()            # 1. clear old gradients
        loss = loss_fn(model(xb), yb)    # 2. forward pass and loss
        loss.backward()                  # 3. backpropagation
        optimizer.step()                 # 4. update weights

    model.eval()
    with torch.no_grad():
        val_logits = model(X_val)
        val_loss = loss_fn(val_logits, y_val).item()
        val_acc = (val_logits.argmax(dim=1) == y_val).float().mean().item()
    if epoch in (1, 5, 10, 20, 30):
        print(f"epoch {epoch:>2}: val loss {val_loss:.3f}, val accuracy {val_acc:.3f}")

Output

epoch  1: val loss 0.575, val accuracy 0.804
epoch  5: val loss 0.232, val accuracy 0.923
epoch 10: val loss 0.132, val accuracy 0.951
epoch 20: val loss 0.104, val accuracy 0.965
epoch 30: val loss 0.101, val accuracy 0.958

Those four lines (zero gradients, forward and loss, backward, step) are the heart of every PyTorch training loop, from this tiny model to large language models.

model.train() and model.eval() switch layers such as dropout and batch normalisation between training and evaluation behaviour.
torch.no_grad() turns off gradient tracking during evaluation, saving memory and time.
argmax(dim=1) turns logits into predicted classes; torch.softmax(logits, dim=1) gives probabilities.

Overfitting and early stopping#

Train long enough and a network memorises its training data: training loss keeps falling while validation loss starts rising. Early stopping keeps the weights from the epoch with the best validation loss:

Python

import copy
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

X, y = load_breast_cancer(return_X_y=True)
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.25, random_state=0, stratify=y)
scaler = StandardScaler().fit(X_train)
X_train = torch.tensor(scaler.transform(X_train), dtype=torch.float32)
X_val = torch.tensor(scaler.transform(X_val), dtype=torch.float32)
y_train, y_val = torch.tensor(y_train), torch.tensor(y_val)

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(30, 128), nn.ReLU(), nn.Linear(128, 128), nn.ReLU(), nn.Linear(128, 2))
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)

best_loss, best_state, best_epoch, patience, waited = float("inf"), None, 0, 10, 0
for epoch in range(1, 201):
    model.train()
    for xb, yb in loader:
        optimizer.zero_grad()
        loss_fn(model(xb), yb).backward()
        optimizer.step()
    model.eval()
    with torch.no_grad():
        val_loss = loss_fn(model(X_val), y_val).item()
    if val_loss < best_loss:
        best_loss, best_state, best_epoch, waited = val_loss, copy.deepcopy(model.state_dict()), epoch, 0
    else:
        waited += 1
        if waited >= patience:
            print(f"stopped at epoch {epoch}; best epoch {best_epoch} with val loss {best_loss:.3f}")
            break

model.load_state_dict(best_state)
with torch.no_grad():
    print("restored model val accuracy:", round((model(X_val).argmax(1) == y_val).float().mean().item(), 3))

Output

stopped at epoch 21; best epoch 11 with val loss 0.116
restored model val accuracy: 0.958

Other ways to fight overfitting: dropout layers (nn.Dropout(0.2), randomly switching off neurons during training), weight decay (Adam(..., weight_decay=1e-4), like L2 regularisation), and more data.

Common mistakes#

Practice#

Practice exercise

Train a network on the iris data (4 features, 3 classes): one hidden layer of 16 ReLU neurons, Adam with lr=0.01, full-batch training (no DataLoader) for 200 epochs with torch.manual_seed(0), scaled features, and a stratified 30% test split (random_state=0). Print the test accuracy.

Show hint

With full-batch training, each epoch is one zero_grad / loss / backward / step on all training rows. Convert labels with torch.tensor(y_train) (int64).

Show solution

Python

import torch
from torch import nn
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0, stratify=y)
scaler = StandardScaler().fit(X_train)
X_train = torch.tensor(scaler.transform(X_train), dtype=torch.float32)
X_test = torch.tensor(scaler.transform(X_test), dtype=torch.float32)
y_train, y_test = torch.tensor(y_train), torch.tensor(y_test)

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()
for epoch in range(200):
    optimizer.zero_grad()
    loss_fn(model(X_train), y_train).backward()
    optimizer.step()

model.eval()
with torch.no_grad():
    print("test accuracy:", round((model(X_test).argmax(1) == y_test).float().mean().item(), 3))

Output

test accuracy: 1.0

Key takeaway#

Next: Image Classification with CNNs.

Finished this lesson?

Machine LearningPyTorch Tensors and Autograd for BeginnersGet started with PyTorch: installing it, creating tensors, shapes and dtypes, tensor maths and broadcasting, moving between NumPy and PyTorch, GPUs with .to(device), and automatic gradients with autograd.Intermediate8 min read Machine LearningImage Classification with CNNs in PyTorchClassify images with convolutional neural networks in PyTorch: how convolutions and pooling work, building a small CNN, training on handwritten digits, comparing with a plain network, data augmentation, and transfer learning with pretrained models.Advanced11 min read Machine LearningHow Neural Networks Learn: Neurons, Layers and BackpropagationUnderstand neural networks from the ground up: neurons, weights and biases, activation functions, layers, the loss, gradient descent and backpropagation, a tiny network trained in NumPy, and scikit-learn's MLPClassifier.Intermediate8 min read

Defining a model#

Data in mini-batches#

The training loop#

Overfitting and early stopping#

Common mistakes#

Practice#

Key takeaway#

Related lessons