Building models with nn.Module

What is nn.Module?

Why Subclass Instead of Functions?

Function approach (naive):

def my_layer(x, w, b):
    return x @ w + b

❌ Problems:

Parameters (w, b) must be passed explicitly everywhere
No automatic gradient tracking
Can't easily save/load models
No way to switch between train/eval modes

Module approach:

class MyLayer(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(in_features, out_features))
        self.bias = nn.Parameter(torch.randn(out_features))
    def forward(self, x):
        return x @ self.weight + self.bias

✅ Benefits:

Parameters stored inside, gradient tracking automatic
Can call .parameters() to get all learnable tensors
Can .to(device) to move everything
Composable with other modules

Core Architecture Pattern

The standard pattern for building a custom module:

import torch.nn as nn
 
class MyNetwork(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        # Step 1: Call parent constructor
        super().__init__()  # or super(MyNetwork, self).__init__()
        
        # Step 2: Define submodules and parameters
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.activation = nn.ReLU()
        self.layer2 = nn.Linear(hidden_dim, output_dim)
        
        # Step 3: (Optional) Custom parameters
        self.custom_param = nn.Parameter(torch.randn(10))
        
    def forward(self, x):
        # Step 4: Define computation graph
        x = self.layer1(x)
        x = self.activation(x)
        x = self.layer2(x)
        return x

Formula The Registration Mechanism

How does PyTorch know what parameters to train?

When you assign self.something = value inside __init__:

If value is an nn.Module, it's registered as a submodule
If value is an nn.Parameter, it's registered as a parameter
If value is an nn.ModuleList or nn.ModuleDict, all contained modules are registered Why this matters: The super().__init__() call sets up __setattr__ magic that hooks into attribute assignment. When you later call .parameters(), it recursively walks the submodule tree and collects all nn.Parameter objects.

Derivation of parameter collection:

Given: Module with submodules [M₁, M₂, ...] and parameters [P₁, P₂, ...]

Mparameters() = [P₁, P₂, ...] ∪ M₁.parameters() ∪ M₂.parameters() ∪ ...
Base case: If M has no submodules, M.parameters() = [Pᵢ] (just own parameters)
Recursive case: Include own parameters + all submodule parameters

Worked Example 1: Simple Feedforward Network

Task: Build a 3-layer feedforward network for binary classification.

import torch
import torch.nn as nn
 
class SimpleClassifier(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        
        # Why nn.Linear and not manual weight matrices?
        # nn.Linear handles initialization (Kaiming/Xavier) and bias correctly
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, output_dim)
        
        # Why ReLU as an attribute?
        # For modes like Dropout/BatchNorm, storing as attribute enables train/eval switching
        # For ReLU (stateless), it's style preference, but consistent pattern
        self.relu = nn.ReLU()
        
    def forward(self, x):
        # x: [batch_size, input_dim]
        # Why this step? Linear transform + nonlinearity
        x = self.relu(self.fc1(x))  # [batch_size, hidden_dim]
        # Why second hidden layer? Increases model capacity for complex boundaries
        x = self.relu(self.fc2(x))  # [batch_size, hidden_dim]
        
        # Why no activation on output? Loss function (BCEWithLogitsLoss) expects logits
        x = self.fc3(x)  # [batch_size, output_dim]
        
        return x
 
# Usage
model = SimpleClassifier(input_dim=784, hidden_dim=128, output_dim=1)
 
# Why this works: Forward is called automatically
input_data = torch.randn(32, 784)  # Batch of 32 images (28×28 flattened)
output = model(input_data)  # Calls model.forward(input_data) via __call__
 
print(f"Parameters: {sum(p.numel() for p in model.parameters())}")
# Output: Parameters: 117121
# Why? (784×128 + 128) + (128×128 + 128) + (128×1 + 1)
#     = 100480      + 16512       + 129        = 117121

Why this step? Breaking down parameter count:

Layer 1 (fc1): Weight matrix $(784 \times 128) = 100352$ parameters, bias $(128) = 128$ parameters → subtotal $100480$
Layer 2 (fc2): Weight matrix $(128 \times 128) = 16384$ parameters, bias $(128) = 128$ parameters → subtotal $16512$
Layer 3 (fc3): Weight matrix $(128 \times 1) = 128$ parameters, bias $(1) = 1$ parameter → subtotal $129$
Total: $100480 + 16512 + 129 = 117121$ parameters ✅

Worked Example 2: Residual Block with Skip Connections

Task: Implement a ResNet-style residual block: $\text{out} = F(x) + x$

class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        # Why two layers in F(x)? Standard residual block pattern
        self.conv1 = nn.Conv2d(dim, dim, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(dim)
        self.conv2 = nn.Conv2d(dim, dim, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(dim)
        self.relu = nn.ReLU(inplace=True)
        
    def forward(self, x):
        # Save input for skip connection
        identity = x  # [batch, dim, height, width]
        
        # Why this path? F(x) = BN(Conv(BN(Conv(x))))
        out = self.conv1(x)  # [batch, dim, H, W]
        out = self.bn1(out)
        out = self.relu(out)
        
        out = self.conv2(out)  # [batch, dim, H, W]
        out = self.bn2(out)
        # Why add before final ReLU? Allows gradient to flow through skip connection
        # Derivation: ∂Loss/∂x = ∂Loss/∂out × (∂F(x)/∂x + ∂identity/∂x)
        #                = ∂Loss/∂out × (∂F(x)/∂x + I)
        # The +I term prevents vanishing gradients!
        out = out + identity
        out = self.relu(out)
        
        return out

Why skip connections mathematically?

Without skip: $\frac{\partial L}{\partial x} = \frac{\partial L}{\partial \text{out}} \cdot \frac{\partial F(x)}{\partial x}$

With skip: $\frac{\partial L}{\partial x} = \frac{\partial L}{\partial \text{out}} \cdot \left(\frac{\partial F(x)}{\partial x} + 1\right)$

The "+1" ensures gradient magnitude≥ 1 even if $\frac{\partial F(x)}{\partial x} \approx 0$ , solving vanishing gradients in deep networks.

Worked Example 3: Variable-Length Module Lists

Task: Build a network with arbitrary depth specified at initialization.

class DeepNetwork(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, output_dim):
        super().__init__()
        
        # Why nn.ModuleList and not Python list?
        # Python list: layers = [nn.Linear(...), nn.Linear(...)]
        # Problem: Parameters NOT registered, won't be trained!
        
        # Why this pattern? First layer changes dimension
        self.input_layer = nn.Linear(input_dim, hidden_dim)
        # Why nn.ModuleList? Registers all contained modules automatically
        self.hidden_layers = nn.ModuleList([
            nn.Linear(hidden_dim, hidden_dim) 
            for _ in range(num_layers)
        ])
        
        self.output_layer = nn.Linear(hidden_dim, output_dim)
        self.relu = nn.ReLU()
        
    def forward(self, x):
        x = self.relu(self.input_layer(x))
        
        # Why iterate through ModuleList? Apply each hidden layer sequentially
        for layer in self.hidden_layers:
            x = self.relu(layer(x))
        
        x = self.output_layer(x)
        return x
 
# Verify registration
model = DeepNetwork(input_dim=10, hidden_dim=64, num_layers=5, output_dim=2)
print(f"Registered modules: {len(list(model.modules()))}")
# Output: Registered modules: 10
# Why? 1 (root) + 1 (input_layer) + 1 (ModuleList container) + 5 (hidden layers)
#      + 1 (output_layer) + 1 (ReLU) = 10

Common mistake Common Mistake: Using Python Lists for Layers

Wrong approach:

class BrokenNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        # This looks fine but is BROKEN
        self.layers = [nn.Linear(10, 10) for _ in range(3)]
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

Why this feels right: Python lists are natural for collections, and the forward pass runs fine.

Why it's wrong:

list is not a PyTorch container → layers not registered as submodules
model.parameters() returns EMPTY → optimizer sees no parameters → nothing trains
model.to(device) doesn't move these layers to GPU

The fix: Use nn.ModuleList or nn.ModuleDict for dynamic collections.

Steel-man the mistake: "But I can train each layer manually!" True, but you lose:

Automatic parameter collection
Unified state management (train/eval)
Model serialization (save/load)
The point of using PyTorch's abstractions

Parameter vs. Buffer Registration

class NormalizationLayer(nn.Module):
    def __init__(self, num_features):
        super().__init__()
        # Learnable parameters
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))
        
        # Non-learnable buffers (running statistics)
        # Why register_buffer? These should move with .to(device) and save with state_dict
        # but NOT be trained by optimizer
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))
        self.register_buffer('num_batches_tracked', torch.tensor(0))
        
    def forward(self, x):
        if self.training:
            # Update running statistics (exponential moving average)
            batch_mean = x.mean(dim=0)
            batch_var = x.var(dim=0)
            # Why 0.1 momentum? Balances responsiveness vs. stability
            self.running_mean = 0.9 * self.running_mean + 0.1 * batch_mean
            self.running_var = 0.9 * self.running_var + 0.1 * batch_var
            # Normalize using batch statistics
            x_norm = (x - batch_mean) / torch.sqrt(batch_var + 1e-5)
        else:
            # Why use running stats in eval? No batch statistics available at inference
            x_norm = (x - self.running_mean) / torch.sqrt(self.running_var + 1e-5)
        
        # Why affine transform? Allows network to undo normalization if needed
        return self.gamma * x_norm + self.beta

