Building models with nn.Module — Hinglish notes

3.3.3 · AI-ML › Deep Learning Frameworks

nn.Module kya hai?

Functions ki jagah Subclass kyun banayein?

Function approach (naive):

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

❌ Problems:

Parameters (w, b) ko har jagah explicitly pass karna padta hai
Automatic gradient tracking nahi hoti
Models ko easily save/load nahi kar sakte
Train/eval modes ke beech switch karne ka koi tarika nahi

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 andar stored hain, gradient tracking automatic hai
.parameters() call karke sabhi learnable tensors mil jaate hain
.to(device) se sab kuch move ho jaata hai
Doosre modules ke saath composable hai

Core Architecture Pattern

Custom module banane ka standard pattern:

import torch.nn as nn
 
class MyNetwork(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        # Step 1: Parent constructor call karo
        super().__init__()  # or super(MyNetwork, self).__init__()
        
        # Step 2: Submodules aur parameters define karo
        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: Computation graph define karo
        x = self.layer1(x)
        x = self.activation(x)
        x = self.layer2(x)
        return x

Formula Registration Mechanism

PyTorch ko kaise pata chalta hai ki kaunse parameters train karne hain?

Jab tum __init__ ke andar self.something = value assign karte ho:

Agar value ek nn.Module hai, toh use submodule ke roop mein register kiya jaata hai
Agar value ek nn.Parameter hai, toh use parameter ke roop mein register kiya jaata hai
Agar value ek nn.ModuleList ya nn.ModuleDict hai, toh usme contained sabhi modules register ho jaate hain Ye kyun important hai: super().__init__() call ek __setattr__ magic set up karta hai jo attribute assignment mein hook karta hai. Jab baad mein tum .parameters() call karte ho, toh ye recursively submodule tree walk karta hai aur sabhi nn.Parameter objects collect karta hai.

Parameter collection ka derivation:

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: Binary classification ke liye 3-layer feedforward network banao.

import torch
import torch.nn as nn
 
class SimpleClassifier(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        
        # nn.Linear kyun, manual weight matrices kyun nahi?
        # nn.Linear initialization (Kaiming/Xavier) aur bias sahi se handle karta hai
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, output_dim)
        
        # ReLU ko attribute ke roop mein kyun?
        # Dropout/BatchNorm jaise modes ke liye, attribute ke roop mein store karna train/eval switching enable karta hai
        # ReLU (stateless) ke liye, ye style preference hai, par consistent pattern hai
        self.relu = nn.ReLU()
        
    def forward(self, x):
        # x: [batch_size, input_dim]
        # Ye step kyun? Linear transform + nonlinearity
        x = self.relu(self.fc1(x))  # [batch_size, hidden_dim]
        # Doosra hidden layer kyun? Complex boundaries ke liye model capacity badhata hai
        x = self.relu(self.fc2(x))  # [batch_size, hidden_dim]
        
        # Output par activation kyun nahi? Loss function (BCEWithLogitsLoss) logits expect karta hai
        x = self.fc3(x)  # [batch_size, output_dim]
        
        return x
 
# Usage
model = SimpleClassifier(input_dim=784, hidden_dim=128, output_dim=1)
 
# Ye kyun kaam karta hai: Forward automatically call hota hai
input_data = torch.randn(32, 784)  # Batch of 32 images (28×28 flattened)
output = model(input_data)  # model.forward(input_data) ko __call__ ke zariye call karta hai
 
print(f"Parameters: {sum(p.numel() for p in model.parameters())}")
# Output: Parameters: 117121
# Kyun? (784×128 + 128) + (128×128 + 128) + (128×1 + 1)
#     = 100480      + 16512       + 129        = 117121

Ye step kyun? Parameter count ka breakdown:

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: ResNet-style residual block implement karo: $out = F (x) + x$

class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        # F(x) mein do layers kyun? Standard residual block pattern hai
        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):
        # Skip connection ke liye input save karo
        identity = x  # [batch, dim, height, width]
        
        # Ye path kyun? 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)
        # Final ReLU se pehle add kyun? Gradient ko skip connection ke through flow karne deta hai
        # Derivation: ∂Loss/∂x = ∂Loss/∂out × (∂F(x)/∂x + ∂identity/∂x)
        #                = ∂Loss/∂out × (∂F(x)/∂x + I)
        # +I term vanishing gradients ko rokta hai!
        out = out + identity
        out = self.relu(out)
        
        return out

