Pytorch Patterns

PyTorch Patterns

Model Definition

CNN with Residual Blocks

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

class ResidualBlock(nn.Module):
    def __init__(self, channels: int):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(channels)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.bn2(self.conv2(x))
        return F.relu(x + residual)

Transformer Classifier (NLP)

class TransformerClassifier(nn.Module):
    def __init__(self, vocab_size: int, embed_dim: int, num_heads: int, num_classes: int) -> None:
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.transformer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=num_heads,
            dim_feedforward=embed_dim * 4,
            batch_first=True,  # (batch, seq, feature)
        )
        self.classifier = nn.Sequential(
            nn.Linear(embed_dim, embed_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(embed_dim // 2, num_classes),
        )

    def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
        emb = self.embedding(x)
        enc = self.transformer(emb, src_key_padding_mask=mask)
        pooled = enc.mean(dim=1)  # mean pooling over sequence
        return self.classifier(pooled)

Custom Dataset & DataLoader

Tabular Data

from torch.utils.data import Dataset, DataLoader
import pandas as pd

class TabularDataset(Dataset):
    def __init__(self, csv_path: str, target_col: str):
        df = pd.read_csv(csv_path)
        self.X = torch.tensor(df.drop(columns=[target_col]).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]

loader = DataLoader(
    TabularDataset("data.csv", target_col="label"),
    batch_size=64,
    shuffle=True,
    num_workers=4,
    pin_memory=True,        # faster GPU transfer
    persistent_workers=True,
)

Text Data with a Tokenizer

class TextDataset(Dataset):
    def __init__(self, texts: list[str], labels: list[int], tokenizer, max_len: int = 128) -> None:
        self.encodings = tokenizer(texts, truncation=True, max_length=max_len, padding="max_length")
        self.labels = labels

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

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        return {
            "input_ids": torch.tensor(self.encodings["input_ids"][idx]),
            "attention_mask": torch.tensor(self.encodings["attention_mask"][idx]),
            "labels": torch.tensor(self.labels[idx]),
        }

Training Loop (Mixed Precision + Scheduler)

from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.cuda.amp import GradScaler, autocast

def train(model, loader, optimizer, criterion, device, epochs=10):
    scaler = GradScaler()                      # mixed precision
    scheduler = CosineAnnealingLR(optimizer, T_max=epochs)
    model.train()

    for epoch in range(epochs):
        total_loss = 0
        for batch_x, batch_y in loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)

            optimizer.zero_grad()

            with autocast():                   # AMP — fp16 forward pass
                output = model(batch_x)
                loss = criterion(output, batch_y)

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()

            total_loss += loss.item()

        scheduler.step()
        print(f"Epoch {epoch+1}: loss={total_loss/len(loader):.4f}")

Distributed Training (DDP)

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def train_ddp(rank, world_size, model_fn, dataset):
    setup(rank, world_size)
    model = model_fn().to(rank)
    model = DDP(model, device_ids=[rank])

    sampler = torch.utils.data.DistributedSampler(
        dataset, num_replicas=world_size, rank=rank
    )
    loader = DataLoader(dataset, batch_size=64, sampler=sampler)

    # Training loop here...
    dist.destroy_process_group()

# Launch: torchrun --nproc-per-node=4 train.py

torch.compile & Model Export

# Compile for speed (PyTorch 2.0+)
model = torch.compile(model, mode="max-autotune")

# Export to ONNX
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model, dummy_input, "model.onnx",
    input_names=["image"], output_names=["logits"],
    dynamic_axes={"image": {0: "batch_size"}},
    opset_version=17,
)

# TorchScript (serializable)
scripted = torch.jit.script(model)
scripted.save("model.pt")
loaded = torch.jit.load("model.pt")

Checkpointing

# Save
torch.save({
    "epoch": epoch,
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
    "loss": loss,
}, "checkpoint.pt")

# Resume
checkpoint = torch.load("checkpoint.pt", map_location=device)
model.load_state_dict(checkpoint["model_state"])
optimizer.load_state_dict(checkpoint["optimizer_state"])
start_epoch = checkpoint["epoch"] + 1

Device Handling

device = torch.device(
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()  # Apple Silicon
    else "cpu"
)
model = model.to(device)
# Always move both model and data to device before the forward pass

Checklist

• pin_memory=True + persistent_workers=True on DataLoader when training on GPU
• Mixed precision (autocast + GradScaler) enabled for GPU training
• Gradient clipping applied before optimizer.step()
• Checkpoints save both model AND optimizer state, plus the epoch
• Device selection falls back cleanly: CUDA → MPS → CPU
• torch.compile benchmarked before adopting — not all models benefit

See Also

• skills/ml-ai/rag-patterns.md
• skills/python-ecosystem/python-ml.md

Merged from the former skills/python-ecosystem/pytorch-patterns.md and skills/ml-ai/pytorch-patterns.md, which had diverged into two different filename-colliding skills. PyTorch is ML/AI-specific, so this canonical version now lives under ml-ai/ only.