Why buffers matter:

.parameters() → only [gamma, beta] (2 tensors)
.state_dict() → includes gamma, beta, running_mean, running_var, num_batches_tracked (5 tensors)
Without register_buffer, running statistics would be lost when saving the model!

The Module Lifecycle

Formula Module Modes and Their Effects

Training Mode (model.train()):

Dropout: Randomly zeros elements with probability $p$
BatchNorm: Uses batch statistics, updates running averages
Gradient computation: Enabled (requires_grad=True) Evaluation Mode (model.eval()):
Dropout: Disabled (multiplies by keep probability instead)
BatchNorm: Uses running statistics (frozen)
Gradient computation: Still enabled unless wrapped in torch.no_grad()

Mathematical effect of Dropout mode: Training: $y_i = \begin{cases} 0 & \text{with probability } p \\ \frac{x_i}{1-p} & \text{with probability } 1-p \end{cases}$

Why divide by $(1-p)$ ? Expected value: $\mathbb{E}[y_i] = 0 \cdot p + \frac{x_i}{1-p} \cdot (1-p) = x_i$ (maintains scale)

Inference: $y_i = x_i$ (no dropout, scale already matched)

model = SimpleClassifier(784, 128, 10)
 
# Training loop
model.train()  # Sets self.training = True for all submodules
for batch in train_loader:
    optimizer.zero_grad()
    output = model(batch)  # Dropout active, BatchNorm updates running stats
    loss = criterion(output, targets)
    loss.backward()
    optimizer.step()
 
