Building a Feedforward Network with nn.Module

nn.Module is PyTorch's base class for all neural network components. Every layer, activation function, loss function, and complete model in PyTorch subclasses nn.Module. It provides parameter management, device movement, serialisation, and training/eval mode switching out of the box. By subclassing it you get all this functionality for free and only need to define the architecture and the forward pass.

Creating a custom network requires two methods. In __init__ you call super().__init__() and define all learnable layers as attributes. In forward you describe how data flows through those layers. PyTorch tracks any nn.Module or nn.Parameter assigned to self as a trainable component. Calling the model like a function (model(x)) automatically invokes forward and runs any registered hooks.

nn.Linear(in_features, out_features) applies the transformation y = x @ W.T + b where W is a weight matrix and b is a bias vector. The layer automatically initialises weights with kaiming uniform distribution and biases with uniform distribution. It is the building block of feedforward networks, also called fully connected or dense layers.

Activation functions introduce non-linearity, allowing the network to learn complex patterns that a stack of linear layers cannot represent. ReLU (max(0, x)) is the default for hidden layers — fast and avoids vanishing gradients. Sigmoid squashes output to [0, 1], used for binary classification outputs. Tanh squashes to [-1, 1], commonly used in RNNs. All are available as modules in nn.

nn.Sequential is a convenient container that passes the output of each module as the input to the next. It is ideal for simple feedforward architectures where data flows linearly. For more complex networks with skip connections, multiple inputs, or branching paths, you need the full nn.Module subclass approach. Sequential networks can still be extended by subclassing and using the Sequential as a sub-block.

A critical debugging technique when building networks is to run a dummy tensor through the model before training. This verifies that all layer dimensions are compatible and reveals shape mismatches immediately. The dummy tensor has the same shape as your real data but contains random values — it just exercises the forward path. Always do this after defining a new architecture.

Understanding the total number of trainable parameters in a model gives you a sense of its capacity and memory requirements. model.parameters() returns an iterator over all learnable tensors; model.named_parameters() pairs each tensor with its name for inspection. A compact parameter counter is one of the first utilities every PyTorch practitioner builds.

Moving the model to the GPU is as simple as calling model.to(device). This transfers all parameters and buffers registered in the module to the target device. After this call, all forward pass computations happen on the GPU automatically — as long as your input tensors are also on the same device. Mixing CPU and GPU tensors raises a runtime error.

Some layers (Dropout, BatchNorm) behave differently during training and inference. Calling model.train() enables stochastic behaviour (random dropout, batch statistics); model.eval() switches to deterministic inference behaviour. Forgetting to call model.eval() before evaluation leads to inconsistent results because Dropout randomly zeros activations. Always pair these calls with torch.no_grad() during inference for maximum efficiency.

Putting it all together: a multi-layer perceptron (MLP) is a feedforward network with one or more hidden layers between input and output. Each hidden layer applies a linear transformation followed by a non-linear activation. The output layer's activation depends on the task — no activation for regression, softmax for multi-class classification, sigmoid for binary classification. The example below builds a 3-layer MLP for 10-class classification.

Neural network training is expensive, so you save the model after training. PyTorch's convention is to save only the state_dict — a dictionary of parameter tensors — rather than the entire model object. This avoids pickle dependency on the class definition. To restore, create a fresh model instance with the same architecture, then load the state dict with load_state_dict. Always set model.eval() after loading for inference.

Test your understanding of Machine Learning with Python concepts from this lesson.

In this lesson you learned: nn.Module is the base class for all PyTorch neural networks, requiring __init__ to define layers and forward to define data flow, nn.Sequential provides a simple container for linear stacks of layers, and model.train() / model.eval() switches behaviour of layers like Dropout and BatchNorm. Next up we write the complete training loop including loss, optimizer, and multiple epochs.