Wesam Khallaf — Author of PyTorch From Ground Up

Posted on Oct 10 AI-assisted

model.train() vs model.eval() vs torch.no_grad() in PyTorch: What Each One Changes

#machinelearning #beginners #deeplearning #pytorch

PyTorch has three calls that are used around evaluation: model.train(), model.eval() and torch.no_grad(). They control two separate things. train() and eval() set a mode flag that changes how some layers compute their output. torch.no_grad() controls whether autograd records the computation. This article explains each one, shows which layers are affected, and gives the standard train-and-evaluate loop.

The training flag

Every nn.Module has a boolean attribute training. model.train() sets it to True on the model and on every sub-module. model.eval() sets it to False. A new model starts in training mode.

import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Sequential(
    nn.Linear(10, 32),
    nn.BatchNorm1d(32),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(32, 3),
)

print(model.training)
model.eval()
print([m.training for m in model])
model.train()
print([m.training for m in model])

Output:

True
[False, False, False, False, False]
[True, True, True, True, True]

model.eval() is the same call as model.train(False). Both return the model itself, so model = model.eval() and model.eval() do the same thing.

Setting the flag does nothing by itself. It only matters for layers that read it in their forward(). In the standard layer set, these are mainly the dropout layers and the normalization layers that keep running statistics (BatchNorm1d, BatchNorm2d, BatchNorm3d). nn.Linear, nn.Conv2d, nn.ReLU, nn.LayerNorm and most other layers compute the same output in both modes.

A model made only of Linear and ReLU layers gives identical outputs in both modes:

plain = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 3))
x = torch.randn(4, 10)

plain.train(); a = plain(x)
plain.eval();  b = plain(x)
print(torch.equal(a, b))

True

Dropout in each mode

In training mode, nn.Dropout(p) sets each element of its input to zero with probability p, and multiplies the remaining elements by 1 / (1 - p). In evaluation mode it returns its input unchanged.

drop = nn.Dropout(p=0.5)
ones = torch.ones(10)

drop.train(); print(drop(ones))
drop.eval();  print(drop(ones))

Output:

tensor([0., 0., 2., 0., 0., 0., 2., 2., 0., 2.])
tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])

With p = 0.5, the surviving elements are multiplied by 2. This keeps the expected value of each element the same in both modes. Over one million ones, the mean of the training-mode output is 1.0004, and the mean of the evaluation-mode output is exactly 1.0. Because of this scaling, the layers after dropout receive inputs of the same average size in training and in evaluation, so no extra correction is needed at test time.

Since the dropped elements are chosen at random on every call, the same input gives a different output on every forward pass in training mode. For a small model with one Dropout(0.5) layer, ten forward passes of one input give:

train mode: 10 different outputs in 10 passes
eval mode:   1 different output  in 10 passes

BatchNorm in each mode

nn.BatchNorm1d normalizes each feature to mean 0 and variance 1, then applies a learned scale and shift. The difference between the modes is where the mean and variance come from.

Training mode: the mean and variance are computed from the current batch. The layer also updates two stored buffers, running_mean and running_var, as a moving average of the batch statistics.
Evaluation mode: the layer uses the stored running_mean and running_var, and does not change them.

The buffers start at mean 0 and variance 1. One forward pass in training mode moves them 10% of the way towards the batch statistics (the default momentum=0.1):

torch.manual_seed(0)
bn = nn.BatchNorm1d(3)
data = torch.randn(64, 3) * 2 + 5          # features with mean about 5, std about 2

print(data.mean(0))
bn.train(); _ = bn(data)
print(bn.running_mean)
print(bn.running_var)

tensor([5.0390, 4.4964, 5.5362])
tensor([0.5039, 0.4496, 0.5536])
tensor([1.2227, 1.2167, 1.3412])

0.5039 is 0.9 * 0 + 0.1 * 5.0390. The running statistics are updated by the forward pass. No backward() call or optimizer step is involved.

In training mode, one sample's output depends on the rest of the batch

Because training mode normalizes with the batch's own statistics, the output for a sample changes when the other samples in the batch change. Below, the same sample s is passed in two different batches, after the layer has seen 200 training batches:

