How to Plot and Analyze Model Results in PyTorch

Analyzing and visualizing model results is crucial for understanding how well a model is performing and where improvements might be needed. Particularly in machine learning with libraries like PyTorch, plotting results can help in interpreting the data and model diagnostics. This guide will walk you through how to plot and analyze model results using PyTorch, with complete code snippets and explanations.

Prerequisites

Before diving into the visualization part, make sure you have the following libraries installed:

pip install torch torchvision matplotlib numpy

We'll use a simple neural network model built with PyTorch and visualize its performance metrics using Python’s popular plotting library, Matplotlib.

1. Build a Simple Neural Network with PyTorch

Firstly, let's create a simple neural network. We'll use a dataset like the MNIST, which is stored in PyTorch's torchvision package, to train this model.

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

# Define a simple feedforward neural network
class SimpleNN(nn.Module):
    def __init__(self, input_size=784, hidden_size=128, num_classes=10):
        super(SimpleNN, 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

Now, let's set up the data loaders for the MNIST dataset:

# Load the MNIST dataset
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
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)

2. Train the Neural Network Model

Let's train our simple neural network using a training loop. Below is an example of a simple training loop:

model = SimpleNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

epochs = 2  # Define a small number for quick example
def train_model():
    model.train()
    for epoch in range(epochs):
        for batch_idx, (images, labels) in enumerate(train_loader):
            # Flatten the images from (batch_size, 1, 28, 28) to (batch_size, 784)
            images = images.reshape(-1, 28*28)
            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)
            # Backward pass and optimization
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if (batch_idx+1) % 100 == 0:
                print(f'Epoch [{epoch+1}/{epochs}], Step [{batch_idx+1}/len(train_loader)], Loss: {loss.item():.4f}')

train_model()

3. Plotting Results Using Matplotlib

After training the model, you may want to see how well its predictions are doing. For this example, we'll visualize a few predictions against their true labels. This simple plot can help you intuitively understand model performance for specific examples.

import matplotlib.pyplot as plt

def plot_some_results():
    # Set the model to evaluation mode
    model.eval()
    # Get images and label from a single batch
    images, labels = next(iter(train_loader))
    images = images.reshape(-1, 28*28)
    # Forward pass
    outputs = model(images)
    _, predicted = torch.max(outputs, 1)
    # Visualizing 6 images
    fig, axes = plt.subplots(1, 6, figsize=(12, 2))
    for i in range(6):
        axes[i].imshow(images[i].reshape(28, 28).detach().numpy(), cmap='gray')
        axes[i].set_title(f'Pred: {predicted[i]}, True: {labels[i]}')
        axes[i].axis('off')
    plt.show()

plot_some_results()

4. Plotting Loss and Accuracy

Tracking the loss and accuracy over epochs is crucial during model development. In this case, we need to modify our training loop to store and later visualize these metrics. A simple modification can help track metrics like this:

epoch_losses = []
accuracy_list = []

for epoch in range(epochs):
    epoch_loss = 0
    correct = 0
    total = 0
    for batch_idx, (images, labels) in enumerate(train_loader):
        images = images.reshape(-1, 28*28)
        outputs = model(images)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()
        total += labels.size(0)
        _, predicted = torch.max(outputs.data, 1)
        correct += (predicted == labels).sum().item()

    epoch_losses.append(epoch_loss/len(train_loader))
    accuracy_list.append(correct/total)

    print(f'Epoch {epoch+1}, Loss: {epoch_losses[-1]:.4f}, Accuracy: {accuracy_list[-1]:.4f}')

With saved metrics, you can plot them using Matplotlib:

plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
plt.plot(range(1, epochs+1), epoch_losses, label='Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Loss Over Time')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(range(1, epochs+1), accuracy_list, label='Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.title('Accuracy Over Time')
plt.legend()

plt.show()

These plots offer insights into the model's learning and effectiveness, offering early cues into overfitting or the need for hyperparameter tuning. Exploring these kinds of visualizations is critical for understanding and fixing training issues, optimizing performance, and ensuring robustness over time. Remember, good practices of data visualization integrate directly with logical analysis for iterative development.