🎓 A1Learner
🔍

PyTorch Train and Eval Modes

Train mode and eval mode are two settings carried by every module. model.train() selects train mode, and model.eval() selects eval mode.

Read the mode flag

The training flag reports the current mode. It flips with each call.

model = nn.Linear(2, 1)
model.train()
print("After train():", model.training)
model.eval()
print("After eval():", model.training)

Evaluate under no_grad

Evaluation needs no gradients. The output of a no_grad pass tracks none.

torch.manual_seed(0)
model = nn.Linear(1, 1)
model.eval()
with torch.no_grad():
    out = model(torch.tensor([[2.0]]))
print("Output:", round(out.item(), 2))
print("Tracks gradients:", out.requires_grad)

Example

import torch
from torch import nn

torch.manual_seed(0)
x = torch.linspace(-2, 2, 40).reshape(40, 1)
y = 3.0 * x - 2.0 + 0.3 * torch.randn(40, 1)

torch.manual_seed(0)
model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

model.train()
for step in range(300):
    pred = model(x)
    loss = loss_fn(pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

torch.manual_seed(7)
fresh_x = torch.tensor([[-1.5], [0.5], [1.5]])
fresh_y = 3.0 * fresh_x - 2.0

model.eval()
with torch.no_grad():
    predictions = model(fresh_x)
    error = torch.mean(torch.abs(predictions - fresh_y)).item()

print("Training flag in eval:", model.training)
print("Predictions:", [round(v, 2) for v in predictions.flatten().tolist()])
print("True values:", fresh_y.flatten().tolist())
print("Mean absolute error:", round(error, 3))

Key points

  • Use model.train() before a training loop
  • Use model.eval() before scoring a model
  • model.training: a True or False flag that reports the current mode
  • Most layers behave the same way in both modes
  • Use torch.no_grad() during evaluation to skip gradient tracking
  • The mode is a switch on the model, not a change to its weights

In this example, seed 0 is set before the data and again before the model. The data is the noisy line from the last lesson. model.train() runs before the training loop. The loop trains for 300 steps with Adam. Seed 7 then makes three fresh input points. Their true values come from the clean line, with no noise. model.eval() switches the model before scoring. The scoring pass runs inside torch.no_grad(). The flag print confirms eval mode is active. The predictions sit close to the true values. The mean absolute error averages the three gaps. It prints 0.036, a small error on fresh points.

Note

A later course meets dropout, a layer that behaves differently in each mode

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What should run before a block of code that scores a trained model?