Skip connections mathematically kyun?

Skip ke bina: $\frac{\partial L}{\partial x} = \frac{\partial L}{\partial out} \cdot \frac{\partial F ( x )}{\partial x}$

Skip ke saath: $\frac{\partial L}{\partial x} = \frac{\partial L}{\partial out} \cdot (\frac{\partial F ( x )}{\partial x} + 1)$

"+1" ensure karta hai ki gradient magnitude ≥ 1 rahe, chahe $\frac{\partial F ( x )}{\partial x} \approx 0$ ho, jo deep networks mein vanishing gradients solve karta hai.

Worked Example 3: Variable-Length Module Lists

Task: Ek aisa network banao jiska depth initialization par specify kiya ja sake.

class DeepNetwork(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, output_dim):
        super().__init__()
        
        # nn.ModuleList kyun, Python list kyun nahi?
        # Python list: layers = [nn.Linear(...), nn.Linear(...)]
        # Problem: Parameters register NAHI hote, train nahi honge!
        
        # Ye pattern kyun? Pehli layer dimension change karti hai
        self.input_layer = nn.Linear(input_dim, hidden_dim)
        # nn.ModuleList kyun? Contained sabhi modules automatically register karta hai
        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))
        
        # ModuleList ke through iterate kyun? Har hidden layer sequentially apply karo
        for layer in self.hidden_layers:
            x = self.relu(layer(x))
        
        x = self.output_layer(x)
        return x
 
# Registration verify karo
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
# Kyun? 1 (root) + 1 (input_layer) + 1 (ModuleList container) + 5 (hidden layers)
#      + 1 (output_layer) + 1 (ReLU) = 10

Common mistake Common Mistake: Layers ke liye Python Lists use karna

Galat approach:

class BrokenNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        # Ye theek lagta hai par BROKEN hai
        self.layers = [nn.Linear(10, 10) for _ in range(3)]
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

Ye sahi kyun lagta hai: Python lists collections ke liye natural hain, aur forward pass theek chalta hai.

Ye galat kyun hai:

list ek PyTorch container nahi hai → layers submodules ke roop mein register nahi hoti
model.parameters() EMPTY return karta hai → optimizer koi parameters nahi dekhta → kuch train nahi hota
model.to(device) in layers ko GPU par move nahi karta

Fix: Dynamic collections ke liye nn.ModuleList ya nn.ModuleDict use karo.

Galti ka steel-man: "Par main har layer ko manually train kar sakta hoon!" Sach hai, par tumhara loss hoga:

Automatic parameter collection
Unified state management (train/eval)
Model serialization (save/load)
PyTorch ke abstractions use karne ka purpose

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)
        # register_buffer kyun? Ye .to(device) ke saath move hone chahiye aur state_dict ke saath save hone chahiye
        # par optimizer dwara train NAHI hone chahiye
        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:
            # Running statistics update karo (exponential moving average)
            batch_mean = x.mean(dim=0)
            batch_var = x.var(dim=0)
            # 0.1 momentum kyun? Responsiveness aur stability ke beech balance
            self.running_mean = 0.9 * self.running_mean + 0.1 * batch_mean
            self.running_var = 0.9 * self.running_var + 0.1 * batch_var
            # Batch statistics se normalize karo
            x_norm = (x - batch_mean) / torch.sqrt(batch_var + 1e-5)
        else:
            # Eval mein running stats kyun? Inference par batch statistics available nahi hoti
            x_norm = (x - self.running_mean) / torch.sqrt(self.running_var + 1e-5)
        
        # Affine transform kyun? Network ko normalization undo karne ki suvidha deta hai agar zarurat ho
        return self.gamma * x_norm + self.beta

