Building a model with nn.Module

The idea

In the autograd module you managed w and b by hand. A real network has many layers and thousands of parameters, so PyTorch provides building blocks in torch.nn:

• layers (nn.Linear, nn.Conv2d...) that own their parameters,
• activation functions (nn.ReLU, nn.Sigmoid...),
• and nn.Module, the base class you subclass to describe your model: which layers it has (in __init__) and how data flows through them (in forward).

It's the classes module from the Python track, put to work.

A Linear layer is X @ W.T + b

The basic building block of a neural network is the linear layer. It's the same "multiply and add" as a linear regression, just with several outputs at once. Here's one that takes 3 features and produces 2 numbers:

import torch
from torch import nn

torch.manual_seed(0)
layer = nn.Linear(in_features=3, out_features=2)
print(layer.weight.shape, layer.bias.shape)

x = torch.tensor([[1.0, 0.0, 2.0]])          # a batch of 1 example with 3 features
print(layer(x))
print(x @ layer.weight.T + layer.bias)       # the same calculation by hand

torch.Size([2, 3]) torch.Size([2])
tensor([[-0.9661,  0.3425]], grad_fn=<AddmmBackward0>)
tensor([[-0.9661,  0.3425]], grad_fn=<AddBackward0>)

The grad_fn in the output shows the result is part of the autograd graph, ready for backward(). nn.Linear(3, 2) holds a weight matrix of shape (2, 3) and a bias of shape (2,), both created randomly and marked requires_grad=True. Calling the layer on a batch computes x @ W.T + b: the "layer by hand" exercise from the operations module.

Activations: the non-linear part

Stacking linear layers alone gives just another linear function (a matrix times a matrix is a matrix). Between layers you put a non-linear activation, which lets the network bend:

z = torch.tensor([-2.0, -0.5, 0.0, 1.5])
print(torch.relu(z))         # max(0, z): the default choice
print(torch.sigmoid(z))      # squashes to (0, 1): for probabilities
print(torch.tanh(z))         # squashes to (-1, 1)

tensor([0.0000, 0.0000, 0.0000, 1.5000])
tensor([0.1192, 0.3775, 0.5000, 0.8176])
tensor([-0.9640, -0.4621,  0.0000,  0.9051])

ReLU (rectified linear unit) is the standard between hidden layers: simple, fast, and it trains well.

Your first model class

A churn model with 5 inputs, one hidden layer of 16 units, and one output:

class ChurnNet(nn.Module):
    def __init__(self, n_features, hidden=16):
        super().__init__()
        self.hidden = nn.Linear(n_features, hidden)
        self.output = nn.Linear(hidden, 1)

    def forward(self, x):
        x = torch.relu(self.hidden(x))
        return self.output(x)          # a raw score ("logit"); sigmoid turns it into a probability


torch.manual_seed(0)
model = ChurnNet(n_features=5)
batch = torch.randn(4, 5)              # 4 customers, 5 features each
logits = model(batch)
print(logits.shape)
print(torch.sigmoid(logits).squeeze().detach().round(decimals=3))

torch.Size([4, 1])
tensor([0.3920, 0.4280, 0.3520, 0.4480])

• __init__ creates the layers and must call super().__init__() first.
• forward(x) describes the computation. You call model(batch), not model.forward(batch); PyTorch wraps the call with extra machinery (hooks).
• The output is a logit (raw score). Keeping the sigmoid out of the model and inside the loss function is more numerically stable (next module).

The untrained model's probabilities sit between 0.35 and 0.45: random weights, no knowledge yet. (.detach() drops the autograd bookkeeping before printing.)

nn.Sequential for simple stacks

When data just flows through layers in order, nn.Sequential saves writing a class:

torch.manual_seed(0)
model = nn.Sequential(
    nn.Linear(5, 16),
    nn.ReLU(),
    nn.Linear(16, 1),
)
print(model)
print(model(batch).shape)

Sequential(
  (0): Linear(in_features=5, out_features=16, bias=True)
  (1): ReLU()
  (2): Linear(in_features=16, out_features=1, bias=True)
)
torch.Size([4, 1])

Use nn.Sequential for straight stacks; write an nn.Module subclass when forward needs logic (several inputs, skip connections, branching).

Parameters

Every layer registered as an attribute is tracked automatically:

for name, p in model.named_parameters():
    print(name, tuple(p.shape))
print("total:", sum(p.numel() for p in model.parameters()))

0.weight (16, 5)
0.bias (16,)
2.weight (1, 16)
2.bias (1,)
total: 113

5×16 + 16 + 16×1 + 1 = 113 parameters. model.parameters() is what you hand to an optimizer, so it can update them all. This works because the layers were assigned to self.something in __init__; a layer stored in a plain Python list wouldn't be found (use nn.ModuleList for lists of layers).

Why this matters for clean code

Model classes are where deep-learning code stays readable or turns into spaghetti. Keep __init__ declarative (just the layers, with clear names like hidden and output), keep forward short and linear, pass sizes as arguments (n_features, hidden) instead of hard-coding them, and return logits rather than mixing probability logic into the model. Then the same class can be tested, reused and reconfigured.

Recap

• nn.Linear(in, out) computes x @ W.T + b with weight (out, in) and bias (out,).
• Non-linear activations (ReLU between layers; sigmoid for probabilities) let networks learn curves.
• Subclass nn.Module: layers in __init__ (after super().__init__()), computation in forward; call model(x).
• nn.Sequential for simple stacks; a custom class when forward needs logic.
• model.parameters() / named_parameters() list every learnable tensor.