Python / PyTorch Fundamentals Interview Questions
What is the standard PyTorch training loop and what does each step do?
The PyTorch training loop follows a fixed five-step pattern repeated for every batch. Understanding exactly what each line does and what happens if you skip or reorder a step is essential for debugging training issues.
import torch import torch.nn as nn import torch.optim as optim model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 2)) optimizer = optim.AdamW(model.parameters(), lr=1e-3) loss_fn = nn.CrossEntropyLoss() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) def train_one_epoch(model, loader, optimizer, loss_fn, device): model.train() # 0. enables Dropout, BatchNorm train mode total_loss = 0.0 for X_batch, y_batch in loader: X_batch = X_batch.to(device) y_batch = y_batch.to(device) optimizer.zero_grad() # 1. clear gradients from previous step logits = model(X_batch) # 2. forward pass loss = loss_fn(logits, y_batch) # 3. compute loss loss.backward() # 4. backpropagate â fills .grad optimizer.step() # 5. update weights using gradients total_loss += loss.item() * X_batch.size(0) return total_loss / len(loader.dataset) @torch.no_grad() # disable gradient tracking for eval def validate(model, loader, loss_fn, device): model.eval() # disables Dropout, BatchNorm uses running stats total_loss, correct = 0.0, 0 for X_batch, y_batch in loader: X_batch, y_batch = X_batch.to(device), y_batch.to(device) logits = model(X_batch) loss = loss_fn(logits, y_batch) total_loss += loss.item() * X_batch.size(0) correct += (logits.argmax(1) == y_batch).sum().item() return total_loss / len(loader.dataset), correct / len(loader.dataset) # Full training loop for epoch in range(10): train_loss = train_one_epoch(model, train_loader, optimizer, loss_fn, device) val_loss, val_acc = validate(model, val_loader, loss_fn, device) print(f"Epoch {epoch}: train_loss={train_loss:.4f} val_loss={val_loss:.4f} val_acc={val_acc:.4f}")
| Step | Call | Purpose |
|---|---|---|
| 0 | model.train() | Enable Dropout, set BatchNorm to use batch statistics |
| 1 | optimizer.zero_grad() | Clear accumulated gradients from the previous step |
| 2 | model(x) | Forward pass compute predictions |
| 3 | loss_fn(pred, target) | Compute scalar loss |
| 4 | loss.backward() | Backpropagate populate .grad on each parameter |
| 5 | optimizer.step() | Update parameters using gradients and the optimizer's rule |
More Related questions...