LeNet5 PyTorch Implementation with Corrected Input Size for Fully Connected Layer

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary


class LeNet5(nn.Module):
    def __init__(self, num_classes):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3,
                               out_channels=6,
                               kernel_size=(3, 3),
                               stride=(1, 1),
                               )
        self.pool1 = nn.MaxPool2d(kernel_size=2,
                                  stride=2)
        self.conv2 = nn.Conv2d(in_channels=6,
                               out_channels=16,
                               kernel_size=3,
                               stride=1)
        self.pool2 = nn.MaxPool2d(kernel_size=2,
                                  stride=2)
        self.conv3 = nn.Conv2d(in_channels=16,
                               out_channels=32,
                               kernel_size=3,
                               stride=1)
        self.pool3 = nn.MaxPool2d(kernel_size=2,
                                  stride=2)
        self.fc1 = nn.Linear(32*32*32, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        print(x.size())
        x = F.relu(self.conv3(x))
        print(x.size())
        x = self.pool3(x)
        print(x.size())
        x = x.view(-1, 32*32*32)

        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


model = LeNet5(num_classes=10)
print(model)
summary(model, (3, 32, 32))

This code provides a corrected implementation of LeNet5 in PyTorch. The critical change is the input size for the first fully connected layer ('fc1'). The original code incorrectly used 36 as the input size. This has been corrected to 323232, which represents the actual output size of the convolutional and pooling layers before flattening.

This correction ensures that the model receives the correct dimensions for its fully connected layers, allowing for proper training and prediction. The code also includes print statements to visualize the intermediate output sizes after each layer, facilitating debugging and understanding the model's architecture.