import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, SubsetRandomSampler
from tqdm import tqdm

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

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,
    transform=train_transform,
    download=True
)

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

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

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

100%|██████████| 170M/170M [00:04<00:00, 36.6MB/s]

class VggCIFAR(nn.Module):
  def __init__(self, num_classes=10, width_multiplier=1.2):
    super().__init__()

    c1 = int(64 * width_multiplier)   # 96
    c2 = int(128 * width_multiplier)  # 192
    c3 = int(256 * width_multiplier)  # 384
    c4 = int(512 * width_multiplier)  # 768

    self.features = nn.Sequential(
      # Block 1
      nn.Conv2d(3, c1, 3, padding=1),
      nn.BatchNorm2d(c1),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(2, 2),

      # Block 2
      nn.Conv2d(c1, c2, 3, padding=1),
      nn.BatchNorm2d(c2),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(2, 2),

      # Block 3
      nn.Conv2d(c2, c3, 3, padding=1),
      nn.BatchNorm2d(c3),
      nn.ReLU(inplace=True),
      nn.Conv2d(c3, c3, 3, padding=1),
      nn.BatchNorm2d(c3),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(2, 2),

      # Block 4
      nn.Conv2d(c3, c4, 3, padding=1),
      nn.BatchNorm2d(c4),
      nn.ReLU(inplace=True),
      nn.Conv2d(c4, c4, 3, padding=1),
      nn.BatchNorm2d(c4),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(2, 2)
    )

    self.classifier = nn.Sequential(
      nn.Dropout(0.5),
      nn.Linear(c4 * 2 * 2, 1024),
      nn.ReLU(inplace=True),
      nn.Dropout(0.5),
      nn.Linear(1024, 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')
      elif isinstance(m, nn.BatchNorm2d):
        nn.init.constant_(m.weight, 1)
        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 = VggCIFAR().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)

best_acc = 0
epochs = 100

for epoch in range(epochs):
  model.train()
  running_loss = 0.0
  correct = 0
  total = 0

  train_pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs} [Train]")
  for inputs, targets in train_pbar:
    inputs, targets = inputs.to(device), targets.to(device)
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

    running_loss += loss.item()
    _, predicted = outputs.max(1)
    total += targets.size(0)
    correct += predicted.eq(targets).sum().item()

    train_pbar.set_postfix(loss=running_loss/(total/targets.size(0)), acc=100.*correct/total)

  train_loss = running_loss / len(train_loader)
  train_acc = 100. * correct / total

  model.eval()
  val_loss = 0.0
  val_correct = 0
  val_total = 0

  with torch.no_grad():
    val_pbar = tqdm(val_loader, desc=f"Epoch {epoch+1}/{epochs} [Val]")
    for inputs, targets in val_pbar:
      inputs, targets = inputs.to(device), targets.to(device)
      outputs = model(inputs)
      loss = criterion(outputs, targets)

      val_loss += loss.item()
      _, predicted = outputs.max(1)
      val_total += targets.size(0)
      val_correct += predicted.eq(targets).sum().item()
      val_pbar.set_postfix(loss=val_loss/(val_total/targets.size(0)), acc=100.*val_correct/val_total)

  validation_loss = val_loss / len(val_loader)
  validation_acc = 100. * val_correct / val_total

  print(f"Epoch {epoch+1}/{epochs} | Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}% | Val Loss: {validation_loss:.4f}, Val Acc: {validation_acc:.2f}%")

  scheduler.step()

  if validation_acc > best_acc:
    best_acc = validation_acc
    torch.save(model.state_dict(), 'best_lenet_classic_model.pth')
    print(f"New best model found! Accuracy: {best_acc:.2f}%. Saving model...")

print("\nTraining finished.")
print(f"Best validation accuracy: {best_acc:.2f}%")

Epoch 1/100 [Train]: 100%|██████████| 391/391 [00:14<00:00, 26.62it/s, acc=13.9, loss=1.4]
Epoch 1/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 96.96it/s, acc=16.5, loss=0.27]

Epoch 1/100 | Train Loss: 2.2451, Train Acc: 13.91% | Val Loss: 2.1343, Val Acc: 16.49%
New best model found! Accuracy: 16.49%. Saving model...

Epoch 2/100 [Train]: 100%|██████████| 391/391 [00:12<00:00, 30.09it/s, acc=21.1, loss=1.26]
Epoch 2/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 108.65it/s, acc=30.6, loss=0.225]

Epoch 2/100 | Train Loss: 2.0116, Train Acc: 21.10% | Val Loss: 1.7839, Val Acc: 30.64%
New best model found! Accuracy: 30.64%. Saving model...

Epoch 3/100 [Train]: 100%|██████████| 391/391 [00:13<00:00, 29.84it/s, acc=33.3, loss=1.09]
Epoch 3/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 111.99it/s, acc=41.3, loss=0.199]

Epoch 3/100 | Train Loss: 1.7362, Train Acc: 33.30% | Val Loss: 1.5765, Val Acc: 41.26%
New best model found! Accuracy: 41.26%. Saving model...

Epoch 4/100 [Train]: 100%|██████████| 391/391 [00:13<00:00, 29.75it/s, acc=46.3, loss=0.937]
Epoch 4/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 111.76it/s, acc=53.7, loss=0.166]

Epoch 4/100 | Train Loss: 1.4973, Train Acc: 46.30% | Val Loss: 1.3144, Val Acc: 53.70%
New best model found! Accuracy: 53.70%. Saving model...

Epoch 5/100 [Train]: 100%|██████████| 391/391 [00:13<00:00, 29.88it/s, acc=56.2, loss=0.779]
Epoch 5/100 [Val]: 100%|██████████| 79/79 [00:00<00:00, 100.80it/s, acc=58.6, loss=0.148]

Epoch 5/100 | Train Loss: 1.2455, Train Acc: 56.17% | Val Loss: 1.1747, Val Acc: 58.63%
New best model found! Accuracy: 58.63%. Saving model...