pytorch

PyTorch Best Practices

This guide outlines the definitive best practices for developing with PyTorch, ensuring your code is readable, performant, and production-ready. We prioritize usability, explicit control, and modern tooling.

1. Code Organization and Structure

Structure your PyTorch projects for clarity, testability, and scalability. Encapsulate logical blocks into distinct functions or classes.

1.1. Modularize Your Codebase

Separate data loading, model definition, training, and evaluation into dedicated modules or functions. This makes components reusable and testable.

❌ BAD: Monolithic script

# train.py
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# ... data loading, model definition, training loop all in one file ...

class MyModel(nn.Module):
    # ...
    pass

def main():
    # Data loading
    train_data = TensorDataset(...)
    train_loader = DataLoader(train_data, batch_size=32)

    # Model, optimizer, loss
    model = MyModel()
    optimizer = torch.optim.Adam(model.parameters())
    criterion = nn.CrossEntropyLoss()

    # Training loop
    for epoch in range(10):
        for batch_idx, (data, target) in enumerate(train_loader):
            # ... training logic ...
            pass

if __name__ == "__main__":
    main()

✅ GOOD: Modularized structure

# src/data.py
import torch
from torch.utils.data import DataLoader, TensorDataset

def get_dataloaders(batch_size: int) -> tuple[DataLoader, DataLoader]:
    # Example: Create synthetic data
    X = torch.randn(1000, 784)
    y = torch.randint(0, 10, (1000,))
    train_dataset = TensorDataset(X, y)
    val_dataset = TensorDataset(X[:100], y[:100]) # Smaller val set
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
    return train_loader, val_loader

# src/model.py
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self, num_classes: int = 10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128), # Assuming 28x28 input, adjust for other sizes
            nn.ReLU(),
            nn.Linear(128, num_classes)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = self.classifier(x)
        return x

# src/train.py
import torch
import torch.nn as nn
from torch.optim import Adam
from src.model import SimpleCNN
from src.data import get_dataloaders

def train_epoch(model: nn.Module, loader: DataLoader, optimizer: Adam, criterion: nn.Module, device: torch.device) -> float:
    model.train()
    total_loss = 0.0
    for data, target in loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)

def evaluate_model(model: nn.Module, loader: DataLoader, device: torch.device) -> float:
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    return 100 * correct / total

# main.py
import torch
from src.model import SimpleCNN
from src.data import get_dataloaders
from src.train import train_epoch, evaluate_model

def run_experiment(epochs: int = 10, batch_size: int = 64, lr: float = 1e-3):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    train_loader, val_loader = get_dataloaders(batch_size)
    model = SimpleCNN().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = torch.nn.CrossEntropyLoss()

    for epoch in range(epochs):
        train_loss = train_epoch(model, train_loader, optimizer, criterion, device)
        val_accuracy = evaluate_model(model, val_loader, device)
        print(f"Epoch {epoch+1}: Train Loss = {train_loss:.4f}, Val Acc = {val_accuracy:.2f}%")

    torch.save(model.state_dict(), "final_model.pth")

if __name__ == "__main__":
    run_experiment()