# Evaluation
model.eval()  # Sets self.training = False for all submodules
with torch.no_grad():  # Disables gradient computation (saves memory)
    for batch in test_loader:
        output = model(batch)  # Dropout off, BatchNorm uses running stats
        # No backward pass

Worked example Complete Training Script with nn.Module

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
# Define model
class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, 10)
        )
    def forward(self, x):
        return self.layers(x)
# Initialize
model = Classifier()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
# Why model.parameters()? Gives optimizer all learnable weights
# Alternative broken approach: optim.Adam([model.layers[0].weight]) 
# → Only trains first layer!
 
# Dummy data
X_train = torch.randn(1000, 784)
y_train = torch.randint(0, 10, (1000,))
train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32)
 
# Training
model.train()
for epoch in range(10):
    for batch_x, batch_y in train_loader:
        # Forward
        outputs = model(batch_x)  # Shape: [32, 10]
        loss = criterion(outputs, batch_y)
        
        # Backward
        optimizer.zero_grad()  # Why? Gradients accumulate by default
        loss.backward()
        optimizer.step()
 
# Save model
torch.save(model.state_dict(), 'model.pth')
# Load model
new_model = Classifier()
new_model.load_state_dict(torch.load('model.pth'))
new_model.eval()  # Don't forget this!

Advanced Pattern: Custom Initialization

class CustomInitNetwork(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        
        # Why custom initialization? Default may not suit your problem
        self._initialize_weights()
    def _initialize_weights(self):
        # Why iterate modules()? Gets all submodules recursively
        for m in self.modules():
            if isinstance(m, nn.Linear):
                # Xavier initialization for linear layers
                # Why Xavier? Maintains variance across layers
                # Derivation: For linear layer Y = WX, Var(Y) = n_in × Var(W) × Var(X)
                # To keep Var(Y) = Var(X), need Var(W) = 1/n_in
                nn.init.xavier_uniform_(m.weight)
                # Why constant bias? Start with no bias shift
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Conv2d):
                # He initialization for ReLU activations
                # Why He? ReLU zeros half the activations, needs √2 scaling
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