Buffers kyun important hain:

.parameters() → sirf [gamma, beta] (2 tensors)
.state_dict() → gamma, beta, running_mean, running_var, num_batches_tracked include karta hai (5 tensors)
register_buffer ke bina, running statistics model save karne par kho jaati!

Module Lifecycle

model = SimpleClassifier(784, 128, 10)
 
# Training loop
model.train()  # Sabhi submodules ke liye self.training = True set karta hai
for batch in train_loader:
    optimizer.zero_grad()
    output = model(batch)  # Dropout active, BatchNorm running stats update karta hai
    loss = criterion(output, targets)
    loss.backward()
    optimizer.step()
 
# Evaluation
model.eval()  # Sabhi submodules ke liye self.training = False set karta hai
with torch.no_grad():  # Gradient computation disable karta hai (memory bachata hai)
    for batch in test_loader:
        output = model(batch)  # Dropout off, BatchNorm running stats use karta hai
        # Koi backward pass nahi

Worked example nn.Module ke saath Complete Training Script

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
# Model define karo
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 karo
model = Classifier()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
# model.parameters() kyun? Optimizer ko sabhi learnable weights deta hai
# Alternative broken approach: optim.Adam([model.layers[0].weight]) 
# → Sirf pehli layer train hogi!
 
# 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()  # Kyun? Gradients by default accumulate hote hain
        loss.backward()
        optimizer.step()
 
# Model save karo
torch.save(model.state_dict(), 'model.pth')
# Model load karo
new_model = Classifier()
new_model.load_state_dict(torch.load('model.pth'))
new_model.eval()  # Ye mat bhuolo!

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)
        
        # Custom initialization kyun? Default tumhari problem ke liye suit nahi kar sakta
        self._initialize_weights()
    def _initialize_weights(self):
        # modules() iterate kyun? Sabhi submodules recursively milte hain
        for m in self.modules():
            if isinstance(m, nn.Linear):
                # Linear layers ke liye Xavier initialization
                # Xavier kyun? Layers ke across variance maintain karta hai
                # Derivation: Linear layer Y = WX ke liye, Var(Y) = n_in × Var(W) × Var(X)
                # Var(Y) = Var(X) rakhne ke liye, Var(W) = 1/n_in chahiye
                nn.init.xavier_uniform_(m.weight)
                # Constant bias kyun? Koi bias shift ke bina start karo
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Conv2d):
                # ReLU activations ke liye He initialization
                # He kyun? ReLU aadhe activations zero kar deta hai, √2 scaling chahiye
                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:

Xavier ke liye: Given $Y = W X$ jahan $W \in R^{n_{out} \times n_{in}}$

$Var (Y_{j}) = \sum_{i = 1}^{n_{in}} Var (W_{j i} X_{i}) = n_{in} \cdot Var (W) \cdot Var (X)$

$Var (Y) = Var (X)$ maintain karne ke liye, $Var (W) = \frac{1}{n _{in}}$ chahiye.

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

Set karo $\frac{a ^{2}}{3} = \frac{1}{n _{in}} \Rightarrow a = \frac{3}{n _{in}}$

Isliye: $W \sim U (- \frac{3}{n _{in}}, \frac{3}{n _{in}})$

Recall Ek 12 saal ke bacche ko samjhao

Socho tum ek LEGO castle bana rahe ho. Har LEGO piece ek nn.Module hai:

Chhoti bricks (jaise nn.Linear) simple pieces hain jo ek kaam karti hain
Wall sections (jaise ek ResidualBlock) chhoti bricks combine karke banti hain
Towers (jaise ek CNN