PyTorch Neural Networks with nn.Module

Real-World Scenario

A computer vision engineer at a retail company needs to build a product classification model. They use PyTorch’s nn.Module to define a multi-layer network, DataLoader to batch their product images, and the standard training loop to minimize cross-entropy loss. This is the workflow behind every production PyTorch model, and each piece of it fits in a small example.

The nn.Module API

Subclass nn.Module, create layers in __init__, and describe how data flows through them in forward.

import torch
import torch.nn as nn
import torch.nn.functional as F

class TwoLayerNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)   # assigning a layer registers it
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        self.dropout = nn.Dropout(p=0.3)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        return self.fc2(x)            # raw logits: the loss applies softmax

model = TwoLayerNet(784, 256, 10)
print(model)
print(f"parameters: {sum(p.numel() for p in model.parameters()):,}")
print(model(torch.randn(32, 784)).shape)

$ python module.py
TwoLayerNet(
  (fc1): Linear(in_features=784, out_features=256, bias=True)
  (fc2): Linear(in_features=256, out_features=10, bias=True)
  (dropout): Dropout(p=0.3, inplace=False)
)
parameters: 203,530
torch.Size([32, 10])

Printing a module lists every layer assigned as an attribute, in the order they were assigned — not the order forward uses them. The parameter count is 784·256 + 256 for fc1 plus 256·10 + 10 for fc2: 203,530. Calling model(x) rather than model.forward(x) runs forward through PyTorch’s hooks, and the output has one row of 10 logits per input.

Only Registered Layers Are Parameters

nn.Module finds parameters by looking at attributes. A layer it cannot see is a layer the optimizer never trains.

import torch.nn as nn

