PyTorch nn.Module: Building Neural Networks

PyTorch nn.Module is the base class for every neural network. You subclass it, create layers in __init__, and describe how data flows through them in forward. PyTorch then tracks all learnable parameters for you, moves them between devices and saves them.

Your First nn.Module

import torch
from torch import nn

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(4, 8)
        self.out = nn.Linear(8, 3)

    def forward(self, x):
        x = torch.relu(self.hidden(x))
        return self.out(x)

model = TinyNet()
x = torch.randn(5, 4)          # batch of 5 samples, 4 features each
print(model(x).shape)

torch.Size([5, 3])

nn.Linear(4, 8) computes x @ W.T + b with a weight of shape [8, 4] and bias of shape [8]. Always call model(x) rather than model.forward(x): the call runs hooks and other bookkeeping.

Activation Functions

Without non-linear activations, stacked linear layers collapse into a single linear function. Common choices:

• nn.ReLU: max(0, x), simple and the usual default for hidden layers.
• nn.GELU and nn.SiLU: smooth variants used in transformers and modern CNNs.
• nn.Sigmoid: squashes to 0-1, used for binary outputs.
• nn.Softmax: turns a vector into probabilities, but usually left to the loss function.

nn.Sequential

For a simple stack of layers you can skip writing a class:

model = nn.Sequential(
    nn.Linear(4, 16),
    nn.ReLU(),
    nn.Linear(16, 16),
    nn.ReLU(),
    nn.Linear(16, 3),
)
print(model)

Sequential(
  (0): Linear(in_features=4, out_features=16, bias=True)
  (1): ReLU()
  (2): Linear(in_features=16, out_features=16, bias=True)
  (3): ReLU()
  (4): Linear(in_features=16, out_features=3, bias=True)
)

Inspecting Parameters

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

total = sum(p.numel() for p in model.parameters())
print("total params:", total)

0.weight (16, 4)
0.bias (16,)
2.weight (16, 16)
2.bias (16,)
4.weight (3, 16)
4.bias (3,)
total params: 403

The math: 4x16+16 = 80, 16x16+16 = 272, 16x3+3 = 51, which adds up to 403.

train() and eval() Modes

Some layers behave differently while training. nn.Dropout randomly zeros activations during training and does nothing in evaluation; nn.BatchNorm uses batch statistics in training and running averages in evaluation. Switch modes explicitly:

model.train()   # before training steps
model.eval()    # before validation or inference
print(model.training)

False

Frequently asked questions

What is nn.Module in PyTorch?

The base class for all neural network components. It registers layers and parameters, provides .to(device), .train()/.eval(), state_dict saving and more.

When should I use nn.Sequential instead of a custom nn.Module?

Use Sequential for a straight stack of layers. Write a custom Module when forward needs branches, skip connections, multiple inputs or any custom logic.

Why call model(x) instead of model.forward(x)?

model(x) runs __call__, which executes registered hooks and other internal logic before and after forward. Calling forward directly skips them.

How do I count the parameters in a PyTorch model?

Use sum(p.numel() for p in model.parameters()), adding if p.requires_grad to count only trainable ones.