Building Networks with nn.Module

nn.Sequential is great for straight-line networks, but real architectures branch, reuse blocks, and take configuration arguments. For that, PyTorch has one universal pattern: subclass nn.Module. Every model you'll ever see — from a two-layer MLP to GPT — is written this way. Since it's built on Python classes, we'll start with a five-minute object-oriented programming (OOP) refresher, then build a complete classifier end to end.

Just enough OOP

A class is a blueprint; an object (instance) is one thing built from it. __init__ runs at construction time and stores data on self; other methods define behavior. This runs right in your browser:

The second idea you need is inheritance: a class can extend another, getting all its attributes and methods for free. super().__init__() runs the parent's constructor first, so the parent's setup happens before yours:

That's genuinely all the OOP you need: class, __init__, self, methods, and super().__init__(). Now look at how PyTorch uses exactly this pattern.

Anatomy of an nn.Module

A PyTorch model is a class that inherits from nn.Module and defines two things:

• __init__ — what parts exist: create the layers and store them on self.
• forward — how data flows: take the input, pass it through the parts, return the output.

import torch
from torch import nn
 
class MLP(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()                   # nn.Module's own setup — never skip
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size // 2)
        self.fc3 = nn.Linear(hidden_size // 2, output_size)
        self.relu = nn.ReLU()
 
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        return self.fc3(x)                   # raw logits — no softmax
 
model = MLP(input_size=8, hidden_size=16, output_size=2)
print(model)

Three details worth pausing on:

• super().__init__() must run before you assign any layers — it's what lets nn.Module detect and register the parameters you attach to self (that's how model.parameters() later finds them for the optimizer).
• Because the class takes arguments, the same blueprint builds networks of any size — try MLP(20, 64, 5).
• You call the model like a function — model(x) — never model.forward(x) directly. The call syntax runs important hooks around your forward.

End to end: classifying the moons dataset

Time to put MCO and nn.Module together on a real (toy) classification problem — two interleaved crescents that no straight line can separate. Run this in Colab:

import torch
from torch import nn, optim
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
 
device = "cuda" if torch.cuda.is_available() else "cpu"
 
# 1. Data: numpy -> scaled -> tensors
X, y = make_moons(n_samples=1000, noise=0.25, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=42
)
 
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)
 
X_train = torch.tensor(X_train, dtype=torch.float32).to(device)
X_test  = torch.tensor(X_test, dtype=torch.float32).to(device)
y_train = torch.tensor(y_train, dtype=torch.long).to(device)   # class labels: long!
y_test  = torch.tensor(y_test, dtype=torch.long).to(device)
 
# 2. MCO
class MLP(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size // 2),
            nn.ReLU(),
            nn.Linear(hidden_size // 2, output_size),
        )
 
    def forward(self, x):
        return self.net(x)
 
model = MLP(2, 32, 2).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
 
# 3. Train
for epoch in range(300):
    model.train()
    output = model(X_train)
    loss = criterion(output, y_train)
 
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
 
    if (epoch + 1) % 100 == 0:
        print(f"epoch {epoch+1} | train loss {loss.item():.4f}")
 
# 4. Evaluate
model.eval()
with torch.inference_mode():
    preds = model(X_test).argmax(1)
acc = (preds == y_test).float().mean().item()
print(f"test accuracy: {acc:.3f}")   # ~0.96-0.98

Note the rhythm: features become float32, class labels become long, model and data both move .to(device), logits come out, argmax(1) turns them into predicted classes. That skeleton carries you through the entire course.

How big should the hidden layers be?

There's no formula, but these rules of thumb serve well:

• Start small (one hidden layer, 16–64 units) and grow only if the model underfits — training loss stuck high.
• Funnel shapes work well: sizes shrinking toward the output (e.g. 64 → 32 → 16), compressing information stage by stage.
• Powers of two (16, 32, 64, 128...) are convention, not magic — they're just easy to reason about and hardware-friendly.
• More width/depth means more capacity — and more overfitting risk on small data. Which brings us to dropout.

Dropout: organized forgetting

Dropout randomly zeroes a fraction of a layer's activations during training (e.g. p=0.2 drops 20%). Each step, a different random subset of units vanishes, so no unit can rely on a specific neighbor — the network is forced to learn redundant, robust features instead of brittle co-adaptations. It's one of the cheapest and most effective ways to fight overfitting.

Placement: after the activation, on hidden layers only — never on the output:

self.net = nn.Sequential(
    nn.Linear(2, 32),
    nn.ReLU(),
    nn.Dropout(0.2),        # after the activation
    nn.Linear(32, 16),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(16, 2),       # output layer: no dropout
)

Dropout behaves differently in training and evaluation — it drops units while learning but must use all of them when predicting. That's exactly what model.train() and model.eval() toggle. Forgetting model.eval() before evaluation is a classic bug: your test predictions become randomly noisy and accuracy mysteriously fluctuates.

Next up: right now we feed the entire dataset through the model every epoch — fine for 1,000 points, impossible for 1,000,000 images. Minibatches, Dataset, and DataLoader fix that.