Implementation exercises for PyTorch mechanics — building custom Datasets, complete training loops, and debugging common implementation bugs.
🟢 Problem 1: Build a custom Dataset and DataLoader
Task: Implement a custom Dataset for a simple in-memory list of (feature, label) pairs, then wrap it in a DataLoader and iterate over a batch.
from torch.utils.data import Dataset, DataLoader
import torch
class SimpleDataset(Dataset):
def __init__(self, features, labels):
self.features = features
self.labels = labels
def __len__(self):
return len(self.features)
def __getitem__(self, idx):
return self.features[idx], self.labels[idx]
features = torch.randn(100, 10)
labels = torch.randint(0, 2, (100,))
dataset = SimpleDataset(features, labels)
loader = DataLoader(dataset, batch_size=16, shuffle=True)
x_batch, y_batch = next(iter(loader))
print(x_batch.shape, y_batch.shape) # torch.Size([16, 10]) torch.Size([16])🟡 Problem 2: Write a complete training loop with validation, from scratch
Task: Write a full training loop that trains for multiple epochs, tracks training and validation loss, and prints progress each epoch.
def train_model(model, train_loader, val_loader, epochs=5, lr=0.001):
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
loss_fn = torch.nn.CrossEntropyLoss()
for epoch in range(epochs):
model.train()
train_loss = 0.0
for x_batch, y_batch in train_loader:
optimizer.zero_grad()
loss = loss_fn(model(x_batch), y_batch)
loss.backward()
optimizer.step()
train_loss += loss.item()
model.eval()
val_loss = 0.0
with torch.no_grad():
for x_batch, y_batch in val_loader:
val_loss += loss_fn(model(x_batch), y_batch).item()
print(f"Epoch {epoch+1}: train_loss={train_loss/len(train_loader):.4f}, "
f"val_loss={val_loss/len(val_loader):.4f}")Hint if stuck: Notice model.train() and model.eval() are called at the start of each phase — forgetting this is a very common bug that silently produces incorrect results if the model uses Dropout or BatchNorm.
🔴 Problem 3: Find and fix the bugs in this broken training loop
Task: The following training loop has three separate bugs. Find and fix all three before reading the solution below it.
# BUGGY CODE -- find the 3 bugs before scrolling to the fixed version
def broken_train(model, loader, epochs=5):
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = torch.nn.CrossEntropyLoss()
for epoch in range(epochs):
for x_batch, y_batch in loader:
loss = loss_fn(model(x_batch), y_batch)
loss.backward()
optimizer.step()# FIXED VERSION
def fixed_train(model, loader, epochs=5):
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = torch.nn.CrossEntropyLoss()
model.train() # BUG 1 FIX: ensure train mode (matters if using Dropout/BatchNorm)
for epoch in range(epochs):
for x_batch, y_batch in loader:
optimizer.zero_grad() # BUG 2 FIX: missing zero_grad() -- gradients would
# accumulate across batches, corrupting every update
loss = loss_fn(model(x_batch), y_batch)
loss.backward()
optimizer.step()
print(f"Epoch {epoch}: loss={loss.item():.4f}") # BUG 3 FIX: original had no
# progress visibility at allHint if stuck: The most impactful bug is the missing zero_grad() — without it, gradients accumulate across every batch of every epoch instead of being freshly computed each step, causing the effective gradient magnitude to grow uncontrollably as training proceeds.
🟡 Problem 4: Implement a custom loss function
Task: Implement a custom weighted loss function that penalizes false negatives more heavily than false positives (useful for imbalanced classification), as a subclass of nn.Module.
class WeightedBCELoss(torch.nn.Module):
def __init__(self, false_negative_weight=3.0):
super().__init__()
self.fn_weight = false_negative_weight
def forward(self, predictions, targets):
predictions = torch.clamp(predictions, min=1e-7, max=1-1e-7) # avoid log(0)
loss = -(self.fn_weight * targets * torch.log(predictions) +
(1 - targets) * torch.log(1 - predictions))
return loss.mean()
criterion = WeightedBCELoss(false_negative_weight=3.0)
predictions = torch.sigmoid(torch.randn(8))
targets = torch.randint(0, 2, (8,)).float()
loss = criterion(predictions, targets)
print(f"Weighted loss: {loss.item():.4f}")