Pytorch Patterns

1---2name: pytorch-patterns3description: PyTorch deep learning patterns and best practices for building robust, efficient, and reproducible training pipelines, model architectures, and data loading.4origin: ECC5---6 7# PyTorch Development Patterns8 9Idiomatic PyTorch patterns and best practices for building robust, efficient, and reproducible deep learning applications.10 11## When to Activate12 13- Writing new PyTorch models or training scripts14- Reviewing deep learning code15- Debugging training loops or data pipelines16- Optimizing GPU memory usage or training speed17- Setting up reproducible experiments18 19## Core Principles20 21### 1. Device-Agnostic Code22 23Always write code that works on both CPU and GPU without hardcoding devices.24 25```python26# Good: Device-agnostic27device = torch.device("cuda" if torch.cuda.is_available() else "cpu")28model = MyModel().to(device)29data = data.to(device)30 31# Bad: Hardcoded device32model = MyModel().cuda()  # Crashes if no GPU33data = data.cuda()34```35 36### 2. Reproducibility First37 38Set all random seeds for reproducible results.39 40```python41# Good: Full reproducibility setup42def set_seed(seed: int = 42) -> None:43    torch.manual_seed(seed)44    torch.cuda.manual_seed_all(seed)45    np.random.seed(seed)46    random.seed(seed)47    torch.backends.cudnn.deterministic = True48    torch.backends.cudnn.benchmark = False49 50# Bad: No seed control51model = MyModel()  # Different weights every run52```53 54### 3. Explicit Shape Management55 56Always document and verify tensor shapes.57 58```python59# Good: Shape-annotated forward pass60def forward(self, x: torch.Tensor) -> torch.Tensor:61    # x: (batch_size, channels, height, width)62    x = self.conv1(x)    # -> (batch_size, 32, H, W)63    x = self.pool(x)     # -> (batch_size, 32, H//2, W//2)64    x = x.view(x.size(0), -1)  # -> (batch_size, 32*H//2*W//2)65    return self.fc(x)    # -> (batch_size, num_classes)66 67# Bad: No shape tracking68def forward(self, x):69    x = self.conv1(x)70    x = self.pool(x)71    x = x.view(x.size(0), -1)  # What size is this?72    return self.fc(x)           # Will this even work?73```74 75## Model Architecture Patterns76 77### Clean nn.Module Structure78 79```python80# Good: Well-organized module81class ImageClassifier(nn.Module):82    def __init__(self, num_classes: int, dropout: float = 0.5) -> None:83        super().__init__()84        self.features = nn.Sequential(85            nn.Conv2d(3, 64, kernel_size=3, padding=1),86            nn.BatchNorm2d(64),87            nn.ReLU(inplace=True),88            nn.MaxPool2d(2),89        )90        self.classifier = nn.Sequential(91            nn.Dropout(dropout),92            nn.Linear(64 * 16 * 16, num_classes),93        )94 95    def forward(self, x: torch.Tensor) -> torch.Tensor:96        x = self.features(x)97        x = x.view(x.size(0), -1)98        return self.classifier(x)99 100# Bad: Everything in forward101class ImageClassifier(nn.Module):102    def __init__(self):103        super().__init__()104 105    def forward(self, x):106        x = F.conv2d(x, weight=self.make_weight())  # Creates weight each call!107        return x108```109 110### Proper Weight Initialization111 112```python113# Good: Explicit initialization114def _init_weights(self, module: nn.Module) -> None:115    if isinstance(module, nn.Linear):116        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")117        if module.bias is not None:118            nn.init.zeros_(module.bias)119    elif isinstance(module, nn.Conv2d):120        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")121    elif isinstance(module, nn.BatchNorm2d):122        nn.init.ones_(module.weight)123        nn.init.zeros_(module.bias)124 125model = MyModel()126model.apply(model._init_weights)127```128 129## Training Loop Patterns130 131### Standard Training Loop132 133```python134# Good: Complete training loop with best practices135def train_one_epoch(136    model: nn.Module,137    dataloader: DataLoader,138    optimizer: torch.optim.Optimizer,139    criterion: nn.Module,140    device: torch.device,141    scaler: torch.amp.GradScaler | None = None,142) -> float:143    model.train()  # Always set train mode144    total_loss = 0.0145 146    for batch_idx, (data, target) in enumerate(dataloader):147        data, target = data.to(device), target.to(device)148 149        optimizer.zero_grad(set_to_none=True)  # More efficient than zero_grad()150 151        # Mixed precision training152        with torch.amp.autocast("cuda", enabled=scaler is not None):153            output = model(data)154            loss = criterion(output, target)155 156        if scaler is not None:157            scaler.scale(loss).backward()158            scaler.unscale_(optimizer)159            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)160            scaler.step(optimizer)161            scaler.update()162        else:163            loss.backward()164            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)165            optimizer.step()166 167        total_loss += loss.item()168 169    return total_loss / len(dataloader)170```171 172### Validation Loop173 174```python175# Good: Proper evaluation176@torch.no_grad()  # More efficient than wrapping in torch.no_grad() block177def evaluate(178    model: nn.Module,179    dataloader: DataLoader,180    criterion: nn.Module,181    device: torch.device,182) -> tuple[float, float]:183    model.eval()  # Always set eval mode — disables dropout, uses running BN stats184    total_loss = 0.0185    correct = 0186    total = 0187 188    for data, target in dataloader:189        data, target = data.to(device), target.to(device)190        output = model(data)191        total_loss += criterion(output, target).item()192        correct += (output.argmax(1) == target).sum().item()193        total += target.size(0)194 195    return total_loss / len(dataloader), correct / total196```197 198## Data Pipeline Patterns199 200### Custom Dataset201 202```python203# Good: Clean Dataset with type hints204class ImageDataset(Dataset):205    def __init__(206        self,207        image_dir: str,208        labels: dict[str, int],209        transform: transforms.Compose | None = None,210    ) -> None:211        self.image_paths = list(Path(image_dir).glob("*.jpg"))212        self.labels = labels213        self.transform = transform214 215    def __len__(self) -> int:216        return len(self.image_paths)217 218    def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]:219        img = Image.open(self.image_paths[idx]).convert("RGB")220        label = self.labels[self.image_paths[idx].stem]221 222        if self.transform:223            img = self.transform(img)224 225        return img, label226```227 228### Efficient DataLoader Configuration229 230```python231# Good: Optimized DataLoader232dataloader = DataLoader(233    dataset,234    batch_size=32,235    shuffle=True,            # Shuffle for training236    num_workers=4,           # Parallel data loading237    pin_memory=True,         # Faster CPU->GPU transfer238    persistent_workers=True, # Keep workers alive between epochs239    drop_last=True,          # Consistent batch sizes for BatchNorm240)241 242# Bad: Slow defaults243dataloader = DataLoader(dataset, batch_size=32)  # num_workers=0, no pin_memory244```245 246### Custom Collate for Variable-Length Data247 248```python249# Good: Pad sequences in collate_fn250def collate_fn(batch: list[tuple[torch.Tensor, int]]) -> tuple[torch.Tensor, torch.Tensor]:251    sequences, labels = zip(*batch)252    # Pad to max length in batch253    padded = nn.utils.rnn.pad_sequence(sequences, batch_first=True, padding_value=0)254    return padded, torch.tensor(labels)255 256dataloader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)257```258 259## Checkpointing Patterns260 261### Save and Load Checkpoints262 263```python264# Good: Complete checkpoint with all training state265def save_checkpoint(266    model: nn.Module,267    optimizer: torch.optim.Optimizer,268    epoch: int,269    loss: float,270    path: str,271) -> None:272    torch.save({273        "epoch": epoch,274        "model_state_dict": model.state_dict(),275        "optimizer_state_dict": optimizer.state_dict(),276        "loss": loss,277    }, path)278 279def load_checkpoint(280    path: str,281    model: nn.Module,282    optimizer: torch.optim.Optimizer | None = None,283) -> dict:284    checkpoint = torch.load(path, map_location="cpu", weights_only=True)285    model.load_state_dict(checkpoint["model_state_dict"])286    if optimizer:287        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])288    return checkpoint289 290# Bad: Only saving model weights (can't resume training)291torch.save(model.state_dict(), "model.pt")292```293 294## Performance Optimization295 296### Mixed Precision Training297 298```python299# Good: AMP with GradScaler300scaler = torch.amp.GradScaler("cuda")301for data, target in dataloader:302    with torch.amp.autocast("cuda"):303        output = model(data)304        loss = criterion(output, target)305    scaler.scale(loss).backward()306    scaler.step(optimizer)307    scaler.update()308    optimizer.zero_grad(set_to_none=True)309```310 311### Gradient Checkpointing for Large Models312 313```python314# Good: Trade compute for memory315from torch.utils.checkpoint import checkpoint316 317class LargeModel(nn.Module):318    def forward(self, x: torch.Tensor) -> torch.Tensor:319        # Recompute activations during backward to save memory320        x = checkpoint(self.block1, x, use_reentrant=False)321        x = checkpoint(self.block2, x, use_reentrant=False)322        return self.head(x)323```324 325### torch.compile for Speed326 327```python328# Good: Compile the model for faster execution (PyTorch 2.0+)329model = MyModel().to(device)330model = torch.compile(model, mode="reduce-overhead")331 332# Modes: "default" (safe), "reduce-overhead" (faster), "max-autotune" (fastest)333```334 335## Quick Reference: PyTorch Idioms336 337| Idiom | Description |338|-------|-------------|339| `model.train()` / `model.eval()` | Always set mode before train/eval |340| `torch.no_grad()` | Disable gradients for inference |341| `optimizer.zero_grad(set_to_none=True)` | More efficient gradient clearing |342| `.to(device)` | Device-agnostic tensor/model placement |343| `torch.amp.autocast` | Mixed precision for 2x speed |344| `pin_memory=True` | Faster CPU→GPU data transfer |345| `torch.compile` | JIT compilation for speed (2.0+) |346| `weights_only=True` | Secure model loading |347| `torch.manual_seed` | Reproducible experiments |348| `gradient_checkpointing` | Trade compute for memory |349 350## Anti-Patterns to Avoid351 352```python353# Bad: Forgetting model.eval() during validation354model.train()355with torch.no_grad():356    output = model(val_data)  # Dropout still active! BatchNorm uses batch stats!357 358# Good: Always set eval mode359model.eval()360with torch.no_grad():361    output = model(val_data)362 363# Bad: In-place operations breaking autograd364x = F.relu(x, inplace=True)  # Can break gradient computation365x += residual                  # In-place add breaks autograd graph366 367# Good: Out-of-place operations368x = F.relu(x)369x = x + residual370 371# Bad: Moving data to GPU inside the training loop repeatedly372for data, target in dataloader:373    model = model.cuda()  # Moves model EVERY iteration!374 375# Good: Move model once before the loop376model = model.to(device)377for data, target in dataloader:378    data, target = data.to(device), target.to(device)379 380# Bad: Using .item() before backward381loss = criterion(output, target).item()  # Detaches from graph!382loss.backward()  # Error: can't backprop through .item()383 384# Good: Call .item() only for logging385loss = criterion(output, target)386loss.backward()387print(f"Loss: {loss.item():.4f}")  # .item() after backward is fine388 389# Bad: Not using torch.save properly390torch.save(model, "model.pt")  # Saves entire model (fragile, not portable)391 392# Good: Save state_dict393torch.save(model.state_dict(), "model.pt")394```395 396__Remember__: PyTorch code should be device-agnostic, reproducible, and memory-conscious. When in doubt, profile with `torch.profiler` and check GPU memory with `torch.cuda.memory_summary()`.