Mini Project: CIFAR-10 Image Classifier

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as T
from torch.utils.data import DataLoader

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

# ─── DATA ─────────────────────────────────────────────
train_tfm = T.Compose([
    T.RandomCrop(32, padding=4),
    T.RandomHorizontalFlip(),
    T.AutoAugment(T.AutoAugmentPolicy.CIFAR10),  # learned augmentation
    T.ToTensor(),
    T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
])
val_tfm = T.Compose([
    T.ToTensor(),
    T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
])

train_ds = torchvision.datasets.CIFAR10("./data", train=True,  download=True, transform=train_tfm)
val_ds   = torchvision.datasets.CIFAR10("./data", train=False, download=True, transform=val_tfm)
train_loader = DataLoader(train_ds, 128, shuffle=True,  num_workers=4, pin_memory=True)
val_loader   = DataLoader(val_ds,   256, shuffle=False, num_workers=4, pin_memory=True)

# ─── MODEL ────────────────────────────────────────────
def conv_bn(in_c, out_c, k=3, s=1, p=1):
    return nn.Sequential(nn.Conv2d(in_c, out_c, k, s, p, bias=False), nn.BatchNorm2d(out_c), nn.ReLU(inplace=True))

class ResBlock(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.net = nn.Sequential(conv_bn(c, c), conv_bn(c, c))
    def forward(self, x): return x + self.net(x)

class CIFAR10Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            conv_bn(3, 64), ResBlock(64), nn.MaxPool2d(2),      # 16x16
            conv_bn(64, 128), ResBlock(128), nn.MaxPool2d(2),   # 8x8
            conv_bn(128, 256), ResBlock(256), nn.MaxPool2d(2),  # 4x4
            nn.AdaptiveAvgPool2d(1), nn.Flatten(),
            nn.Linear(256, 10),
        )
    def forward(self, x): return self.net(x)

model = CIFAR10Net().to(device)
print(f"Params: {sum(p.numel() for p in model.parameters()):,}")

# ─── TRAINING ─────────────────────────────────────────
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.1,
    steps_per_epoch=len(train_loader), epochs=100)
loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)  # label smoothing reduces overconfidence

best_acc = 0.0
for epoch in range(100):
    model.train()
    for imgs, labels in train_loader:
        imgs, labels = imgs.to(device), labels.to(device)
        optimizer.zero_grad()
        loss_fn(model(imgs), labels).backward()
        optimizer.step()
        scheduler.step()

    model.eval()
    correct = total = 0
    with torch.no_grad():
        for imgs, labels in val_loader:
            imgs, labels = imgs.to(device), labels.to(device)
            correct += (model(imgs).argmax(1) == labels).sum().item()
            total   += labels.size(0)

    acc = correct / total
    if acc > best_acc:
        best_acc = acc
        torch.save(model.state_dict(), "cifar10_best.pt")
    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1:3d} | Acc={acc:.2%} | Best={best_acc:.2%}")

# Expected: ~93-95% accuracy. Human-level on CIFAR-10 ≈ 94%.
print(f"Final Best: {best_acc:.2%}")