PyTorch Neural Networks with nn.Module

Real-World Scenario

A computer vision engineer at a retail company needs to build a product classification model. They use PyTorch’s nn.Module to define a multi-layer network, DataLoader to batch their 50,000 product images, and the standard training loop to minimize cross-entropy loss. This is the workflow behind every production PyTorch model.

The nn.Module API

import torch
import torch.nn as nn
import torch.nn.functional as F

class TwoLayerNet(nn.Module):
    """A simple two-layer feedforward network."""

    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
        super().__init__()
        # Layers are registered automatically as submodules
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        self.dropout = nn.Dropout(p=0.3)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Define the forward pass — autograd builds the computation graph
        x = F.relu(self.fc1(x))   # linear → ReLU activation
        x = self.dropout(x)        # randomly zero 30% of neurons during training
        x = self.fc2(x)            # final linear layer (no activation — loss handles it)
        return x


# Instantiate
model = TwoLayerNet(input_dim=784, hidden_dim=256, output_dim=10)
print(model)

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

# Single forward pass
x = torch.randn(32, 784)   # batch of 32 samples
output = model(x)           # calls forward() internally
print(output.shape)         # (32, 10)

nn.Sequential — Layers Without Boilerplate

import torch
import torch.nn as nn

# Quick way to stack layers when you don't need branching
model = nn.Sequential(
    nn.Linear(784, 512),
    nn.BatchNorm1d(512),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(512, 256),
    nn.BatchNorm1d(256),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(256, 10),
)

x = torch.randn(16, 784)
print(model(x).shape)  # (16, 10)

Optimizers and Loss Functions

import torch
import torch.nn as nn
import torch.optim as optim

model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))

# Loss functions
criterion_ce  = nn.CrossEntropyLoss()     # multi-class classification
criterion_bce = nn.BCEWithLogitsLoss()    # binary classification (numerically stable)
criterion_mse = nn.MSELoss()              # regression

# Optimizers
optimizer_sgd  = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)
optimizer_adam = optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-8)
optimizer_adamw = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)  # Adam + decoupled decay

# Learning rate schedulers
scheduler_step = optim.lr_scheduler.StepLR(optimizer_adam, step_size=10, gamma=0.5)
scheduler_cos  = optim.lr_scheduler.CosineAnnealingLR(optimizer_adam, T_max=100)
scheduler_reduce = optim.lr_scheduler.ReduceLROnPlateau(optimizer_adam, patience=5, factor=0.5)

The Standard Training Loop

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Synthetic dataset — replace with real data
X = torch.randn(2000, 20)
y = (X[:, 0] + X[:, 1] > 0).long()   # binary label based on first two features

train_size = 1600
X_train, X_test = X[:train_size], X[train_size:]
y_train, y_test = y[:train_size], y[train_size:]

train_ds = TensorDataset(X_train, y_train)
test_ds  = TensorDataset(X_test,  y_test)
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
test_loader  = DataLoader(test_ds,  batch_size=64, shuffle=False)

# Model
model = nn.Sequential(
    nn.Linear(20, 64), nn.ReLU(), nn.Dropout(0.2),
    nn.Linear(64, 32), nn.ReLU(),
    nn.Linear(32, 2),  # 2 output logits for binary classification
).to(device)

optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
criterion = nn.CrossEntropyLoss()
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=30)

def train_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss = 0.0
    correct = 0

    for X_batch, y_batch in loader:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)

        optimizer.zero_grad()          # clear gradients from previous step
        logits = model(X_batch)        # forward pass
        loss   = criterion(logits, y_batch)  # compute loss
        loss.backward()               # compute gradients (backward pass)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # gradient clipping
        optimizer.step()              # update weights

        total_loss += loss.item() * len(y_batch)
        correct    += (logits.argmax(dim=1) == y_batch).sum().item()

    return total_loss / len(loader.dataset), correct / len(loader.dataset)


def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss = 0.0
    correct = 0

    with torch.no_grad():   # no gradient tracking during evaluation
        for X_batch, y_batch in loader:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            logits = model(X_batch)
            loss   = criterion(logits, y_batch)
            total_loss += loss.item() * len(y_batch)
            correct    += (logits.argmax(dim=1) == y_batch).sum().item()

    return total_loss / len(loader.dataset), correct / len(loader.dataset)


# Training loop
for epoch in range(30):
    train_loss, train_acc = train_epoch(model, train_loader, optimizer, criterion, device)
    val_loss,   val_acc   = evaluate(model, test_loader, criterion, device)
    scheduler.step()

    if (epoch + 1) % 5 == 0:
        print(f"Epoch {epoch+1:2d} | "
              f"Train Loss: {train_loss:.4f}  Acc: {train_acc:.3f} | "
              f"Val Loss: {val_loss:.4f}  Acc: {val_acc:.3f} | "
              f"LR: {scheduler.get_last_lr()[0]:.6f}")

Saving and Loading Models

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
optimizer = torch.optim.Adam(model.parameters())

# Save model weights only — preferred for inference
torch.save(model.state_dict(), "model_weights.pt")

# Load weights back
model_loaded = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
model_loaded.load_state_dict(torch.load("model_weights.pt", map_location="cpu"))
model_loaded.eval()

# Save full training checkpoint — for resuming training
torch.save({
    "epoch":                30,
    "model_state_dict":     model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "loss":                 0.0431,
}, "checkpoint.pt")

# Resume from checkpoint
checkpoint = torch.load("checkpoint.pt", map_location="cpu")
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
start_epoch = checkpoint["epoch"] + 1