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

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

torch.manual_seed(42);
np.random.seed(42)

CIFAR10_MEAN = [0.4914, 0.4822, 0.4465]
CIFAR10_STD = [0.2470, 0.2435, 0.2616]

train_transform = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(
        brightness=0.2,
        contrast=0.2,
        saturation=0.2,
        hue=0.1
    ),
    transforms.ToTensor(),
    transforms.Normalize(CIFAR10_MEAN, CIFAR10_STD),
    transforms.RandomErasing(p=0.5, scale=(0.02, 0.2), ratio=(0.3, 3.3))
])

test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(CIFAR10_MEAN, CIFAR10_STD)
])

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

test_dataset = datasets.CIFAR10(
    root='./data',
    train=False,
    download=True,
    transform=test_transform
)

print(f"Training dataset size: {len(train_dataset)}")
print(f"Test dataset size: {len(test_dataset)}")
print(f"Number of classes: {len(train_dataset.classes)}")
print(f"Classes: {train_dataset.classes}")

train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=128,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
    persistent_workers=True
)

test_loader = DataLoader(
    dataset=test_dataset,
    batch_size=128,
    shuffle=False,
    num_workers=4,
    pin_memory=True,
    persistent_workers=True
)

100%|██████████| 170M/170M [00:14<00:00, 11.9MB/s]

Training dataset size: 50000
Test dataset size: 10000
Number of classes: 10
Classes: ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

class AlexNetCIFAR(nn.Module):
  def __init__(self, num_classes=10):
    super(AlexNetCIFAR, self).__init__()

    self.features = nn.Sequential(
      # Conv1: 32x32x3 -> 32x32x64
      nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
      nn.BatchNorm2d(64),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(kernel_size=2, stride=2),  # 32x32 -> 16x16

      # Conv2: 16x16x64 -> 16x16x192
      nn.Conv2d(64, 192, kernel_size=3, padding=1),
      nn.BatchNorm2d(192),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(kernel_size=2, stride=2),  # 16x16 -> 8x8

      # Conv3: 8x8x192 -> 8x8x384
      nn.Conv2d(192, 384, kernel_size=3, padding=1),
      nn.BatchNorm2d(384),
      nn.ReLU(inplace=True),

      # Conv4: 8x8x384 -> 8x8x256
      nn.Conv2d(384, 256, kernel_size=3, padding=1),
      nn.BatchNorm2d(256),
      nn.ReLU(inplace=True),

      # Conv5: 8x8x256 -> 8x8x256
      nn.Conv2d(256, 256, kernel_size=3, padding=1),
      nn.BatchNorm2d(256),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(kernel_size=2, stride=2),  # 8x8 -> 4x4
    )

    self.classifier = nn.Sequential(
      nn.Dropout(0.5),
      nn.Linear(256 * 4 * 4, 1024),
      nn.BatchNorm1d(1024),
      nn.ReLU(inplace=True),
      nn.Dropout(0.5),
      nn.Linear(1024, 512),
      nn.BatchNorm1d(512),
      nn.ReLU(inplace=True),
      nn.Linear(512, num_classes),
    )

    self._initialize_weights()

  def forward(self, x):
    x = self.features(x)
    x = torch.flatten(x, 1)
    x = self.classifier(x)
    return x


  def _initialize_weights(self):
    for m in self.modules():
      if isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        if m.bias is not None:
          nn.init.constant_(m.bias, 0)
      elif isinstance(m, nn.Linear):
        nn.init.normal_(m.weight, 0, 0.01)
        nn.init.constant_(m.bias, 0)

model = AlexNetCIFAR(num_classes=10).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30, 60, 80], gamma=0.2)

for epoch in range(100):
    model.train()
    train_loss = 0.0
    train_correct = 0
    train_total = 0

    train_loop = tqdm(train_loader, desc=f"Epoch {epoch} Training")
    for batch_idx, (inputs, targets) in enumerate(train_loop):
        inputs, targets = inputs.to(device), targets.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        train_total += targets.size(0)
        train_correct += predicted.eq(targets).sum().item()

        train_loop.set_postfix({
            'loss': f'{train_loss/(batch_idx+1):.3f}',
            'acc': f'{100.*train_correct/train_total:.2f}%'
        })

    model.eval()
    test_correct = 0
    test_total = 0
    test_loss = 0.0

    with torch.no_grad():
        test_loop = tqdm(test_loader, desc=f"Epoch {epoch} Validation")
        for batch_idx, (inputs, targets) in enumerate(test_loop):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            test_loss += loss.item()
            _, predicted = outputs.max(1)
            test_total += targets.size(0)
            test_correct += predicted.eq(targets).sum().item()

            test_loop.set_postfix({
                'loss': f'{test_loss/(batch_idx+1):.3f}',
                'acc': f'{100.*test_correct/test_total:.2f}%'
            })

    test_acc = 100. * test_correct / test_total
    print(f'Epoch {epoch}: Test Accuracy: {test_acc:.2f}%')
    scheduler.step()

Epoch 0 Training: 100%|██████████| 391/391 [00:15<00:00, 26.02it/s, loss=1.760, acc=37.10%]
Epoch 0 Validation: 100%|██████████| 79/79 [00:00<00:00, 94.40it/s, loss=1.312, acc=52.83%]

Epoch 0: Test Accuracy: 52.83%

Epoch 1 Training: 100%|██████████| 391/391 [00:13<00:00, 29.78it/s, loss=1.310, acc=53.09%]
Epoch 1 Validation: 100%|██████████| 79/79 [00:00<00:00, 113.45it/s, loss=1.144, acc=60.14%]

Epoch 1: Test Accuracy: 60.14%

Epoch 2 Training: 100%|██████████| 391/391 [00:13<00:00, 29.45it/s, loss=1.127, acc=59.91%]
Epoch 2 Validation: 100%|██████████| 79/79 [00:00<00:00, 107.99it/s, loss=0.964, acc=66.21%]

Epoch 2: Test Accuracy: 66.21%

Epoch 3 Training: 100%|██████████| 391/391 [00:13<00:00, 29.83it/s, loss=1.013, acc=64.26%]
Epoch 3 Validation: 100%|██████████| 79/79 [00:00<00:00, 110.76it/s, loss=0.886, acc=67.96%]

Epoch 3: Test Accuracy: 67.96%

Epoch 4 Training: 100%|██████████| 391/391 [00:13<00:00, 29.81it/s, loss=0.934, acc=67.30%]
Epoch 4 Validation: 100%|██████████| 79/79 [00:00<00:00, 112.89it/s, loss=0.935, acc=67.54%]

Epoch 4: Test Accuracy: 67.54%