Skip to content

← nn.Module Mechanics step 2 of 7

Easy Primitives

Compare dropout in training and evaluation modes

Jump to the task

model.train() enables training behavior. model.eval() is equivalent to model.train(False): it sets the training flag to False on the module and its registered children. Each layer decides how to use that flag.

Dropout gives a simple example. With probability one, its training output is always zero, so there is no randomness to hide the difference:

import torch
import torch.nn as nn

model = nn.Sequential(nn.Dropout(p=1.0))
x = torch.ones(3)
model.train()
print(model(x).tolist())  # [0.0, 0.0, 0.0]
model.eval()
print(model(x).tolist())  # [1.0, 1.0, 1.0]
print(model[0].training)  # False

Layers that change behavior

Dropout usually uses a probability between zero and one. During training it randomly zeroes elements and scales the survivors by 1/(1-p), preserving the expected value. In evaluation it returns the input unchanged. The p=1 case above is a special endpoint that returns zeros in training.

BatchNorm, with its default track_running_stats=True, uses current-batch statistics during training and updates running estimates. In evaluation it uses those estimates. Leaving it in training mode during validation can change predictions when the batch composition changes and update its stored statistics with validation data. If track_running_stats=False, BatchNorm uses batch statistics in both modes.

Other built-in and custom modules can also use the flag. Linear layers do not change their calculation simply because you call eval().

Evaluation mode does not disable gradients

Module mode and gradient recording are separate controls:

layer = nn.Linear(2, 1).eval()
x = torch.ones(1, 2)
print(layer(x).requires_grad)  # True: the weight and bias require gradients
with torch.no_grad():
    print(layer(x).requires_grad)  # False

Use eval() with no_grad() for ordinary validation that needs no derivatives. Evaluation with gradients can also be intentional, such as when computing input sensitivity. Conversely, no_grad() alone does not switch off a Dropout layer's training behavior or stop BatchNorm's running-stat updates. Call train() again before resuming ordinary training after validation.

Your task

def both_modes(n: int) -> dict

Build a module containing Dropout(p=1.0). Run a vector of n ones through it in training mode, then in evaluation mode. Capture the training flag after each switch, rather than reading both flags after the final switch.

Return train_values, train_flag, eval_values and eval_flag. For n=2, the values are [0.0, 0.0], True, [1.0, 1.0] and False, respectively. Return the tensor values as lists.