iwantcoding.com
🔥 Daily 👥 Rooms 🏆 Top Log in Sign up

Examples

By · Reviewed & last updated · Part of PyTorch Tutorial

A small gallery of working PyTorch patterns — a typed training loop, an inference pipeline, a tiny custom dataset, and a Lightning-free DDP launcher. Each is short enough to lift into your own project without a framework.

Four idiomatic PyTorch examples

EXAMPLE
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from pathlib import Path

# ===== 1) A small typed training loop =====
device = 'cuda' if torch.cuda.is_available() else 'cpu'

class MLP(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, hidden: int = 128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Linear(hidden, out_dim),
        )
    def forward(self, x): return self.net(x)

def train(model, loader, opt, scheduler, epochs):
    model.train()
    scaler = torch.amp.GradScaler('cuda', enabled=device == 'cuda')
    for ep in range(epochs):
        for x, y in loader:
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            opt.zero_grad(set_to_none=True)
            with torch.amp.autocast(device_type=device, dtype=torch.float16, enabled=device == 'cuda'):
                logits = model(x)
                loss = F.cross_entropy(logits, y)
            scaler.scale(loss).backward()
            scaler.step(opt)
            scaler.update()
        scheduler.step()
        print(f'epoch {ep}: loss={loss.item():.4f}')

# ===== 2) Inference pipeline — clean separation from training =====
@torch.inference_mode()
def predict(model, x):
    model.eval()
    return model(x.to(device)).argmax(dim=-1).cpu()

# ===== 3) A custom Dataset for image-folder data =====
class TinyImageDataset(Dataset):
    def __init__(self, root: str, transform=None):
        self.paths = list(Path(root).rglob('*.jpg'))
        self.labels = sorted({p.parent.name for p in self.paths})
        self.label2idx = {n: i for i, n in enumerate(self.labels)}
        self.transform = transform

    def __len__(self): return len(self.paths)

    def __getitem__(self, i):
        from PIL import Image
        path = self.paths[i]
        img = Image.open(path).convert('RGB')
        if self.transform: img = self.transform(img)
        return img, self.label2idx[path.parent.name]

# ===== 4) Distributed training (DDP) — torchrun launcher =====
# train_ddp.py
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def run_ddp():
    dist.init_process_group(backend='nccl')
    rank = dist.get_rank()
    world = dist.get_world_size()
    local_rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(local_rank)

    model = MLP(32, 10).to(local_rank)
    model = DDP(model, device_ids=[local_rank])

    # IMPORTANT: per-rank shard of the dataset, not the whole thing
    sampler = torch.utils.data.distributed.DistributedSampler(
        TinyImageDataset('/data/train'), num_replicas=world, rank=rank, shuffle=True)
    loader = DataLoader(TinyImageDataset('/data/train'), batch_size=64, sampler=sampler,
                        num_workers=4, pin_memory=True)

    opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=10)

    for ep in range(10):
        sampler.set_epoch(ep)            # so shuffling differs across epochs
        # ... same training loop body as above
        pass

    if rank == 0:
        torch.save(model.module.state_dict(), '/ckpt/last.pt')

    dist.destroy_process_group()

if __name__ == '__main__':
    run_ddp()

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

# ===== 5) Checkpoints that survive everything =====
def save_checkpoint(path, model, opt, scheduler, scaler, epoch, best_val):
    torch.save({
        'epoch': epoch, 'best_val': best_val,
        'model':    model.state_dict(),
        'opt':      opt.state_dict(),
        'scheduler': scheduler.state_dict(),
        'scaler':   scaler.state_dict(),
    }, path)

def load_checkpoint(path, model, opt=None, scheduler=None, scaler=None):
    ck = torch.load(path, map_location=device)
    model.load_state_dict(ck['model'])
    if opt and 'opt' in ck:          opt.load_state_dict(ck['opt'])
    if scheduler and 'scheduler' in ck: scheduler.load_state_dict(ck['scheduler'])
    if scaler and 'scaler' in ck:    scaler.load_state_dict(ck['scaler'])
    return ck['epoch'], ck['best_val']

Why it matters

Use @torch.inference_mode() instead of @torch.no_grad() for prediction. It is a stricter version that also disables view tracking, gives a small speedup, and prevents accidental gradient ops on the inference path. It is the right default for any production model serving code.

Tip: Tweak the snippet with Try it Yourself », then sit the quiz at the bottom of the page.

Example

Example
# MNIST → CIFAR-10 → fine-tune BERT on a small text dataset.
Try it Yourself »

Discussion

Loading…