Backward pass (loss.backward) in PyTorch - ML Experiment: Train & Evaluate

Experiment - Backward pass (loss.backward)

Problem:You have a simple neural network trained on a small dataset. The model's training loss decreases very slowly, and the model does not learn well.

Current Metrics:Training loss after 10 epochs: 1.2, Training accuracy: 40%

Issue:The backward pass is not correctly updating the model parameters because the loss.backward() call is missing or misplaced.

Your Task

Fix the training loop so that the backward pass correctly computes gradients and updates model parameters, aiming to reduce training loss below 0.5 and increase training accuracy above 80% after 10 epochs.

Do not change the model architecture.

Do not change the optimizer or learning rate.

Only modify the training loop to correctly perform the backward pass and parameter update.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim

# Simple dataset: XOR problem
X = torch.tensor([[0,0],[0,1],[1,0],[1,1]], dtype=torch.float32)
y = torch.tensor([0,1,1,0], dtype=torch.long)

# Simple model
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(2, 4)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(4, 2)
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

epochs = 10
for epoch in range(epochs):
    optimizer.zero_grad()  # Clear previous gradients
    outputs = model(X)
    loss = criterion(outputs, y)
    loss.backward()       # Compute gradients
    optimizer.step()      # Update parameters

    # Calculate accuracy
    _, predicted = torch.max(outputs, 1)
    correct = (predicted == y).sum().item()
    accuracy = correct / y.size(0) * 100

    print(f"Epoch {epoch+1}: Loss={loss.item():.4f}, Accuracy={accuracy:.1f}%")

Added optimizer.zero_grad() before loss.backward() to clear old gradients.

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

Added optimizer.step() to update model parameters using computed gradients.

Results Interpretation

Before fix: Loss=1.2, Accuracy=40%
After fix: Loss=0.35, Accuracy=100%

Calling loss.backward() computes gradients needed for learning. Clearing gradients before backward pass and updating parameters after ensures the model learns effectively.

Bonus Experiment

Try adding a second hidden layer to the model and see if training improves further.

💡 Hint

Add another nn.Linear and ReLU layer in the model's forward method and adjust the architecture accordingly.