Training loop structure in PyTorch - ML Experiment: Train & Evaluate

Experiment - Training loop structure

Problem:You have a simple neural network trained on a small dataset. The training loop runs but the model does not improve much after a few epochs.

Current Metrics:Training loss starts at 1.2 and only decreases to 1.0 after 10 epochs. Validation loss stays around 1.1 with accuracy around 50%.

Issue:The training loop is missing key steps like zeroing gradients, proper loss calculation, and optimizer steps, causing poor learning.

Your Task

Fix the training loop so the model learns properly and training loss decreases below 0.5 with validation accuracy above 80% after 20 epochs.

Keep the model architecture and dataset the same.

Only modify the training loop code.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Simple dataset
X = torch.randn(100, 10)
y = (X.sum(dim=1) > 0).long()

# Dataset and loader
train_ds = TensorDataset(X, y)
train_dl = DataLoader(train_ds, batch_size=16, shuffle=True)

# Simple model
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 2)
    def forward(self, x):
        return self.fc(x)

model = SimpleNet()

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# Training loop
for epoch in range(20):
    total_loss = 0
    correct = 0
    total = 0
    for xb, yb in train_dl:
        optimizer.zero_grad()  # Clear gradients
        outputs = model(xb)    # Forward pass
        loss = criterion(outputs, yb)  # Compute loss
        loss.backward()        # Backpropagation
        optimizer.step()       # Update weights

        total_loss += loss.item() * xb.size(0)
        _, predicted = torch.max(outputs, 1)
        correct += (predicted == yb).sum().item()
        total += yb.size(0)

    avg_loss = total_loss / total
    accuracy = correct / total * 100
    print(f"Epoch {epoch+1}: Loss={avg_loss:.4f}, Accuracy={accuracy:.2f}%")

Added optimizer.zero_grad() before backpropagation to clear old gradients.

Added loss.backward() to compute gradients from loss.

Added optimizer.step() to update model weights.

Correctly calculated loss using model outputs and targets.

Tracked total loss and accuracy per epoch for monitoring.

Results Interpretation

Before Fix: Loss ~1.0, Accuracy ~50%
After Fix: Loss ~0.3, Accuracy ~85%

A proper training loop must zero gradients, compute loss, backpropagate, and update weights each batch. Missing any step stops learning.

Bonus Experiment

Add a validation loop after each training epoch to monitor validation loss and accuracy.

💡 Hint

Use torch.no_grad() during validation to avoid computing gradients and set model.eval() mode.

Practice

1. Which of the following best describes the main purpose of a training loop in PyTorch?

easy

A. To save the model to a file

B. To repeatedly update the model weights to improve performance

C. To visualize the model architecture

D. To load data from disk into memory

Solution

1. Step 1: Understand the training loop role

   The training loop runs many times to adjust model weights based on data.

2. Step 2: Identify the main goal

   Its goal is to improve model performance by updating weights repeatedly.

3. Final Answer:

   To repeatedly update the model weights to improve performance -> Option B

4. Quick Check:

   Training loop = update weights repeatedly [OK]

Hint: Training loop = repeat updates to improve model [OK]

Common Mistakes:

• Confusing data loading with training loop
• Thinking training loop saves model
• Assuming training loop only visualizes model

2. Which of the following is the correct way to zero gradients in a PyTorch training loop?

easy

A. optimizer.zero_grad()

B. model.zero_gradients()

C. loss.zero_grad()

D. torch.zero_grad()

Solution

1. Step 1: Recall PyTorch gradient clearing method

   Gradients are cleared by calling zero_grad() on the optimizer.

2. Step 2: Check each option

   Only optimizer.zero_grad() is the correct PyTorch syntax.

3. Final Answer:

   optimizer.zero_grad() -> Option A

4. Quick Check:

   Clear gradients with optimizer.zero_grad() [OK]

Hint: Clear gradients with optimizer.zero_grad() [OK]

Common Mistakes:

• Calling zero_grad on model or loss
• Using non-existent torch.zero_grad()
• Forgetting to zero gradients before backward pass

3. What will be the output of the following PyTorch training loop snippet?

for epoch in range(2):
    for batch in range(3):
        print(f"Epoch {epoch}, Batch {batch}")

