AlexNet

import time
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import torch
from torch import nn, optim
from torch.utils.data import DataLoader, Subset
from torchvision.datasets.cifar import CIFAR10
from torchvision import transforms as T

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

train_transform = T.Compose([T.Resize((50, 50)), 
                             T.ToTensor(),
                             T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

train_val_dataset = CIFAR10(root='../datasets', download=True, train=True, transform=train_transform)

# split dataset into train and validate
indices = list(range(len(train_val_dataset)))
train_idxs, val_idxs = train_test_split(
    indices, test_size=0.1, stratify=train_val_dataset.targets
)

train_dataset = Subset(train_val_dataset, train_idxs)
val_dataset = Subset(train_val_dataset, val_idxs)

# The batch size of 128 images is taken from the original paper.
batch_size=128
train_dataloader = DataLoader(
    dataset=train_dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=4,
    drop_last=True,
)
val_dataloader = DataLoader(
    dataset=val_dataset,
    batch_size=batch_size,
    shuffle=False,
    num_workers=4,
    drop_last=False,
)

class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.feature_extractor = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=2, bias=False),
            nn.BatchNorm2d(num_features=96),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.ReLU(),
            nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, stride=1, padding=2, bias=False),
            nn.BatchNorm2d(num_features=256),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.ReLU(),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(p=0.5),
            # in AlexNet we would have used 6x6x256
            nn.Linear(256, 4096),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Linear(4096, 10)
        )
        
    def forward(self, x):
        x = self.feature_extractor(x)
        x = self.classifier(x)
        return x

def track_performance(dataloader, model, criterion):
    # switch to evaluation mode
    model.eval()
    num_samples = 0
    num_correct = 0
    loss_sum = 0

    # no need to calculate gradients
    with torch.inference_mode():
        for _, (features, labels) in enumerate(dataloader):
            with torch.autocast(device_type="cuda", dtype=torch.float16):
                features = features.to(device)
                labels = labels.to(device)
                logits = model(features)

                predictions = logits.max(dim=1)[1]
                num_correct += (predictions == labels).sum().item()

                loss = criterion(logits, labels)
                loss_sum += loss.cpu().item()
                num_samples += len(features)

    # we return the average loss and the accuracy
    return loss_sum / num_samples, num_correct / num_samples

def train(
    num_epochs,
    train_dataloader,
    val_dataloader,
    model,
    criterion,
    optimizer,
    scheduler=None,
):
    model.to(device)
    scaler = torch.cuda.amp.GradScaler()
    for epoch in range(num_epochs):
        start_time = time.time()
        for _, (features, labels) in enumerate(train_dataloader):
            model.train()
            features = features.to(device)
            labels = labels.to(device)

            # Empty the gradients
            optimizer.zero_grad()

            with torch.autocast(device_type="cuda", dtype=torch.float16):
                # Forward Pass
                logits = model(features)
                # Calculate Loss
                loss = criterion(logits, labels)

            # Backward Pass
            scaler.scale(loss).backward()

            # Gradient Descent
            scaler.step(optimizer)
            scaler.update()

        val_loss, val_acc = track_performance(val_dataloader, model, criterion)
        end_time = time.time()

        s = (
            f"Epoch: {epoch+1:>2}/{num_epochs} | "
            f"Epoch Duration: {end_time - start_time:.3f} sec | "
            f"Val Loss: {val_loss:.5f} | "
            f"Val Acc: {val_acc:.3f} |"
        )
        print(s)

        if scheduler:
            scheduler.step(val_loss)

model = Model()
optimizer = optim.Adam(params=model.parameters(), lr=0.001)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, factor=0.1, patience=2, verbose=True
)
criterion = nn.CrossEntropyLoss(reduction="sum")

train(
    num_epochs=30,
    train_dataloader=train_dataloader,
    val_dataloader=val_dataloader,
    model=model,
    criterion=criterion,
    optimizer=optimizer,
    scheduler=scheduler,
)

