pytorch-patterns

name	pytorch-patterns
description	Idiomatic PyTorch for robust, reproducible, memory-conscious training pipelines: device-agnostic placement, full seed control, explicit tensor shape tracking, clean nn.Module construction, weight initialisation, correct train/eval mode discipline, the standard training and validation loops, efficient Dataset/DataLoader configuration, variable-length collation, resumable checkpointing, and the performance levers (mixed precision, gradient checkpointing, torch.compile). Front-loads the autograd- and correctness-breaking anti-patterns: forgetting eval() at validation, in-place ops severing the graph, .item() before backward, moving the model to GPU inside the loop, and the pickle-RCE risk of loading checkpoints without weights_only. Use when writing or reviewing models, training scripts, or data pipelines, or when tuning GPU memory and speed.
version	1.0.0
metadata	{"activation_triggers":["torch|pytorch","nn\\.Module","DataLoader|Dataset","optimizer|zero_grad|backward","autocast|GradScaler|amp","torch\\.compile|checkpoint","cuda|state_dict|weights_only"],"paths":["**/*.py"],"risk_class":"low","domain":"data-ml"}
source	affaan-m/ecc@81af4076 skills/pytorch-patterns/
license	MIT

PyTorch Patterns

Idiomatic PyTorch for training code that is portable across hardware, reproducible run-to-run, and honest about GPU memory. The through-line: keep the model device-agnostic, control every source of randomness, track tensor shapes explicitly, and never let a convenience call quietly detach the autograd graph.

When to Activate

• Writing a new model, training script, or data pipeline.
• Reviewing deep-learning code for correctness or efficiency.
• Debugging a training loop, a data loader, or a loss that will not move.
• Tuning GPU memory footprint or training throughput.
• Setting up an experiment that must be reproducible.

Version Note

The AMP and compilation surfaces have moved. This skill uses the current torch.amp namespace (torch.amp.autocast("cuda", ...), torch.amp.GradScaler("cuda")) rather than the older torch.cuda.amp, and torch.compile (available from 2.0). If you are on an older build, the concepts hold but the import paths differ — check torch.__version__.

Type hints below use typing.Optional[...] / typing.Tuple[...] rather than PEP 604 X | None, so the snippets stay valid on Python 3.9. Add from __future__ import annotations at the top of a module if you prefer the | spelling on 3.9.

Core Principles

1. Device-agnostic placement

Resolve the device once, then move the model and every batch onto it. Never hardcode .cuda() — it turns "no GPU available" into a crash instead of a CPU fallback.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
# per batch:
data = data.to(device)

# Avoid: hard-fails on any CPU-only machine
model = MyModel().cuda()

2. Reproducibility first

Seed every RNG your pipeline touches — PyTorch (CPU and CUDA), NumPy, and Python's random. For bit-reproducible runs also force deterministic cuDNN, accepting the throughput cost.

def set_seed(seed: int = 42) -> None:
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False  # benchmark=True is faster but non-deterministic

cudnn.benchmark = True auto-tunes convolution algorithms and is faster for fixed input sizes — but it makes runs non-deterministic. Pick one; do not expect both.

3. Explicit shape management

Annotate the shape after each transform in forward. A one-line comment per step turns a silent view/reshape mismatch into something you can read.

def forward(self, x: torch.Tensor) -> torch.Tensor:
    # x: (N, C, H, W)
    x = self.conv1(x)            # (N, 32, H, W)
    x = self.pool(x)             # (N, 32, H//2, W//2)
    x = x.view(x.size(0), -1)    # (N, 32 * H//2 * W//2)
    return self.fc(x)            # (N, num_classes)

Prefer reshape over view when a tensor may be non-contiguous (e.g. after a permute or transpose); view raises on non-contiguous memory.

Model Architecture

Build submodules in __init__, not forward

Every learnable layer must be constructed once in __init__ so its parameters register with the module. Constructing a layer (or its weights) inside forward creates fresh, untrained parameters on every call and they never appear in model.parameters().

class ImageClassifier(nn.Module):
    def __init__(self, num_classes: int, dropout: float = 0.5) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(64 * 16 * 16, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

Weight initialisation

Initialise explicitly rather than trusting layer defaults, especially for deep nets. model.apply(fn) walks every submodule.

def init_weights(module: nn.Module) -> None:
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
    elif isinstance(module, nn.BatchNorm2d):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)

