# Pytorch Training Pipeline

> Standard PyTorch training patterns with AMP, DDP, checkpointing, and best practices.

- Skill: `aselimc/pytorch-training-pipeline` (Agent Skill)
- Install (CLI): `npx skillmds@latest add aselimc/pytorch-training-pipeline`
- Raw SKILL.md: https://api.skillmd.com/api/skills/aselimc/pytorch-training-pipeline/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: DevOps & Infra
- Author: aselimc (https://skillmd.com/u/aselimc)
- Updated: 2026-09-10
- Page: https://skillmd.com/skills/aselimc/pytorch-training-pipeline

---


# PyTorch Training Pipeline

## Standard Training Loop
```python
scaler = torch.amp.GradScaler()
for epoch in range(epochs):
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        with torch.amp.autocast(device_type="cuda"):
            loss = model(batch)
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()
```

## Checklist
1. **Sanity check**: overfit on 1 batch first
2. **AMP**: `torch.amp.autocast` + `GradScaler`
3. **Gradient clipping**: `clip_grad_norm_` (default 1.0)
4. **LR schedule**: cosine with warmup
5. **Checkpointing**: save model + optimizer + scheduler + epoch + best metric
6. **EMA**: optional but recommended for SSL/generative models
7. **Logging**: train/val loss, LR, grad norm per step
8. **Reproducibility**: `torch.manual_seed()`, `deterministic=True` for debugging

## Multi-GPU (DDP)
```python
model = DDP(model, device_ids=[local_rank])
sampler = DistributedSampler(dataset)
```

## Key Libraries
PyTorch, PyTorch Lightning, torchvision, timm