Epoch:  1/30 | Epoch Duration: 6.760 sec | Val Loss: 1.61129 | Val Acc: 0.389 |
Epoch:  2/30 | Epoch Duration: 5.634 sec | Val Loss: 1.30560 | Val Acc: 0.522 |
Epoch:  3/30 | Epoch Duration: 5.695 sec | Val Loss: 1.24179 | Val Acc: 0.551 |
Epoch:  4/30 | Epoch Duration: 5.671 sec | Val Loss: 1.18613 | Val Acc: 0.593 |
Epoch:  5/30 | Epoch Duration: 5.686 sec | Val Loss: 1.07944 | Val Acc: 0.624 |
Epoch:  6/30 | Epoch Duration: 5.742 sec | Val Loss: 1.01913 | Val Acc: 0.642 |
Epoch:  7/30 | Epoch Duration: 5.729 sec | Val Loss: 1.00196 | Val Acc: 0.649 |
Epoch:  8/30 | Epoch Duration: 5.728 sec | Val Loss: 0.96261 | Val Acc: 0.682 |
Epoch:  9/30 | Epoch Duration: 5.761 sec | Val Loss: 0.92294 | Val Acc: 0.689 |
Epoch: 10/30 | Epoch Duration: 5.785 sec | Val Loss: 0.95872 | Val Acc: 0.690 |
Epoch: 11/30 | Epoch Duration: 5.744 sec | Val Loss: 0.93677 | Val Acc: 0.692 |
Epoch: 12/30 | Epoch Duration: 5.781 sec | Val Loss: 1.03572 | Val Acc: 0.669 |
Epoch 00012: reducing learning rate of group 0 to 1.0000e-04.
Epoch: 13/30 | Epoch Duration: 5.758 sec | Val Loss: 0.86987 | Val Acc: 0.734 |
Epoch: 14/30 | Epoch Duration: 5.712 sec | Val Loss: 0.87811 | Val Acc: 0.732 |
Epoch: 15/30 | Epoch Duration: 5.748 sec | Val Loss: 0.91566 | Val Acc: 0.731 |
Epoch: 16/30 | Epoch Duration: 5.750 sec | Val Loss: 0.95146 | Val Acc: 0.730 |
Epoch 00016: reducing learning rate of group 0 to 1.0000e-05.
Epoch: 17/30 | Epoch Duration: 5.758 sec | Val Loss: 0.95101 | Val Acc: 0.733 |
Epoch: 18/30 | Epoch Duration: 5.755 sec | Val Loss: 0.95624 | Val Acc: 0.734 |
Epoch: 19/30 | Epoch Duration: 5.753 sec | Val Loss: 0.96301 | Val Acc: 0.732 |
Epoch 00019: reducing learning rate of group 0 to 1.0000e-06.
Epoch: 20/30 | Epoch Duration: 5.735 sec | Val Loss: 0.95980 | Val Acc: 0.735 |
Epoch: 21/30 | Epoch Duration: 5.730 sec | Val Loss: 0.96182 | Val Acc: 0.733 |
Epoch: 22/30 | Epoch Duration: 5.762 sec | Val Loss: 0.96726 | Val Acc: 0.733 |
Epoch 00022: reducing learning rate of group 0 to 1.0000e-07.
Epoch: 23/30 | Epoch Duration: 5.747 sec | Val Loss: 0.96622 | Val Acc: 0.733 |
Epoch: 24/30 | Epoch Duration: 5.780 sec | Val Loss: 0.96463 | Val Acc: 0.734 |
Epoch: 25/30 | Epoch Duration: 5.787 sec | Val Loss: 0.96622 | Val Acc: 0.733 |
Epoch 00025: reducing learning rate of group 0 to 1.0000e-08.
Epoch: 26/30 | Epoch Duration: 5.790 sec | Val Loss: 0.96433 | Val Acc: 0.733 |
Epoch: 27/30 | Epoch Duration: 5.737 sec | Val Loss: 0.96556 | Val Acc: 0.732 |
Epoch: 28/30 | Epoch Duration: 5.785 sec | Val Loss: 0.96416 | Val Acc: 0.733 |
Epoch: 29/30 | Epoch Duration: 5.795 sec | Val Loss: 0.96417 | Val Acc: 0.734 |
Epoch: 30/30 | Epoch Duration: 5.799 sec | Val Loss: 0.96674 | Val Acc: 0.734 |