class PlainList(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [nn.Linear(4, 4) for _ in range(3)]      # plain Python list

class Registered(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList(nn.Linear(4, 4) for _ in range(3))

for cls in (PlainList, Registered):
    params = list(cls().parameters())
    print(f"{cls.__name__:10s} {len(params)} parameter tensors")

for name, p in Registered().named_parameters():
    print(name, tuple(p.shape))

$ python registration.py
PlainList  0 parameter tensors
Registered 6 parameter tensors
layers.0.weight (4, 4)
layers.0.bias (4,)
layers.1.weight (4, 4)
layers.1.bias (4,)
layers.2.weight (4, 4)
layers.2.bias (4,)

The three nn.Linear layers inside a plain Python list are invisible: parameters() returns nothing, so an optimizer built from it would update nothing, and no error would say so. nn.ModuleList registers each layer under an indexed name (layers.0.weight), which is also the key used in the saved state_dict.

train() and eval() Change What Layers Do

import torch
import torch.nn as nn

torch.manual_seed(0)
layer = nn.Dropout(p=0.5)
x = torch.ones(8)

print(layer.training)
print(layer(x))       # train mode: random zeros, survivors scaled by 1/(1-p)
print(layer(x))       # a different mask on every call

layer.eval()
print(layer(x))       # eval mode: identity

$ python dropout_modes.py
True
tensor([0., 0., 2., 0., 0., 0., 2., 2.])
tensor([0., 2., 2., 2., 2., 0., 2., 2.])
tensor([1., 1., 1., 1., 1., 1., 1., 1.])

A new module starts in training mode. There, Dropout(p=0.5) zeroes each value with probability 0.5 and multiplies the survivors by 1 / (1 - 0.5) = 2, so the expected sum stays the same — and draws a new mask on every call. After eval() it passes input through untouched. Forgetting model.eval() before evaluation gives noisy, pessimistic numbers.

nn.Sequential — Layers Without Boilerplate

When data flows straight through a stack of layers, nn.Sequential saves you writing forward.

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(784, 128), nn.BatchNorm1d(128), nn.ReLU(), nn.Dropout(0.3),
    nn.Linear(128, 10),
)
print(model(torch.randn(16, 784)).shape)

single = torch.randn(1, 784)
try:
    model(single)
except ValueError as e:
    print("train mode:", e)

model.eval()
print("eval mode:", model(single).shape)

$ python batchnorm_one.py
torch.Size([16, 10])
train mode: Expected more than 1 value per channel when training, got input size torch.Size([1, 128])
eval mode: torch.Size([1, 10])

A batch of 16 works. A single example fails in training mode, because BatchNorm1d normalizes each feature using the batch’s mean and variance, and one value has no variance. In eval() mode BatchNorm uses running statistics instead, so single-example inference works — another reason to always switch modes before predicting.

Loss Functions

import torch
import torch.nn as nn

logits = torch.tensor([[2.0, 0.5, -1.0],
                       [0.1, 0.2, 3.0]])
targets = torch.tensor([0, 2])          # class indices, not one-hot

ce = nn.CrossEntropyLoss()
print(f"logits into CE:  {ce(logits, targets).item():.4f}")
print(f"softmax into CE: {ce(logits.softmax(dim=1), targets).item():.4f}")

bce = nn.BCEWithLogitsLoss()
print(f"BCE with logits: {bce(torch.tensor([3.0, -2.0]), torch.tensor([1.0, 0.0])).item():.4f}")

mse = nn.MSELoss()
print(f"MSE:             {mse(torch.tensor([2.5, 0.0]), torch.tensor([3.0, -0.5])).item():.4f}")

$ python losses.py
logits into CE:  0.1755
softmax into CE: 0.6611
BCE with logits: 0.0878
MSE:             0.2500

CrossEntropyLoss takes raw logits and integer class indices. Feeding it softmax probabilities doesn’t raise an error; it applies softmax a second time and reports 0.6611 instead of 0.1755 for the same predictions — a flatter loss with weaker gradients. BCEWithLogitsLoss is the binary equivalent and also wants raw scores; MSELoss averages squared errors: (0.5² + 0.5²) / 2 = 0.25.

What One Optimizer Step Does

Every training step is the same three calls: clear gradients, backpropagate, step.

import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Linear(3, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

x = torch.tensor([[1.0, 2.0, 3.0]])
target = torch.tensor([[1.0]])

before = model.weight.detach().clone()
optimizer.zero_grad()
loss = nn.functional.mse_loss(model(x), target)
loss.backward()
optimizer.step()

print("loss:   ", round(loss.item(), 4))
print("grad:   ", model.weight.grad)
print("change: ", model.weight.detach() - before)
print("-lr*grad", -0.1 * model.weight.grad)

$ python one_step.py
loss:    4.9968
grad:    tensor([[ -4.4707,  -8.9414, -13.4120]])
change:  tensor([[0.4471, 0.8941, 1.3412]])
-lr*grad tensor([[0.4471, 0.8941, 1.3412]])

For plain SGD, the weight change is exactly -lr × grad, as the last two lines show. The gradient is proportional to the input [1, 2, 3], so the weight attached to the largest input moved furthest. Adam and AdamW scale each step using running averages of past gradients, but the three calls are identical.

Learning Rate Schedulers

import torch
import torch.nn as nn

model = nn.Linear(10, 1)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
step = torch.optim.lr_scheduler.StepLR(opt, step_size=3, gamma=0.5)

for epoch in range(8):
    # ... one epoch of training would run here ...
    opt.step()
    print(f"epoch {epoch}: lr {step.get_last_lr()[0]:.6f}")
    step.step()

$ python schedulers.py
epoch 0: lr 0.001000
epoch 1: lr 0.001000
epoch 2: lr 0.001000
epoch 3: lr 0.000500
epoch 4: lr 0.000500
epoch 5: lr 0.000500
epoch 6: lr 0.000250
epoch 7: lr 0.000250

StepLR(step_size=3, gamma=0.5) halves the learning rate every three scheduler.step() calls. Call it once per epoch, after the optimizer has stepped — calling it before optimizer.step() makes PyTorch warn that the two calls are in the wrong order, and the first value of the schedule is skipped.

The Standard Training Loop

Putting it together: a DataLoader for batches, a training phase, and an evaluation phase on held-out data.

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

torch.manual_seed(42)
X = torch.randn(2000, 20)
y = (X[:, 0] + X[:, 1] > 0).long()           # label depends on two features
train = DataLoader(TensorDataset(X[:1600], y[:1600]), batch_size=32, shuffle=True)
val_X, val_y = X[1600:], y[1600:]

model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Dropout(0.2), nn.Linear(64, 2))
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
criterion = nn.CrossEntropyLoss()

for epoch in range(1, 16):
    model.train()
    total = 0.0
    for xb, yb in train:
        optimizer.zero_grad()                  # 1. clear old gradients
        loss = criterion(model(xb), yb)        # 2. forward
        loss.backward()                        # 3. backward
        optimizer.step()                       # 4. update
        total += loss.item() * len(yb)
    model.eval()
    with torch.no_grad():
        val_logits = model(val_X)
        val_loss = criterion(val_logits, val_y).item()
        val_acc = (val_logits.argmax(1) == val_y).float().mean().item()
    if epoch % 3 == 0:
        print(f"epoch {epoch:2d}  train_loss {total / 1600:.4f}  val_loss {val_loss:.4f}  val_acc {val_acc:.3f}")

$ python train_loop.py
epoch  3  train_loss 0.3544  val_loss 0.3116  val_acc 0.928
epoch  6  train_loss 0.1738  val_loss 0.1695  val_acc 0.957
epoch  9  train_loss 0.1134  val_loss 0.1200  val_acc 0.978
epoch 12  train_loss 0.0897  val_loss 0.1016  val_acc 0.970
epoch 15  train_loss 0.0636  val_loss 0.0887  val_acc 0.973

The four numbered lines are the whole training step. The evaluation phase switches to model.eval() and wraps the forward pass in torch.no_grad(): the first turns dropout off, the second skips building a graph. Validation accuracy climbs to about 97% because the label is a simple function of two of the twenty features. Validation loss sits below training loss early on for two reasons: dropout is active while the training loss is measured, and the training loss is averaged over an epoch during which the model was still improving.

Saving and Loading Weights

Save the state_dict — a dictionary of named tensors — rather than the model object.

import torch
import torch.nn as nn

def make_model():
    return nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))

