pytorch-patterns

name pytorch-patterns description PyTorch deep learning patterns and best practices for building robust, efficient, and reproducible training pipelines, model architectures, and data loading. origin ECC PyTorch Development Patterns Idiomatic PyTorch patterns and best practices for building robust, efficient, and reproducible deep learning applications. When to Activate Writing new PyTorch models or training scripts Reviewing deep learning code Debugging training loops or data pipelines Optimizing GPU memory usage or training speed Setting up reproducible experiments Core Principles 1. Device-Agnostic Code Always write code that works on both CPU and GPU without hardcoding devices. # Good: Device-agnostic device = torch.device( "cuda" if torch.cuda.is_available() else "cpu" ) model = MyModel().to(device) data = data.to(device) # Bad: Hardcoded device model = MyModel().cuda() # Crashes if no GPU data = data.cuda() 2. Reproducibility First Set all random seeds for reproducible results. # Good: Full reproducibility setup 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 # Bad: No seed control model = MyModel() # Different weights every run 3. Explicit Shape Management Always document and verify tensor shapes. # Good: Shape-annotated forward pass def forward ( self, x: torch.Tensor ) -> torch.Tensor: # x: (batch_size, channels, height, width) x = self .conv1(x) # -> (batch_size, 32, H, W) x = self .pool(x) # -> (batch_size, 32, H//2, W//2) x = x.view(x.size( 0 ), - 1 ) # -> (batch_size, 32*H//2*W//2) return self .fc(x) # -> (batch_size, num_classes) # Bad: No shape tracking def forward ( self, x ): x = self .conv1(x) x = self .pool(x) x = x.view(x.size( 0 ), - 1 ) # What size is this? return self .fc(x) # Will this even work? Model Architecture Patterns Clean nn.Module Structure # Good: Well-organized module 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) # Bad: Everything in forward class ImageClassifier (nn.Module): def __init__ ( self ): super ().__init__() def forward ( self, x ): x = F.conv2d(x, weight= self .make_weight()) # Creates weight each call! return x Proper Weight Initialization # Good: Explicit initialization def _init_weights ( self, 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 = MyModel() model.apply(model._init_weights) Training Loop Patterns Standard Training Loop # Good: Complete training loop with best practices def train_one_epoch ( model: nn.Module, dataloader: DataLoader, optimizer: torch.optim.Optimizer, criterion: nn.Module, device: torch.device, scaler: torch.amp.GradScaler | None = None , ) -> float : model.train() # Always set train mode total_loss = 0.0 for batch_idx, (data, target) in enumerate (dataloader): data, target = data.to(device), target.to(device) optimizer.zero_grad(set_to_none= True ) # More efficient than zero_grad() # Mixed precision training 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) 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) Validation Loop # Good: Proper evaluation @torch.no_grad() # More efficient than wrapping in torch.no_grad() block def evaluate ( model: nn.Module, dataloader: DataLoader, criterion: nn.Module, device: torch.device, ) -> tuple [ float , float ]: model. eval () # Always set eval mode — disables dropout, uses running BN stats total_loss = 0.0 correct = 0 total = 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 Data Pipeline Patterns Custom Dataset # Good: Clean Dataset with type hints class ImageDataset ( Dataset ): def __init__ ( self, image_dir: str , labels: dict [ str , int ], transform: transforms.Compose | None = 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 Efficient DataLoader Configuration # Good: Optimized DataLoader dataloader = DataLoader( dataset, batch_size= 32 , shuffle= True , # Shuffle for training num_workers= 4 , # Parallel data loading pin_memory= True , # Faster CPU->GPU transfer persistent_workers= True , # Keep workers alive between epochs drop_last= True , # Consistent batch sizes for BatchNorm ) # Bad: Slow defaults dataloader = DataLoader(dataset, batch_size= 32 ) # num_workers=0, no pin_memory Custom Collate for Variable-Length Data # Good: Pad sequences in collate_fn def collate_fn ( batch: list [ tuple [torch.Tensor, int ]] ) -> tuple [torch.Tensor, torch.Tensor]: sequences, labels = zip (*batch) # Pad to max length in 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 Patterns Save and Load Checkpoints # Good: Complete checkpoint with all training state def save_checkpoint ( model: nn.Module, optimizer: torch.optim.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: nn.Module, optimizer: torch.optim.Optimizer | None = None , ) -> dict : checkpoint = torch.load(path, map_location= "cpu" , weights_only= True ) model.load_state_dict(checkpoint[ "model_state_dict" ]) if optimizer: optimizer.load_state_dict(checkpoint[ "optimizer_state_dict" ]) return checkpoint # Bad: Only saving model weights (can't resume training) torch.save(model.state_dict(), "model.pt" ) Performance Optimization Mixed Precision Training # Good: AMP with GradScaler 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 for Large Models # Good: Trade compute for memory from torch.utils.checkpoint import checkpoint class LargeModel (nn.Module): def forward ( self, x: torch.Tensor ) -> torch.Tensor: # Recompute activations during backward to save memory x = checkpoint( self .block1, x, use_reentrant= False ) x = checkpoint( self .block2, x, use_reentrant= False ) return self .head(x) torch.compile for Speed # Good: Compile the model for faster execution (PyTorch 2.0+) model = MyModel().to(device) model = torch. compile (model, mode= "reduce-overhead" ) # Modes: "default" (safe), "reduce-overhead" (faster), "max-autotune" (fastest) Quick Reference: PyTorch Idioms Idiom Description model.train() / model.eval() Always set mode before train/eval torch.no_grad() Disable gradients for inference optimizer.zero_grad(set_to_none=True) More efficient gradient clearing .to(device) Device-agnostic tensor/model placement torch.amp.autocast Mixed precision for 2x speed pin_memory=True Faster CPU→GPU data transfer torch.compile JIT compilation for speed (2.0+) weights_only=True Secure model loading torch.manual_seed Reproducible experiments gradient_checkpointing Trade compute for memory Anti-Patterns to Avoid # Bad: Forgetting model.eval() during validation model.train() with torch.no_grad(): output = model(val_data) # Dropout still active! BatchNorm uses batch stats! # Good: Always set eval mode model. eval () with torch.no_grad(): output = model(val_data) # Bad: In-place operations breaking autograd x = F.relu(x, inplace= True ) # Can break gradient computation x += residual # In-place add breaks autograd graph # Good: Out-of-place operations x = F.relu(x) x = x + residual # Bad: Moving data to GPU inside the training loop repeatedly for data, target in dataloader: model = model.cuda() # Moves model EVERY iteration! # Good: Move model once before the loop model = model.to(device) for data, target in dataloader: data, target = data.to(device), target.to(device) # Bad: Using .item() before backward loss = criterion(output, target).item() # Detaches from graph! loss.backward() # Error: can't backprop through .item() # Good: Call .item() only for logging loss = criterion(output, target) loss.backward() print ( f"Loss: {loss.item(): .4 f} " ) # .item() after backward is fine # Bad: Not using torch.save properly torch.save(model, "model.pt" ) # Saves entire model (fragile, not portable) # Good: Save state_dict torch.save(model.state_dict(), "model.pt" ) 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() .