Build a small PyTorch model and inspect its training, validation, and failure behavior.

First neural network training loop

Prerequisites

Learning objectives

  • identify tensors, parameters, predictions, and loss;
  • separate training from evaluation;
  • record a baseline and validation result; and
  • inspect errors rather than celebrating a single score.

A small, reproducible model

This example uses a tiny numeric dataset so the mechanics are visible. It is not evidence that a neural network is needed for a real project.

import torch
from torch import nn

torch.manual_seed(7)
X = torch.tensor([[0.0], [1.0], [2.0], [3.0]])
y = torch.tensor([[1.0], [3.0], [5.0], [7.0]])

model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

for epoch in range(100):
    prediction = model(X)
    loss = loss_fn(prediction, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

with torch.no_grad():
    print(model(torch.tensor([[4.0]])).round(decimals=2))

The model learns an approximate relationship in the small training set. In a serious project, keep separate validation and test data, compare with a simpler baseline, and test inputs outside the training examples.

Training versus evaluation

During training, gradients update parameters. During evaluation, disable updates and measure performance without changing the model:

model.eval()
with torch.no_grad():
    validation_prediction = model(X)

For models with dropout or batch normalization, model.eval() changes evaluation behavior. Always record whether a result came from training or evaluation mode.

Practice

  1. Change the data so one point is noisy.
  2. Compare the neural-network result with a constant baseline.
  3. Plot loss by epoch.
  4. Test a new input and state why extrapolation may be unreliable.

Safety and limitations

Do not train on identifiable student data without approval. A low loss on toy data says nothing about fairness, robustness, or suitability for a consequential decision.

References