PyTorch for Real Model Development

PyTorch's design maps cleanly to how a software engineer thinks: eager execution, Python-native control flow, explicit data movement. There are no sessions, no graph compilations, no symbolic graph nodes. You write Python, you get a trained model. This is not an accident - PyTorch was designed for researchers who needed to debug, and that design philosophy transfers directly to production engineering.

This module is not a "hello world" tutorial. It is the training loop patterns, data loading idioms, and operational habits you need to ship ML models that actually work.

The Canonical Training Loop

Every PyTorch training script has the same structure. Internalize this before adding anything else:

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

def train(model, train_loader, val_loader, n_epochs=10, lr=3e-4):
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=lr, steps_per_epoch=len(train_loader), epochs=n_epochs
    )
    criterion = nn.CrossEntropyLoss()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    for epoch in range(n_epochs):
        # --- Train phase ---
        model.train()
        train_loss = 0.0
        for X_batch, y_batch in train_loader:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            optimizer.zero_grad()
            logits = model(X_batch)
            loss = criterion(logits, y_batch)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            scheduler.step()
            train_loss += loss.item()

        # --- Validation phase ---
        model.eval()
        val_loss, correct = 0.0, 0
        with torch.no_grad():
            for X_batch, y_batch in val_loader:
                X_batch, y_batch = X_batch.to(device), y_batch.to(device)
                logits = model(X_batch)
                val_loss += criterion(logits, y_batch).item()
                correct += (logits.argmax(1) == y_batch).sum().item()

        n_val = len(val_loader.dataset)
        print(f"Epoch {epoch+1}: train_loss={train_loss/len(train_loader):.4f} "
              f"val_loss={val_loss/len(val_loader):.4f} "
              f"val_acc={correct/n_val:.4f}")

    return model

Key habits baked into this loop: optimizer.zero_grad() before backward (not after), gradient clipping before step, model.eval() + torch.no_grad() for validation, explicit .to(device) for both model and batches.

The Dataset and DataLoader Pattern

Never load your entire dataset into memory and slice it manually. Use Dataset + DataLoader:

pythonCopy
from torch.utils.data import Dataset
import pandas as pd
import numpy as np

class TabularDataset(Dataset):
    def __init__(self, df: pd.DataFrame, target_col: str):
        feature_cols = [c for c in df.columns if c != target_col]
        self.X = torch.tensor(df[feature_cols].values, dtype=torch.float32)
        self.y = torch.tensor(df[target_col].values, dtype=torch.long)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

# Usage
train_ds = TabularDataset(train_df, target_col="label")
val_ds = TabularDataset(val_df, target_col="label")

train_loader = DataLoader(train_ds, batch_size=256, shuffle=True, num_workers=4, pin_memory=True)
val_loader = DataLoader(val_ds, batch_size=512, shuffle=False, num_workers=4, pin_memory=True)

pin_memory=True speeds up host-to-GPU transfers. num_workers > 0 enables parallel data loading. Both are low-cost and should be default.

Building Models with nn.Module

Subclass nn.Module for any non-trivial architecture:

pythonCopy
class MLP(nn.Module):
    def __init__(self, input_dim: int, hidden_dims: list[int], output_dim: int, dropout: float = 0.1):
        super().__init__()
        dims = [input_dim] + hidden_dims
        layers = []
        for in_d, out_d in zip(dims[:-1], dims[1:]):
            layers.extend([nn.Linear(in_d, out_d), nn.LayerNorm(out_d), nn.GELU(), nn.Dropout(dropout)])
        layers.append(nn.Linear(dims[-1], output_dim))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

model = MLP(input_dim=128, hidden_dims=[512, 256], output_dim=10)
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

Saving and Loading Models

pythonCopy
# Save checkpoint (preferred over state_dict-only for resumable training)
torch.save({
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "val_loss": best_val_loss,
}, "checkpoint.pt")

# Load 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

# For inference only - lighter and sufficient
torch.save(model.state_dict(), "model_weights.pt")
model.load_state_dict(torch.load("model_weights.pt", map_location="cpu"))
model.eval()

Always use map_location="cpu" on load. This allows loading a GPU-trained model on a CPU machine (e.g., at inference time).

Mixed Precision Training

On modern GPUs (Ampere, H100), float16/bfloat16 computation is 2-8x faster than float32. PyTorch's torch.cuda.amp makes this a ~5-line change:

pythonCopy
from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()

for X_batch, y_batch in train_loader:
    X_batch, y_batch = X_batch.to(device), y_batch.to(device)
    optimizer.zero_grad()
    with autocast():                         # Forward in float16
        logits = model(X_batch)
        loss = criterion(logits, y_batch)
    scaler.scale(loss).backward()            # Scaled backward
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()

Use bfloat16 on H100/A100 (more numerically stable than float16). On older GPUs, float16 with GradScaler is the right choice.

Reproducibility

pythonCopy
import random
import numpy as np

def set_seed(seed: int = 42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False  # Slower but deterministic

set_seed(42)

cudnn.benchmark = False reduces throughput ~5-10% but eliminates non-determinism from CUDA algorithm selection. Accept this trade for experiment reproducibility; disable for final production training.

Debugging PyTorch Models

pythonCopy
# 1. Shape-check each layer manually
x = torch.randn(4, 128)   # batch of 4, 128 features
h = model.net[0](x)       # check output after first linear
print(h.shape)

# 2. Gradient flow check: are all layers actually learning?
for name, param in model.named_parameters():
    if param.grad is not None:
        print(f"{name}: grad_norm={param.grad.norm().item():.4f}")

# 3. NaN/Inf detection
torch.autograd.set_detect_anomaly(True)  # Expensive - dev only

# 4. Profile a forward pass
with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as prof:
    model(x)
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

Common Mistakes and Bad Instincts

Forgetting optimizer.zero_grad(). Gradients accumulate across calls to backward() by default. Calling step() without zeroing first adds the previous batch's gradient to the current one, producing incorrect updates that are hard to debug because the loss usually still decreases (just suboptimally).

Using .cuda() hard-coded instead of .to(device). The device abstraction makes your code portable. Hard-coding .cuda() breaks on machines without GPU.

Saving only the model, not the optimizer state. When resuming a training run, the optimizer's moment estimates encode the curvature information accumulated over many steps. Discarding them forces the optimizer to restart from scratch and often requires additional warmup.

Building datasets as lists of tensors in RAM. For large datasets, use torch.utils.data.Dataset with lazy loading from disk. Loading everything into RAM upfront causes OOM errors and does not scale.

Where to Go Next

• representation-learning-embeddings-and-similarity: apply PyTorch to build embedding models, triplet loss training, and approximate nearest-neighbor retrieval
• transformers-and-modern-nlp-for-engineers: use these PyTorch patterns inside the Transformer architecture
• fine-tuning-adaptation-and-when-not-to-fine-tune: load pretrained model weights, freeze layers, and fine-tune with the loop patterns built here