PyTorch Training Loop: From Data to Weights

1. The Anatomy of a Training Loop

Training a neural network is an iterative process. We repeatedly show the model our data and adjust its parameters slightly each time to improve its predictions. This process is structured into epochs and batches.

So, a training loop typically looks like this:

# Loop over the entire dataset multiple times (epochs)
for epoch in range(num_epochs):
    # Loop over the dataset in smaller chunks (batches)
    for batch in data_loader:
        # ... perform training steps on this batch ...

2. Step 1: Loading Data in Batches with DataLoader

Manually splitting data into batches can be tedious. PyTorch provides a powerful utility called torch.utils.data.DataLoader to handle this for us. It wraps a Dataset (like the MNIST dataset we loaded in the previous lesson) and provides an iterator to loop over the data in mini-batches automatically. It can also shuffle the data at every epoch to prevent the model from learning the order of the data.

Building on our last lesson, here’s how you would create DataLoader instances for the MNIST training and test sets:

import torch
import torchvision
import torchvision.transforms as transforms

# Assuming you have the MNIST dataset downloaded
# (This is from the previous lesson's setup)
train_dataset = torchvision.datasets.MNIST(root='./data', 
                                           train=True, 
                                           transform=transforms.ToTensor(),  
                                           download=True)

test_dataset = torchvision.datasets.MNIST(root='./data', 
                                          train=False, 
                                          transform=transforms.ToTensor())

# Create the DataLoaders
batch_size = 100
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, 
                                           batch_size=batch_size, 
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset, 
                                          batch_size=batch_size, 
                                          shuffle=False) # No need to shuffle the test set

Now, we can simply loop over train_loader to get our batches of images and labels.

3. Step 2: The Five Core Training Steps

For each batch of data, we perform five essential steps. These five steps are the heart of the training process and are repeated for every batch in every epoch.

The article "01. PyTorch Workflow Fundamentals" provides a fantastic, clear summary of these steps in a table.

Here are the five steps for each batch:

1. Forward Pass: Feed the input data from the batch to the model to get its predictions.
   y_pred = model(X_batch)

2. Calculate Loss: Compare the model's predictions with the true labels using the loss function.
   loss = criterion(y_pred, y_labels)

3. Zero Gradients: Reset the optimizer's gradients. PyTorch accumulates gradients by default, so we need to clear them before the next backpropagation.
   optimizer.zero_grad()

4. Backward Pass (Backpropagation): Calculate the gradients of the loss with respect to all of the model's parameters.
   loss.backward()

5. Update Weights (Gradient Descent): The optimizer uses the computed gradients to update the model's parameters.
   optimizer.step()

This five-step process is the core of how a neural network "learns".

4. Step 3: Evaluation Mode - Checking Performance

During training, we also need to check how well our model is performing on data it has never seen before (the test set). This is crucial to ensure the model is generalizing and not just memorizing the training data.

When evaluating, we switch the model to evaluation mode using model.eval(). This is important because some layers, like Dropout and BatchNorm, behave differently during training and testing. We also want to disable gradient calculation to save memory and computation, which we do using the torch.inference_mode() context manager.

The evaluation loop is simpler; it only contains the forward pass and loss calculation.

1. Forward Pass: test_pred = model(X_test_batch)
2. Calculate Loss: test_loss = criterion(test_pred, y_test_labels)

5. Putting It All Together: A Complete Training Script

Now, let's combine everything into a complete training script. We will continue with the MNIST classifier from the previous lesson.

This full workflow is demonstrated beautifully in the "Putting it all together" section of the "01. PyTorch Workflow Fundamentals" article, which integrates data preparation, model building, and a device-agnostic training loop. We will adapt that structure for our FNN.

Here is the full code to train our NeuralNet on the MNIST dataset.

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

# 1. Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 2. Hyperparameters
input_size = 784  # 28x28
hidden_size = 100
num_classes = 10
num_epochs = 5
batch_size = 100
learning_rate = 0.001

# 3. Load MNIST dataset and create DataLoaders
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

# 4. Define the Neural Network model (from previous lesson)
class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size) 
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, num_classes)  
    
    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out

model = NeuralNet(input_size, hidden_size, num_classes).to(device)

# 5. Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  

# 6. The Training Loop
n_total_steps = len(train_loader)
for epoch in range(num_epochs):
    # Set model to training mode
    model.train()
    
    for i, (images, labels) in enumerate(train_loader):  
        # Reshape images to (batch_size, input_size) and move to device
        images = images.reshape(-1, 28*28).to(device)
        labels = labels.to(device)
        
        # --- Core Training Steps ---
        # 1. Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # 2. Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if (i+1) % 100 == 0:
            print (f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}')

# 7. The Testing/Evaluation Loop
with torch.inference_mode():
    n_correct = 0
    n_samples = 0
    # Set model to evaluation mode
    model.eval()
    
    for images, labels in test_loader:
        images = images.reshape(-1, 28*28).to(device)
        labels = labels.to(device)
        
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1) # Get the index of the max log-probability
        n_samples += labels.size(0)
        n_correct += (predicted == labels).sum().item()

    acc = 100.0 * n_correct / n_samples
    print(f'Accuracy of the network on the 10000 test images: {acc} %')

Running this script will train your model and print its final accuracy on the test set. You have now built a complete deep learning workflow!

Conclusion

Congratulations! You have successfully implemented a complete training loop in PyTorch. This is a monumental step in your deep learning journey. The structure you learned today is the foundation for training almost any neural network, from simple FNNs to the most complex models.

Key Takeaways:

• A training process is organized into epochs (full passes over the data) and batches (smaller subsets of data).
• DataLoader is the standard PyTorch tool for creating iterable batches from a Dataset.
• The core training step for each batch involves five steps: forward pass, loss calculation, zeroing gradients, backward pass, and optimizer step.
• It is crucial to switch the model between training (model.train()) and evaluation (model.eval()) modes.
• During evaluation, use torch.inference_mode() to disable gradient calculations for efficiency and safety.

Preview of the Next Lesson:
The training loop we've built is highly modular. We can swap out the model architecture without changing the loop itself. In our next lesson, "Describe the architecture of Convolutional Neural Networks (CNNs) for image data," we will do just that. We will learn about a new, powerful type of model specifically designed for image data, and you'll see how it can be trained using the very same loop you built today.