train mode, s in batch A:  [ 0.4099, -0.3283,  0.1971]
train mode, s in batch B:  [-1.6630, -1.8497, -1.7948]
eval mode,  s in batch A:  [-0.1599, -0.2482, -0.1628]
eval mode,  s in batch B:  [-0.1599, -0.2482, -0.1628]
eval mode,  s alone:       [-0.1599, -0.2482, -0.1628]

In evaluation mode the output for s is the same in any batch and on its own, because the statistics are fixed.

A batch of one sample in training mode

A batch of one sample has no variance to compute, so BatchNorm1d refuses it in training mode:

bn.train()
bn(torch.tensor([[5.0, 5.0, 5.0]]))

ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 3])

This error usually appears when a single prediction is made on a model that is still in training mode, or when the last batch of an epoch has one sample. In the first case, the fix is model.eval(). In the second, DataLoader(..., drop_last=True) drops the incomplete final batch.

What happens when evaluation runs in training mode

The effect of a missing model.eval() can be measured directly.

A model with Dropout. This is a binary classifier with one Dropout(0.3) layer, trained for 20 epochs on 480 samples and evaluated on 120:

torch.manual_seed(0)
x = torch.randn(600, 10)
y = (x[:, 0] + x[:, 1] > 0).long()

model = nn.Sequential(
    nn.Linear(10, 32), nn.ReLU(), nn.Dropout(0.3),
    nn.Linear(32, 2))

After training, the same validation set was evaluated once in evaluation mode and five times in training mode:

eval mode:   val loss 0.0462   accuracy 0.992
train mode:  val loss 0.0520   accuracy 0.967
train mode:  val loss 0.0772   accuracy 0.967
train mode:  val loss 0.0728   accuracy 0.967
train mode:  val loss 0.0831   accuracy 0.967
train mode:  val loss 0.0520   accuracy 0.983

In training mode, the measured loss is higher and changes from run to run, because dropout removes a different 30% of the hidden units each time. The weights are the same in all six rows.

A model with BatchNorm. This is a 3-class classifier with BatchNorm1d and Dropout(0.3), trained for 30 epochs. Its validation accuracy was measured with different validation batch sizes:

validation batch size     train-mode accuracy    eval-mode accuracy
--------------------------------------------------------------------
600                              0.913                 0.932
64                               0.903                 0.932
8                                0.795                 0.932
2                                0.622                 0.932

In evaluation mode the accuracy does not depend on the batch size. In training mode it falls as the batch gets smaller, because the mean and variance of 2 or 8 samples are poor estimates of the statistics the model was trained with.

Validating in training mode also changes the model. After one pass over the validation set in training mode, num_batches_tracked went from 1140 to 1150 and the largest change in running_mean was 0.1975. The validation data has been mixed into the stored statistics, and these are the statistics the model will use later for real predictions.

What torch.no_grad() does

model.eval() changes how layers compute their output. It does not stop autograd. In evaluation mode, the forward pass still builds a computation graph:

model.eval()
out = model(x)
print(out.requires_grad, out.grad_fn)

with torch.no_grad():
    out = model(x)
print(out.requires_grad, out.grad_fn)

True <AddmmBackward0 object at 0x...>
False None

To compute gradients later, autograd stores some intermediate results of the forward pass, such as the input to each linear layer and the output of each ReLU. These stored tensors stay in memory until backward() runs or the output is deleted. Inside torch.no_grad(), nothing is recorded and nothing is stored.

For a network with layers 1024 -> 4096 -> 4096 -> 10, the activations stored for backward are:

batch size     stored with graph     stored under no_grad
---------------------------------------------------------
256                  17.0 MiB              0.0 MiB
4096                272.0 MiB              0.0 MiB

The stored memory grows with the batch size. On a GPU, this is the memory that a validation loop without no_grad() uses for no purpose, and it limits how large the validation batch can be.

A tensor produced under no_grad() has no graph, so calling backward() on it raises:

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

Why both are needed during evaluation

The two calls do not overlap. no_grad() does not switch layers to evaluation mode. A model in training mode, run under no_grad(), still applies dropout:

model.train()
with torch.no_grad():
    o1, o2 = model(x), model(x)
print(torch.allclose(o1, o2))

False

no_grad() also does not stop BatchNorm from updating its running statistics. Those updates happen in the forward pass and depend only on the training flag:

bn.train()
with torch.no_grad():
    _ = bn(data)          # running_mean still changes

And eval() does not stop the graph from being built, as shown above. The two calls are separate because there are cases where only one is wanted. For example, some interpretability methods compute gradients with respect to the input of a model in evaluation mode, which needs eval() without no_grad().

Call	What it changes	What it does not change
`model.train()`	sets `training=True`: dropout active, BatchNorm uses batch statistics and updates running statistics	gradient recording
`model.eval()`	sets `training=False`: dropout off, BatchNorm uses stored running statistics	gradient recording
`torch.no_grad()`	no graph is built, no activations are stored	layer behaviour, BatchNorm running statistics

torch.inference_mode()

torch.inference_mode() is a stricter version of no_grad(). It also skips some internal bookkeeping (version counters and view tracking), so it can be slightly faster. The cost is that tensors created inside it can never be used in a computation that autograd records:

model.eval()
with torch.inference_mode():
    out = model(x)                       # shape (4, 3)

w = torch.randn(3, requires_grad=True)
(out * w).sum().backward()

RuntimeError: Inference tensors cannot be saved for backward. Please do not use Tensors created in inference mode in computation tracked by autograd. ...

Use inference_mode() for prediction code whose outputs are not used for training. Use no_grad() in a validation loop or anywhere the results may be used in a later computation that needs gradients. Like no_grad(), inference_mode() does not change the training flag, so model.eval() is still required.

The train-and-evaluate loop

The training phase starts with model.train(). The validation phase starts with model.eval() and runs inside torch.no_grad().

for epoch in range(num_epochs):

    # training phase
    model.train()
    train_loss = 0.0
    for bx, by in train_loader:
        pred = model(bx)
        loss = loss_fn(pred, by)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * bx.size(0)

    # validation phase
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for bx, by in val_loader:
            pred = model(bx)
            val_loss += loss_fn(pred, by).item() * bx.size(0)

    t = train_loss / len(train_loader.dataset)
    v = val_loss / len(val_loader.dataset)
    print(f"Epoch {epoch:2d}  Train: {t:.4f}  Val: {v:.4f}")

model.train() is at the top of every epoch, not once before the loop. The validation phase leaves the model in evaluation mode, so without this line every epoch after the first would train with dropout off and with BatchNorm frozen at its running statistics.

On the Dropout example above, this loop gives:

Epoch  0  Train: 0.6110  Val: 0.4610
Epoch  5  Train: 0.1028  Val: 0.0940
Epoch 10  Train: 0.0489  Val: 0.0609
Epoch 15  Train: 0.0390  Val: 0.0574
Epoch 19  Train: 0.0367  Val: 0.0462

The training loss is computed in training mode, with dropout active, while the validation loss is computed in evaluation mode. The training loss is also averaged over the whole epoch, while the weights are still improving, and the validation loss is measured at the end of the epoch. Both reasons explain why the validation loss can be lower than the training loss.

For predictions after training, the pattern is the same as the validation phase:

model.eval()
with torch.no_grad():
    logits = model(new_x)
    preds = logits.argmax(dim=1)

Summary

model.train() and model.eval() set the training flag on the model and all its sub-modules. The flag only affects layers that read it, mainly Dropout and BatchNorm.
Dropout in training mode zeroes elements with probability p and scales the rest by 1 / (1 - p). In evaluation mode it does nothing.
BatchNorm in training mode uses the current batch's statistics and updates running_mean and running_var. In evaluation mode it uses the stored values and does not change them.
torch.no_grad() stops autograd from building a graph and storing activations. It does not change layer behaviour and does not stop BatchNorm's running statistics from updating.
Evaluation needs both: model.eval() and torch.no_grad(). Call model.train() at the start of every training epoch.
torch.inference_mode() is a stricter no_grad() for prediction code whose outputs will not be used with autograd.

This is one chapter's worth of an idea from my book, PyTorch From Ground Up, which builds everything from tensors upward so nothing stays vague. If it helped: 8 chapters are free, no email required, there's a free one-page tensor cheat-sheet here, every example runs in the companion notebooks on GitHub, and the full book is on Leanpub or in paperback and Kindle on Amazon.