Forward pass, loss, backward, step in PyTorch - ML Experiment: Train & Evaluate

Experiment - Forward pass, loss, backward, step

Problem:Train a simple neural network to classify handwritten digits from the MNIST dataset.

Current Metrics:Training loss: 0.05, Training accuracy: 98%, Validation loss: 0.10, Validation accuracy: 95%

Issue:The model trains but the learner does not understand how the forward pass, loss calculation, backward pass, and optimizer step work together.

Your Task

Explain and implement the forward pass, loss calculation, backward pass, and optimizer step in PyTorch to train the model correctly.

Use PyTorch framework.

Use a simple feedforward neural network.

Use CrossEntropyLoss for classification.

Use SGD optimizer.

Hint 1

Hint 2

Hint 3

Hint 4

Hint 5

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Define simple feedforward network
class SimpleNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear = nn.Linear(28*28, 10)

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear(x)
        return logits

# Prepare data
transform = transforms.ToTensor()
train_dataset = datasets.MNIST(root='.', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# Initialize model, loss, optimizer
model = SimpleNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# Training loop for 1 epoch
model.train()
for images, labels in train_loader:
    # Forward pass: compute predictions
    outputs = model(images)
    # Compute loss
    loss = criterion(outputs, labels)
    # Zero gradients before backward pass
    optimizer.zero_grad()
    # Backward pass: compute gradients
    loss.backward()
    # Update parameters
    optimizer.step()

    # Print loss for first batch only
    print(f'Loss: {loss.item():.4f}')
    break

Implemented forward pass by calling model(images) to get predictions.

Calculated loss using CrossEntropyLoss comparing predictions and true labels.

Called optimizer.zero_grad() to clear old gradients.

Called loss.backward() to compute gradients.

Called optimizer.step() to update model parameters.

Results Interpretation

Before: Learner did not understand the sequence of forward pass, loss, backward, and step.

After: Learner implemented and ran code that performs these steps correctly, seeing loss value printed.

The forward pass computes predictions, loss measures error, backward pass calculates gradients, and optimizer step updates model weights. This cycle trains the model.

Bonus Experiment

Try adding a hidden layer with ReLU activation to the network and observe how loss changes.

💡 Hint

Add nn.Linear and nn.ReLU layers between input and output layers, then rerun training.