8# PyTorch Development Patterns
9
10Idiomatic PyTorch patterns and best practices for building robust, efficient, and reproducible deep learning applications.
11
12## When to Activate
13
14- Writing new PyTorch models or training scripts
15- Reviewing deep learning code
16- Debugging training loops or data pipelines
17- Optimizing GPU memory usage or training speed
18- Setting up reproducible experiments
19
20## Core Principles
21
22### 1. Device-Agnostic Code
23
24Always write code that works on both CPU and GPU without hardcoding devices.
25
26```python
27# Good: Device-agnostic
28device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
29model = MyModel().to(device)
30data = data.to(device)
31
32# Bad: Hardcoded device
33model = MyModel().cuda() # Crashes if no GPU
34data = data.cuda()
35```
36
37### 2. Reproducibility First
38
39Set all random seeds for reproducible results.
40
41```python
42# Good: Full reproducibility setup
43def set_seed(seed: int = 42) -> None:
44 torch.manual_seed(seed)
45 torch.cuda.manual_seed_all(seed)
46 np.random.seed(seed)
47 random.seed(seed)
48 torch.backends.cudnn.deterministic = True
49 torch.backends.cudnn.benchmark = False
50
51# Bad: No seed control
52model = MyModel() # Different weights every run
53```
54
55### 3. Explicit Shape Management
56
57Always document and verify tensor shapes.
58
59```python
60# Good: Shape-annotated forward pass
61def forward(self, x: torch.Tensor) -> torch.Tensor:
62 # x: (batch_size, channels, height, width)
63 x = self.conv1(x) # -> (batch_size, 32, H, W)
64 x = self.pool(x) # -> (batch_size, 32, H//2, W//2)
65 x = x.view(x.size(0), -1) # -> (batch_size, 32*H//2*W//2)
66 return self.fc(x) # -> (batch_size, num_classes)
67
68# Bad: No shape tracking
69def forward(self, x):
70 x = self.conv1(x)
71 x = self.pool(x)
72 x = x.view(x.size(0), -1) # What size is this?
73 return self.fc(x) # Will this even work?
74```
75
76## Model Architecture Patterns
77
78### Clean nn.Module Structure
79
80```python
81# Good: Well-organized module
82class ImageClassifier(nn.Module):
83 def __init__(self, num_classes: int, dropout: float = 0.5) -> None:
84 super().__init__()
85 self.features = nn.Sequential(
86 nn.Conv2d(3, 64, kernel_size=3, padding=1),
87 nn.BatchNorm2d(64),
88 nn.ReLU(inplace=True),
89 nn.MaxPool2d(2),
90 )
91 self.classifier = nn.Sequential(
92 nn.Dropout(dropout),
93 nn.Linear(64 * 16 * 16, num_classes),
94 )
95
96 def forward(self, x: torch.Tensor) -> torch.Tensor:
97 x = self.features(x)
98 x = x.view(x.size(0), -1)
99 return self.classifier(x)
100
101# Bad: Everything in forward
102class ImageClassifier(nn.Module):
103 def __init__(self):
104 super().__init__()
105
106 def forward(self, x):
107 x = F.conv2d(x, weight=self.make_weight()) # Creates weight each call!
108 return x
109```
110
111### Proper Weight Initialization
112
113```python
114# Good: Explicit initialization
115def _init_weights(self, module: nn.Module) -> None:
116 if isinstance(module, nn.Linear):
117 nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
118 if module.bias is not None:
119 nn.init.zeros_(module.bias)
120 elif isinstance(module, nn.Conv2d):
121 nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
122 elif isinstance(module, nn.BatchNorm2d):
123 nn.init.ones_(module.weight)
124 nn.init.zeros_(module.bias)
125
126model = MyModel()
127model.apply(model._init_weights)
128```
129
130## Training Loop Patterns
131
132### Standard Training Loop
133
134```python
135# Good: Complete training loop with best practices
136def train_one_epoch(
137 model: nn.Module,
138 dataloader: DataLoader,
139 optimizer: torch.optim.Optimizer,
140 criterion: nn.Module,
141 device: torch.device,
142 scaler: torch.amp.GradScaler | None = None,
143) -> float:
144 model.train() # Always set train mode
145 total_loss = 0.0
146
147 for batch_idx, (data, target) in enumerate(dataloader):
148 data, target = data.to(device), target.to(device)
149
150 optimizer.zero_grad(set_to_none=True) # More efficient than zero_grad()
151
152 # Mixed precision training
153 with torch.amp.autocast("cuda", enabled=scaler is not None):
154 output = model(data)
155 loss = criterion(output, target)
156
157 if scaler is not None:
158 scaler.scale(loss).backward()
159 scaler.unscale_(optimizer)
160 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
161 scaler.step(optimizer)
162 scaler.update()
163 else:
164 loss.backward()
165 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
166 optimizer.step()
167
168 total_loss += loss.item()
169
170 return total_loss / len(dataloader)
171```
172
173### Validation Loop
174
175```python
176# Good: Proper evaluation
177@torch.no_grad() # More efficient than wrapping in torch.no_grad() block
178def evaluate(
179 model: nn.Module,
180 dataloader: DataLoader,
181 criterion: nn.Module,
182 device: torch.device,
183) -> tuple[float, float]:
184 model.eval() # Always set eval mode — disables dropout, uses running BN stats
185 total_loss = 0.0
186 correct = 0
187 total = 0
188
189 for data, target in dataloader:
190 data, target = data.to(device), target.to(device)
191 output = model(data)
192 total_loss += criterion(output, target).item()
193 correct += (output.argmax(1) == target).sum().item()
194 total += target.size(0)
195
196 return total_loss / len(dataloader), correct / total
197```
198
199## Data Pipeline Patterns
200
201### Custom Dataset
202
203```python
204# Good: Clean Dataset with type hints
205class ImageDataset(Dataset):
206 def __init__(
207 self,
208 image_dir: str,
209 labels: dict[str, int],
210 transform: transforms.Compose | None = None,
211 ) -> None:
212 self.image_paths = list(Path(image_dir).glob("*.jpg"))
213 self.labels = labels
214 self.transform = transform
215
216 def __len__(self) -> int:
217 return len(self.image_paths)
218
219 def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]:
220 img = Image.open(self.image_paths[idx]).convert("RGB")
221 label = self.labels[self.image_paths[idx].stem]
222
223 if self.transform:
224 img = self.transform(img)
225
226 return img, label
227```
228
229### Efficient DataLoader Configuration
230
231```python
232# Good: Optimized DataLoader
233dataloader = DataLoader(
234 dataset,
235 batch_size=32,
236 shuffle=True, # Shuffle for training
237 num_workers=4, # Parallel data loading
238 pin_memory=True, # Faster CPU->GPU transfer
239 persistent_workers=True, # Keep workers alive between epochs
240 drop_last=True, # Consistent batch sizes for BatchNorm
241)
242
243# Bad: Slow defaults
244dataloader = DataLoader(dataset, batch_size=32) # num_workers=0, no pin_memory
245```
246
247### Custom Collate for Variable-Length Data
248
249```python
250# Good: Pad sequences in collate_fn
251def collate_fn(batch: list[tuple[torch.Tensor, int]]) -> tuple[torch.Tensor, torch.Tensor]:
252 sequences, labels = zip(*batch)
253 # Pad to max length in batch
254 padded = nn.utils.rnn.pad_sequence(sequences, batch_first=True, padding_value=0)
255 return padded, torch.tensor(labels)
256
257dataloader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)
258```
259
260## Checkpointing Patterns
261
262### Save and Load Checkpoints
263
264```python
265# Good: Complete checkpoint with all training state
266def save_checkpoint(
267 model: nn.Module,
268 optimizer: torch.optim.Optimizer,
269 epoch: int,
270 loss: float,
271 path: str,
272) -> None:
273 torch.save({
274 "epoch": epoch,
275 "model_state_dict": model.state_dict(),
276 "optimizer_state_dict": optimizer.state_dict(),
277 "loss": loss,
278 }, path)
279
280def load_checkpoint(
281 path: str,
282 model: nn.Module,
283 optimizer: torch.optim.Optimizer | None = None,
284) -> dict:
285 checkpoint = torch.load(path, map_location="cpu", weights_only=True)
286 model.load_state_dict(checkpoint["model_state_dict"])
287 if optimizer:
288 optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
289 return checkpoint
290
291# Bad: Only saving model weights (can't resume training)
292torch.save(model.state_dict(), "model.pt")
293```
294
295## Performance Optimization
296
297### Mixed Precision Training
298
299```python
300# Good: AMP with GradScaler
301scaler = torch.amp.GradScaler("cuda")
302for data, target in dataloader:
303 with torch.amp.autocast("cuda"):
304 output = model(data)
305 loss = criterion(output, target)
306 scaler.scale(loss).backward()
307 scaler.step(optimizer)
308 scaler.update()
309 optimizer.zero_grad(set_to_none=True)
310```
311
312### Gradient Checkpointing for Large Models
313
314```python
315# Good: Trade compute for memory
316from torch.utils.checkpoint import checkpoint
317
318class LargeModel(nn.Module):
319 def forward(self, x: torch.Tensor) -> torch.Tensor:
320 # Recompute activations during backward to save memory
321 x = checkpoint(self.block1, x, use_reentrant=False)
322 x = checkpoint(self.block2, x, use_reentrant=False)
323 return self.head(x)
324```
325
326### torch.compile for Speed
327
328```python
329# Good: Compile the model for faster execution (PyTorch 2.0+)
330model = MyModel().to(device)
331model = torch.compile(model, mode="reduce-overhead")
332
333# Modes: "default" (safe), "reduce-overhead" (faster), "max-autotune" (fastest)
334```
335
336## Quick Reference: PyTorch Idioms
337
338| Idiom | Description |
339|-------|-------------|
340| model.train() / model.eval() | Always set mode before train/eval |
341| torch.no_grad() | Disable gradients for inference |
342| optimizer.zero_grad(set_to_none=True) | More efficient gradient clearing |
343| .to(device) | Device-agnostic tensor/model placement |
344| torch.amp.autocast | Mixed precision for 2x speed |
345| pin_memory=True | Faster CPU→GPU data transfer |
346| torch.compile | JIT compilation for speed (2.0+) |
347| weights_only=True | Secure model loading |
348| torch.manual_seed | Reproducible experiments |
349| gradient_checkpointing | Trade compute for memory |
350
351## Anti-Patterns to Avoid
352
353```python
354# Bad: Forgetting model.eval() during validation
355model.train()
356with torch.no_grad():
357 output = model(val_data) # Dropout still active! BatchNorm uses batch stats!
358
359# Good: Always set eval mode
360model.eval()
361with torch.no_grad():
362 output = model(val_data)
363
364# Bad: In-place operations breaking autograd
365x = F.relu(x, inplace=True) # Can break gradient computation
366x += residual # In-place add breaks autograd graph
367
368# Good: Out-of-place operations
369x = F.relu(x)
370x = x + residual
371
372# Bad: Moving data to GPU inside the training loop repeatedly
373for data, target in dataloader:
374 model = model.cuda() # Moves model EVERY iteration!
375
376# Good: Move model once before the loop
377model = model.to(device)
378for data, target in dataloader:
379 data, target = data.to(device), target.to(device)
380
381# Bad: Using .item() before backward
382loss = criterion(output, target).item() # Detaches from graph!
383loss.backward() # Error: can't backprop through .item()
384
385# Good: Call .item() only for logging
386loss = criterion(output, target)
387loss.backward()
388print(f"Loss: {loss.item():.4f}") # .item() after backward is fine
389
390# Bad: Not using torch.save properly
391torch.save(model, "model.pt") # Saves entire model (fragile, not portable)
392
393# Good: Save state_dict
394torch.save(model.state_dict(), "model.pt")
395```
396
397__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().
398