Initialization derivation:

For Xavier: Given $Y = WX$ where $W \in \mathbb{R}^{n_{\text{out}} \times n_{\text{in}}}$

$\text{Var}(Y_j) = \sum_{i=1}^{n_{\text{in}}} \text{Var}(W_{ji} X_i) = n_{\text{in}} \cdot \text{Var}(W) \cdot \text{Var}(X)$

To maintain $\text{Var}(Y) = \text{Var}(X)$ , need $\text{Var}(W) = \frac{1}{n_{\text{in}}}$ .

Uniform distribution: $W \sim U(-a, a)$ , $\text{Var}(W) = \frac{(2a)^2}{12} = \frac{a^2}{3}$

Set $\frac{a^2}{3} = \frac{1}{n_{\text{in}}} \Rightarrow a = \sqrt{\frac{3}{n_{\text{in}}}}$

Therefore: $W \sim U\left(-\sqrt{\frac{3}{n_{\text{in}}}}, \sqrt{\frac{3}{n_{\text{in}}}}\right)$

Recall Explain to a 12-year-old

Imagine you're building a LEGO castle. Each LEGO piece is an nn.Module:

Tiny bricks (like nn.Linear) are simple pieces that do one thing
Wall sections (like a ResidualBlock) are made by combining tiny bricks
Towers (like a CNN

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Hinglish (regional understanding)

Intuition Hinglish mein samjho

chalo, isko simple tareeke se samajhte hain. Dekho, PyTorch mein neural network banane ka jo fundamental building block hai, usko bolte hain nn.Module. Ise samajhne ka sabse aasaan tareeka hai Lego bricks wala example. Har ek module ek chota sa self-contained piece hai jo teen kaam karna jaanta hai: apne learnable parameters (jaise weights aur biases) store karna, input data pe ek transformation apply karna jise forward pass bolte hain, aur dusre modules ke saath jud kar bade complex networks banana. Core insight yeh hai ki neural network basically ek nested function hai jaise f(x) = f3(f2(f1(x))) — har chota function ek module hai, aur pura network bhi ek module hai. Isi composability wali baat mein saari khoobsurti chhupi hai.

Ab sawaal aata hai ki hum plain functions ki jagah nn.Module ko subclass kyun karte hain? Agar tum simple function likho jaise my_layer(x, w, b), toh problem yeh hai ki tumhe weights aur biases har baar manually pass karne padte hain, gradient tracking automatic nahi hoti, model save/load karna mushkil ho jaata hai, aur train vs eval mode switch karna toh bilkul aasaan nahi rehta. Lekin jab tum nn.Module use karte ho, toh saare parameters module ke andar hi store ho jaate hain, gradient tracking apne aap ho jaati hai, tum .parameters() call karke saare learnable tensors nikaal sakte ho, aur .to(device) se ek hi baar mein sab kuch GPU ya CPU pe move kar sakte ho. Yeh sab magic tab shuru hota hai jab tum super().__init__() call karte ho — yeh ek special __setattr__ mechanism set karta hai jo tumhare har attribute assignment ko track karta hai.

Toh yeh baat matter kyun karti hai? Kyunki jab tum __init__ ke andar self.layer1 = nn.Linear(...) likhte ho, PyTorch automatically samajh jaata hai ki yeh ek submodule hai aur ise register kar leta hai. Baad mein jab .parameters() call hota hai, toh yeh recursively pure submodule tree pe chal kar saare parameters collect kar leta hai. Iska matlab tumhe manually kabhi bhi yaad nahi rakhna padta ki kaun se weights train karne hain — framework khud handle kar leta hai. Yehi reason hai ki har real-world deep learning model, chahe simple classifier ho ya massive transformer, yeh sab nn.Module se hi inherit karte hain. Ek baar yeh pattern samajh gaye, toh koi bhi complex architecture banana bas Lego bricks jodne jitna aasaan lagega.

What is nn.Module?

Why Subclass Instead of Functions?

Core Architecture Pattern

Worked Example 1: Simple Feedforward Network

Worked Example 2: Residual Block with Skip Connections

Worked Example 3: Variable-Length Module Lists

Parameter vs. Buffer Registration

The Module Lifecycle

Advanced Pattern: Custom Initialization

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Deep Learning Frameworks