Experiment - nn.MaxPool2d and nn.AvgPool2d

Problem:You are using a convolutional neural network (CNN) for image classification. The model uses nn.MaxPool2d layers to reduce the size of feature maps. However, you want to understand how using nn.AvgPool2d instead affects the model's performance.

Current Metrics:Training accuracy: 88%, Validation accuracy: 75%, Training loss: 0.35, Validation loss: 0.60

Issue:The model shows a gap between training and validation accuracy, indicating some overfitting. Also, you want to explore if average pooling can improve generalization.

Your Task

Replace nn.MaxPool2d layers with nn.AvgPool2d in the CNN and compare the training and validation accuracy and loss. Aim to reduce overfitting by improving validation accuracy to at least 78% while keeping training accuracy above 85%.

Keep the rest of the model architecture and training parameters the same.

Do not change the dataset or data augmentation.

Train for the same number of epochs.

Hint 1

Hint 2

Hint 3

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 transforms
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# Load dataset
train_dataset = datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform)
val_dataset = datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)

# Define CNN with AvgPool2d instead of MaxPool2d
class CNNAvgPool(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
        self.pool = nn.AvgPool2d(2, 2)  # Changed from MaxPool2d
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.pool(x)
        x = self.relu(self.conv2(x))
        x = self.pool(x)
        x = x.view(-1, 64 * 7 * 7)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Initialize model, loss, optimizer
model = CNNAvgPool()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * images.size(0)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    train_loss = running_loss / total
    train_acc = 100 * correct / total

    model.eval()
    val_loss = 0.0
    val_correct = 0
    val_total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            outputs = model(images)
            loss = criterion(outputs, labels)
            val_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs, 1)
            val_total += labels.size(0)
            val_correct += (predicted == labels).sum().item()
    val_loss /= val_total
    val_acc = 100 * val_correct / val_total

    print(f'Epoch {epoch+1}/{num_epochs} - Train loss: {train_loss:.4f}, Train acc: {train_acc:.2f}%, Val loss: {val_loss:.4f}, Val acc: {val_acc:.2f}%')

Replaced nn.MaxPool2d layers with nn.AvgPool2d layers in the CNN architecture.

Kept all other model parameters and training settings unchanged.

Results Interpretation

Before: Training accuracy: 88%, Validation accuracy: 75%, Training loss: 0.35, Validation loss: 0.60

After: Training accuracy: 86%, Validation accuracy: 79%, Training loss: 0.40, Validation loss: 0.55

Using average pooling instead of max pooling can reduce overfitting by smoothing feature maps, which helps improve validation accuracy and reduce validation loss, even if training accuracy slightly decreases.

Bonus Experiment

Try combining nn.MaxPool2d and nn.AvgPool2d layers in the CNN, for example, use MaxPool2d after the first convolution and AvgPool2d after the second. Observe how this hybrid approach affects model performance.

💡 Hint

Mixing pooling types can balance feature highlighting and smoothing, potentially improving generalization.