The Project | About Us | Contribute | Donations | License
HOME My courses Completed
◄ Previous
◄

Training Loop

Next ►
►

In this section, we will delve into the core of training neural networks in PyTorch: the training loop. The training loop is where the model learns from the data by iteratively updating its parameters to minimize the loss function. This process involves several key steps, which we will break down and explain in detail.

Key Concepts

  1. Epoch: One complete pass through the entire training dataset.
  2. Batch: A subset of the training data used to update the model's parameters.
  3. Forward Pass: Calculating the output of the neural network.
  4. Loss Calculation: Measuring the difference between the predicted output and the actual target.
  5. Backward Pass: Computing the gradients of the loss with respect to the model's parameters.
  6. Parameter Update: Adjusting the model's parameters using the computed gradients.

Steps in a Training Loop

  1. Initialize the Model, Loss Function, and Optimizer
  2. Iterate Over the Dataset
  3. Perform Forward Pass
  4. Compute Loss
  5. Perform Backward Pass
  6. Update Parameters
  7. Track and Print Metrics

Example Code

Let's look at a practical example of a training loop in PyTorch. We'll use a simple neural network to demonstrate the process.

Step 1: Initialize the Model, Loss Function, and Optimizer

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

# Define a simple neural network
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(10, 50)
        self.fc2 = nn.Linear(50, 1)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Initialize the model
model = SimpleNN()

# Define the loss function
criterion = nn.MSELoss()

# Define the optimizer
optimizer = optim.SGD(model.parameters(), lr=0.01)

Step 2: Iterate Over the Dataset

# Dummy dataset
data = torch.randn(100, 10)  # 100 samples, 10 features each
targets = torch.randn(100, 1)  # 100 target values

# Number of epochs
num_epochs = 20

# Batch size
batch_size = 10

# Data loader
data_loader = torch.utils.data.DataLoader(
    dataset=list(zip(data, targets)),
    batch_size=batch_size,
    shuffle=True
)

Step 3: Perform Forward Pass

Step 4: Compute Loss

Step 5: Perform Backward Pass

Step 6: Update Parameters

Step 7: Track and Print Metrics

# Training loop
for epoch in range(num_epochs):
    for batch_data, batch_targets in data_loader:
        # Zero the gradients
        optimizer.zero_grad()
        
        # Forward pass
        outputs = model(batch_data)
        
        # Compute loss
        loss = criterion(outputs, batch_targets)
        
        # Backward pass
        loss.backward()
        
        # Update parameters
        optimizer.step()
    
    # Print epoch and loss
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

Common Mistakes and Tips

  1. Forgetting to Zero Gradients: Always call optimizer.zero_grad() before the forward pass to clear the old gradients.
  2. Incorrect Loss Calculation: Ensure that the loss function matches the problem type (e.g., MSE for regression, CrossEntropy for classification).
  3. Learning Rate: Choosing an appropriate learning rate is crucial. Too high can cause divergence, too low can slow down training.
  4. Batch Size: Larger batch sizes can speed up training but require more memory. Smaller batch sizes can provide more accurate gradient estimates.

Practical Exercise

Exercise: Implement a Training Loop

  1. Define a neural network with two hidden layers.
  2. Use the MNIST dataset for training.
  3. Implement the training loop with the following specifications:
    • Use CrossEntropyLoss as the loss function.
    • Use Adam optimizer.
    • Train for 10 epochs.
    • Print the loss every epoch.

Solution

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# Define the neural network
class NeuralNet(nn.Module):
    def __init__(self):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(28*28, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)
    
    def forward(self, x):
        x = x.view(-1, 28*28)  # Flatten the input
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Initialize the model
model = NeuralNet()

# Define the loss function
criterion = nn.CrossEntropyLoss()

# Define the optimizer
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Load the MNIST dataset
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)

# Training loop
num_epochs = 10
for epoch in range(num_epochs):
    for batch_data, batch_targets in train_loader:
        # Zero the gradients
        optimizer.zero_grad()
        
        # Forward pass
        outputs = model(batch_data)
        
        # Compute loss
        loss = criterion(outputs, batch_targets)
        
        # Backward pass
        loss.backward()
        
        # Update parameters
        optimizer.step()
    
    # Print epoch and loss
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

Conclusion

In this section, we covered the essential steps involved in a training loop in PyTorch. We discussed the key concepts, provided a detailed example, and highlighted common mistakes and tips. By understanding and implementing these steps, you can effectively train neural networks using PyTorch. In the next section, we will explore validation and testing to evaluate the performance of your trained models.

◄ Previous
◄
Next ►
►
The Project | About Us | Contribute | Donations | License
© Copyright 2024. All rights reserved