Model Definition

File: model.py

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


class LeNet5Variant(nn.Module):
    def __init__(self, input_channels=1, num_classes=10):
        super().__init__()
        # First convolutional block
        self.conv_block_1 = nn.Sequential(
            nn.Conv2d(input_channels, 6, kernel_size=5, padding=2),
            nn.Sigmoid(),
            nn.AvgPool2d(kernel_size=2, stride=2)
        )
        
        # Second convolutional block
        self.conv_block_2 = nn.Sequential(
            nn.Conv2d(6, 16, kernel_size=5),
            nn.Sigmoid(),
            nn.AvgPool2d(kernel_size=2, stride=2)
        )
        
        # Fully connected layers
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(400, 120),
            nn.Sigmoid(),
            nn.Linear(120, 84),
            nn.Sigmoid(),
            nn.Linear(84, num_classes)
        )

    def forward(self, x):
        x = self.conv_block_1(x)
        x = self.conv_block_2(x)
        x = self.classifier(x)
        return x


if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = LeNet5Variant().to(device)
    print(summary(model, input_size=(1, 28, 28)))

Training Pipeline

File: train.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


def prepare_data(batch_size=128):
    transform = transforms.Compose([
        transforms.Resize(28),
        transforms.ToTensor()
    ])
    
    full_dataset = datasets.FashionMNIST(
        root="./data", train=True, download=True, transform=transform
    )
    
    train_size = int(0.8 * len(full_dataset))
    val_size = len(full_dataset) - train_size
    train_ds, val_ds = random_split(full_dataset, [train_size, val_size])
    
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=4)
    
    return train_loader, val_loader


def train_epoch(model, dataloader, criterion, optimizer, device):
    model.train()
    total_loss, correct, total = 0.0, 0, 0
    
    for inputs, labels in dataloader:
        inputs, labels = inputs.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item() * inputs.size(0)
        _, preds = torch.max(outputs, 1)
        correct += (preds == labels).sum().item()
        total += inputs.size(0)
    
    return total_loss / total, correct / total


def validate_epoch(model, dataloader, criterion, device):
    model.eval()
    total_loss, correct, total = 0.0, 0, 0
    
    with torch.no_grad():
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            total_loss += loss.item() * inputs.size(0)
            _, preds = torch.max(outputs, 1)
            correct += (preds == labels).sum().item()
            total += inputs.size(0)
    
    return total_loss / total, correct / total


def train_model(model, train_loader, val_loader, epochs=20):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    
    history = {
        "train_loss": [], "val_loss": [],
        "train_acc": [], "val_acc": []
    }
    
    best_val_acc = 0.0
    best_weights = None
    
    for epoch in range(epochs):
        train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
        val_loss, val_acc = validate_epoch(model, val_loader, criterion, device)
        
        history["train_loss"].append(train_loss)
        history["val_loss"].append(val_loss)
        history["train_acc"].append(train_acc)
        history["val_acc"].append(val_acc)
        
        print(f"Epoch {epoch+1}/{epochs} | Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} | "
              f"Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}")
        
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_weights = model.state_dict().copy()
    
    if best_weights:
        torch.save(best_weights, "lenet_fashionmnist_best.pth")
    
    return pd.DataFrame(history)


def plot_training_history(df):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    
    ax1.plot(df["train_loss"], "o-", label="Training Loss")
    ax1.plot(df["val_loss"], "s--", label="Validation Loss")
    ax1.set_xlabel("Epoch")
    ax1.set_ylabel("Loss")
    ax1.legend()
    ax1.grid(True)
    
    ax2.plot(df["train_acc"], "o-", label="Training Accuracy")
    ax2.plot(df["val_acc"], "s--", label="Validation Accuracy")
    ax2.set_xlabel("Epoch")
    ax2.set_ylabel("Accuracy")
    ax2.legend()
    ax2.grid(True)
    
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    train_dl, val_dl = prepare_data()
    net = LeNet5Variant()
    history_df = train_model(net, train_dl, val_dl, epochs=20)
    plot_training_history(history_df)

Key Implementation Notes

• The architecture follows the original LeNet-5 design but replaces manual layer sequencing with nn.Sequential blocks for improved readability and maintainability.
• Input dimensions assume grayscale Fashion-MNIST images resized to 28×28; padding ensures spatial dimensions remain consistent after first convolution.
• Average pooling is used instead of max pooling, preserving the historical LeNet design choice.
• Sigmoid activation is applied after each convolutional block and fully connected layer, matching the original nonlinearity.
• The classifier flattens the output from the final poooling layer (16×5×5 → 400) before feeding into dense layers.
• Training uses cross-entropy loss and Adam optimization, with early validation-based checkpointing.
• History tracking and visualization are decoupled into reusable functions for clarity and reproducibility.