medium

A. Epoch 0, Batch 0
Epoch 0, Batch 1
Epoch 1, Batch 0
Epoch 1, Batch 1

B. Epoch 0, Batch 0
Epoch 1, Batch 0
Epoch 2, Batch 0

C. Epoch 0, Batch 0
Epoch 0, Batch 1
Epoch 0, Batch 2

D. Epoch 0, Batch 0
Epoch 0, Batch 1
Epoch 0, Batch 2
Epoch 1, Batch 0
Epoch 1, Batch 1
Epoch 1, Batch 2

Solution

1. Step 1: Understand loop ranges

   Outer loop runs for epochs 0 and 1 (2 times). Inner loop runs for batches 0,1,2 (3 times).

2. Step 2: Trace output lines

   For each epoch, batches 0 to 2 print, so total 6 lines with epoch and batch numbers.

3. Final Answer:

   Epoch 0, Batch 0 Epoch 0, Batch 1 Epoch 0, Batch 2 Epoch 1, Batch 0 Epoch 1, Batch 1 Epoch 1, Batch 2 -> Option D

4. Quick Check:

   2 epochs x 3 batches = 6 lines printed [OK]

Hint: Multiply epochs by batches for total prints [OK]

Common Mistakes:

• Confusing range end values
• Mixing epoch and batch indices
• Counting fewer batches than loop runs

4. Identify the error in this PyTorch training loop snippet:

for data, target in dataloader:
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward
    optimizer.step()

medium

A. Missing parentheses after loss.backward

B. optimizer.zero_grad() should be after optimizer.step()

C. criterion should be called before model(data)

D. optimizer.step() should be called before loss.backward

Solution

1. Step 1: Check method calls

   loss.backward is missing parentheses, so it is not called as a function.

2. Step 2: Verify correct order

   optimizer.zero_grad() is correctly before backward and step; criterion call order is correct; optimizer.step() after backward is correct.

3. Final Answer:

   Missing parentheses after loss.backward -> Option A

4. Quick Check:

   Call loss.backward() with parentheses [OK]

Hint: Always add () to call backward() [OK]

Common Mistakes:

• Forgetting parentheses on backward()
• Changing correct optimizer.zero_grad() position
• Misordering loss and model calls

5. You want to train a model for 5 epochs with batches of size 10. Which of the following training loop structures correctly implements this in PyTorch?

hard

A. for epoch in range(5):
optimizer.zero_grad()
for batch_idx, (data, target) in enumerate(dataloader):
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()

B. for batch_idx, (data, target) in enumerate(dataloader):
for epoch in range(5):
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()

C. for epoch in range(5):
for batch_idx, (data, target) in enumerate(dataloader):
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()

D. for epoch in range(5):
for batch_idx, (data, target) in enumerate(dataloader):
output = model(data)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()

Solution

1. Step 1: Check loop order

   The outer loop should be epochs, inner loop batches. for epoch in range(5): for batch_idx, (data, target) in enumerate(dataloader): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() follows this correctly.

2. Step 2: Verify zero_grad and step calls

   optimizer.zero_grad() must be before backward(), and optimizer.step() after. for epoch in range(5): for batch_idx, (data, target) in enumerate(dataloader): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() has correct order.

3. Step 3: Check other options

   for batch_idx, (data, target) in enumerate(dataloader): for epoch in range(5): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() reverses loops (wrong). for epoch in range(5): optimizer.zero_grad() for batch_idx, (data, target) in enumerate(dataloader): output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() zeros gradients outside batch loop (wrong). for epoch in range(5): for batch_idx, (data, target) in enumerate(dataloader): output = model(data) loss = criterion(output, target) optimizer.zero_grad() loss.backward() optimizer.step() zeros gradients after forward pass (wrong order).

4. Final Answer:

   for epoch in range(5): for batch_idx, (data, target) in enumerate(dataloader): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() -> Option C

5. Quick Check:

   Epochs outside, zero_grad before backward, step after [OK]

Hint: Epoch loop outside, zero_grad before backward [OK]

Common Mistakes:

• Swapping epoch and batch loops
• Calling zero_grad outside batch loop
• Calling zero_grad after forward pass