model.apply(init_weights)

Training Loops

Train step

Set train mode, clear grads with set_to_none=True (frees the buffers and is marginally faster), run the forward under autocast when using mixed precision, clip gradients, and step. Accumulate the loss as a Python float with .item() for logging only.

from typing import Optional

def train_one_epoch(
    model: nn.Module,
    dataloader: DataLoader,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: torch.device,
    scaler: Optional[torch.amp.GradScaler] = None,
) -> float:
    model.train()  # dropout on, BatchNorm updates running stats
    total_loss = 0.0

    for data, target in dataloader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad(set_to_none=True)

        with torch.amp.autocast("cuda", enabled=scaler is not None):
            output = model(data)
            loss = criterion(output, target)

        if scaler is not None:
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)  # unscale before clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

        total_loss += loss.item()

    return total_loss / len(dataloader)

Note the ordering under AMP: scaler.scale(loss).backward(), then scaler.unscale_(optimizer) before clip_grad_norm_, then scaler.step / scaler.update. Clipping scaled gradients clips the wrong magnitude.

Validation step

Set eval mode and disable gradients. @torch.no_grad() as a decorator is cleaner than wrapping the body and covers early returns.

from typing import Tuple

@torch.no_grad()
def evaluate(
    model: nn.Module,
    dataloader: DataLoader,
    criterion: nn.Module,
    device: torch.device,
) -> Tuple[float, float]:
    model.eval()  # dropout off, BatchNorm uses running stats
    total_loss, correct, total = 0.0, 0, 0

    for data, target in dataloader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        total_loss += criterion(output, target).item()
        correct += (output.argmax(1) == target).sum().item()
        total += target.size(0)

    return total_loss / len(dataloader), correct / total

model.eval() and torch.no_grad() are independent and you need both: eval() switches layer behaviour (dropout, BatchNorm), no_grad() stops building the autograd graph. Forgetting eval() is the most common silent validation bug — see Anti-Patterns.

Data Pipeline

Custom Dataset

A Dataset needs __len__ and __getitem__. Keep per-item work light and push heavy transforms into the loader's workers.

from typing import Optional, Tuple

class ImageDataset(Dataset):
    def __init__(
        self,
        image_dir: str,
        labels: dict,
        transform: Optional[transforms.Compose] = None,
    ) -> None:
        self.image_paths = list(Path(image_dir).glob("*.jpg"))
        self.labels = labels
        self.transform = transform

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        img = Image.open(self.image_paths[idx]).convert("RGB")
        label = self.labels[self.image_paths[idx].stem]
        if self.transform:
            img = self.transform(img)
        return img, label

DataLoader configuration

The defaults (num_workers=0, no pinned memory) leave the GPU starved. Load in parallel workers, pin host memory for a faster host→device copy, and keep workers alive across epochs.

dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,             # training only; never shuffle eval
    num_workers=4,            # parallel loading
    pin_memory=True,          # faster host -> GPU copy
    persistent_workers=True,  # don't respawn workers each epoch
    drop_last=True,           # stable batch size (matters for BatchNorm)
)

Variable-length collation

For sequences, pad to the batch's max length in a collate_fn.

from typing import List, Tuple

def collate_fn(
    batch: List[Tuple[torch.Tensor, int]],
) -> Tuple[torch.Tensor, torch.Tensor]:
    sequences, labels = zip(*batch)
    padded = nn.utils.rnn.pad_sequence(sequences, batch_first=True, padding_value=0)
    return padded, torch.tensor(labels)

dataloader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)

Checkpointing

Save the full training state — model, optimizer, epoch, metrics — so a run can resume, not just weights. On load, pass weights_only=True (see the security note) and map to CPU first so a checkpoint saved on GPU loads on any machine.

def save_checkpoint(model, optimizer, epoch: int, loss: float, path: str) -> None:
    torch.save(
        {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "loss": loss,
        },
        path,
    )

def load_checkpoint(path: str, model, optimizer=None) -> dict:
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    model.load_state_dict(checkpoint["model_state_dict"])
    if optimizer is not None:
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    return checkpoint