torch.manual_seed(0)
model = make_model()
print(list(model.state_dict().keys()))

torch.save(model.state_dict(), "weights.pt")

restored = make_model()                      # same architecture, fresh random weights
restored.load_state_dict(torch.load("weights.pt", weights_only=True))
restored.eval()

x = torch.randn(2, 10)
print(torch.equal(model(x), restored(x)))

wrong = nn.Sequential(nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 1))
try:
    wrong.load_state_dict(torch.load("weights.pt", weights_only=True))
except RuntimeError as e:
    print(str(e).splitlines()[1].strip())

$ python save_load.py
['0.weight', '0.bias', '2.weight', '2.bias']
True
size mismatch for 0.weight: copying a param with shape torch.Size([32, 10]) from checkpoint, the shape in current model is torch.Size([64, 10]).

The keys come from the layer positions inside nn.Sequential (0, 2; the ReLU at 1 has no parameters). Loading them into a freshly built model of the same architecture reproduces the original’s outputs exactly. Loading into a different architecture fails with a precise size mismatch — the state_dict holds only tensors, so you must rebuild the same model in code first. weights_only=True makes torch.load refuse arbitrary pickled objects, which is the safe default for files you did not create.

Checkpoints for Resuming Training

import torch
import torch.nn as nn

model = nn.Linear(4, 1)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
nn.functional.mse_loss(model(torch.randn(8, 4)), torch.randn(8, 1)).backward()
optimizer.step()                              # Adam now holds moment estimates

torch.save({"epoch": 7,
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict()}, "checkpoint.pt")

ckpt = torch.load("checkpoint.pt", weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
print("resume at epoch", ckpt["epoch"] + 1)
print(sorted(ckpt["optimizer"]["state"][0].keys()))

$ python checkpoint.py
resume at epoch 8
['exp_avg', 'exp_avg_sq', 'step']

To resume training rather than just predict, save the optimizer too. Adam keeps two running averages (exp_avg, exp_avg_sq) and a step count for every parameter; restarting without them resets its adaptive step sizes, so the first steps after a resume behave like the start of training rather than a continuation.

Practice

import torch.nn as nn

model = nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10))
for p in model[0].parameters():
    p.requires_grad = False              # freeze the first layer

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"trainable {trainable:,} of {total:,}")

$ python p1_freeze.py
trainable 2,570 of 203,530

import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Dropout(0.5), nn.Linear(64, 2))
x = torch.randn(1, 20)

with torch.no_grad():                       # no_grad does NOT switch dropout off
    for _ in range(2):
        print("train mode:", [round(v, 4) for v in model(x).squeeze().tolist()])
    model.eval()
    for _ in range(2):
        print("eval mode: ", [round(v, 4) for v in model(x).squeeze().tolist()])

$ python p2_forgot_eval.py
train mode: [0.1892, 0.163]
train mode: [0.6342, -0.5095]
eval mode:  [0.57, -0.1991]
eval mode:  [0.57, -0.1991]

import torch
import torch.nn as nn

opt = torch.optim.AdamW(nn.Linear(10, 1).parameters(), lr=1e-3)
cosine = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=8)

for epoch in range(9):
    opt.step()
    print(f"epoch {epoch}: lr {cosine.get_last_lr()[0]:.6f}")
    cosine.step()

$ python p3_cosine.py
epoch 0: lr 0.001000
epoch 1: lr 0.000962
epoch 2: lr 0.000854
epoch 3: lr 0.000691
epoch 4: lr 0.000500
epoch 5: lr 0.000309
epoch 6: lr 0.000146
epoch 7: lr 0.000038
epoch 8: lr 0.000000

import torch
import torch.nn as nn

old = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
torch.save(old.state_dict(), "old.pt")

new = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1), nn.Sigmoid(),
                    nn.Linear(1, 1))
result = new.load_state_dict(torch.load("old.pt", weights_only=True), strict=False)
print("missing:   ", result.missing_keys)
print("unexpected:", result.unexpected_keys)
print(torch.equal(new[0].weight, old[0].weight))

$ python p4_strict.py
missing:    ['4.weight', '4.bias']
unexpected: []
True