Neural Network Architecture with torch.nn

1. The Anatomy of an nn.Module Subclass

When designing custom architectures, you must implement two primary methods:

• __init__(self, ...): Calls super().__init__() and instantiates all stateful layers (Linear, Conv2d, LayerNorm, Embeddings). Any nn.Module assigned as an attribute onself is automatically registered as a model parameter!
• forward(self, x): Defines the forward dataflow and computation graph.Never call model.forward(x) directly in your code; always invoke model(x), which executes forward and backward hooks, handles profiling, and dispatches the call properly.

import torch
import torch.nn as nn

class MultiLayerPerceptron(nn.Module):
    def __init__(self, in_features: int, hidden_dim: int, num_classes: int, dropout_rate: float = 0.2):
        super().__init__()
        
        # 1. Feature projection block
        self.fc1 = nn.Linear(in_features, hidden_dim)
        self.norm = nn.LayerNorm(hidden_dim) # Stabilizes activations
        self.activation = nn.GELU()          # Modern Gaussian Error Linear Unit (used in GPT & BERT)
        self.dropout = nn.Dropout(p=dropout_rate)
        
        # 2. Classification output head
        self.fc2 = nn.Linear(hidden_dim, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Forward dataflow pipeline
        x = self.fc1(x)
        x = self.norm(x)
        x = self.activation(x)
        x = self.dropout(x)
        logits = self.fc2(x)
        return logits

# Instantiate model
model = MultiLayerPerceptron(in_features=784, hidden_dim=256, num_classes=10)
print(model)

2. Modern Layer & Activation Catalog

3. Parameter Inspection & state_dict

PyTorch provides clean APIs to inspect total learnable parameters and serialize model weights:

# 1. Count total trainable parameters:
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total Trainable Parameters: {total_params:,}")

# 2. Inspect named layers and shapes:
for name, param in model.named_parameters():
    print(f"Layer: {name:25} | Shape: {str(list(param.shape)):15} | Requires Grad: {param.requires_grad}")

# 3. Inspect state dictionary (weights & biases mapping):
print("\nState Dict Keys:")
for key in model.state_dict().keys():
    print(" -", key)

4. The Critical Distinction: model.train() vs model.eval()

# Always switch to evaluation mode before testing:
model.eval()

with torch.inference_mode():
    sample_input = torch.randn(1, 784)
    logits = model(sample_input)
    predicted_class = torch.argmax(logits, dim=-1).item()
    print(f"Predicted Class ID: {predicted_class}")