Security: torch.load historically unpickled arbitrary objects, so a malicious checkpoint could execute code on load. Always pass weights_only=True for anything you did not produce yourself, and treat third-party checkpoints as untrusted input. Save with state_dict(), not the whole model object — pickling a full nn.Module binds the file to your exact class layout and is far more fragile.

Performance Levers

Mixed precision (AMP)

Run the forward and loss in lower precision under autocast, and use a GradScaler so small gradients do not underflow to zero. Often ~2x throughput and roughly half the activation memory on supported GPUs.

scaler = torch.amp.GradScaler("cuda")
for data, target in dataloader:
    with torch.amp.autocast("cuda"):
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)

Gradient checkpointing (trade compute for memory)

Recompute intermediate activations during the backward pass instead of holding them — lets a larger model fit at the cost of extra forward compute.

from torch.utils.checkpoint import checkpoint

class LargeModel(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = checkpoint(self.block1, x, use_reentrant=False)
        x = checkpoint(self.block2, x, use_reentrant=False)
        return self.head(x)

Pass use_reentrant=False — the reentrant path is legacy and interacts badly with some inputs.

torch.compile (2.0+)

JIT-compile the model into fused kernels. Wrap once after moving to device.

model = torch.compile(MyModel().to(device), mode="reduce-overhead")
# modes: "default" (safe) · "reduce-overhead" · "max-autotune" (slowest to compile)

The first step or two pays a compilation cost; steady-state is faster.

Quick Reference

Idiom	Purpose
model.train() / model.eval()	Switch layer behaviour before the loop
torch.no_grad() / @torch.no_grad()	Stop building the graph for inference
optimizer.zero_grad(set_to_none=True)	Cheaper, more thorough grad clearing
.to(device)	Device-agnostic placement
torch.amp.autocast + GradScaler	Mixed precision without gradient underflow
pin_memory=True	Faster host→GPU transfer
torch.compile	Kernel fusion for speed (2.0+)
weights_only=True	Safe checkpoint loading
torch.manual_seed (+ NumPy, random)	Reproducible runs
torch.utils.checkpoint	Fit larger models by recomputing activations

Anti-Patterns

# 1) Forgetting eval() — dropout still fires, BatchNorm uses batch stats
model.train()
with torch.no_grad():
    output = model(val_data)          # WRONG: metrics are noisy and wrong
# Right:
model.eval()
with torch.no_grad():
    output = model(val_data)

# 2) In-place ops that sever autograd
x = F.relu(x, inplace=True)           # can corrupt the graph
x += residual                          # in-place add on a graph tensor
# Right:
x = F.relu(x)
x = x + residual

# 3) Moving the model to GPU inside the loop
for data, target in dataloader:
    model = model.cuda()               # re-copies the model every iteration
# Right: move once, before the loop
model = model.to(device)
for data, target in dataloader:
    data, target = data.to(device), target.to(device)

# 4) .item() before backward — detaches from the graph
loss = criterion(output, target).item()
loss.backward()                        # error: a float has no backward()
# Right: keep the tensor; call .item() only to log
loss = criterion(output, target)
loss.backward()
print(f"loss={loss.item():.4f}")

# 5) Saving the whole model object
torch.save(model, "model.pt")          # fragile, tied to class layout
# Right: save the state_dict
torch.save(model.state_dict(), "model.pt")

When memory or speed is the question, measure before guessing: profile with torch.profiler and read the allocator with torch.cuda.memory_summary().

Changelog

• 1.0.0 — Initial clean-room authoring. Covers device-agnostic placement, full seed control, explicit shape tracking, module construction and weight init, train/eval discipline, the standard train and validation loops, Dataset/DataLoader configuration and variable-length collation, resumable checkpointing with the weights_only security note, the performance levers (AMP, gradient checkpointing, torch.compile), and the five autograd- and correctness-breaking anti-patterns. Type hints written for Python 3.9 compatibility (typing.Optional/Tuple, current torch.amp namespace). Skill-Import-Attestation: reviewed-by=AE9B236FDAF0462874060C6BCFCFACF00335DC74; sha256=9bd2169e09e8cfc7677ffbbfd5e7186d1782fa8be257d42eb69a17d2d44f4778