S102 · CNNs — Convolution, Pooling, ImageNet Architectures

# cifar_cnn.py — train a CNN on CIFAR-10, print accuracy per epoch.
# Run: uv run cifar_cnn.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
 
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BATCH = 128
EPOCHS = 5
LR = 1e-3
 
 
class SmallCNN(nn.Module):
    """A 4-block CNN. 3 → 64 → 128 → 256 → 512 channels, pooling halves spatial dims each block."""
 
    def __init__(self, num_classes: int = 10):
        super().__init__()
        # Block 1: 32×32 → 16×16
        self.conv1a = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv1b = nn.Conv2d(64, 64, kernel_size=3, padding=1)
        # Block 2: 16×16 → 8×8
        self.conv2a = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv2b = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        # Block 3: 8×8 → 4×4
        self.conv3a = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.conv3b = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        # Head
        self.pool = nn.MaxPool2d(2, 2)
        self.dropout = nn.Dropout(0.3)
        self.fc1 = nn.Linear(256 * 4 * 4, 512)
        self.fc2 = nn.Linear(512, num_classes)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Block 1
        x = F.relu(self.conv1a(x))
        x = F.relu(self.conv1b(x))
        x = self.pool(x)  # 32 → 16
        # Block 2
        x = F.relu(self.conv2a(x))
        x = F.relu(self.conv2b(x))
        x = self.pool(x)  # 16 → 8
        # Block 3
        x = F.relu(self.conv3a(x))
        x = F.relu(self.conv3b(x))
        x = self.pool(x)  # 8 → 4
        # Head
        x = x.flatten(1)  # [B, 256*4*4]
        x = self.dropout(F.relu(self.fc1(x)))
        return self.fc2(x)
 
 
def get_loaders() -> tuple[DataLoader, DataLoader]:
    tfm_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
    ])
    tfm_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
    ])
    train = datasets.CIFAR10("./data", train=True, download=True, transform=tfm_train)
    test = datasets.CIFAR10("./data", train=False, download=True, transform=tfm_test)
    return (
        DataLoader(train, batch_size=BATCH, shuffle=True, num_workers=2),
        DataLoader(test, batch_size=BATCH, shuffle=False, num_workers=2),
    )
 
 
def evaluate(model: nn.Module, loader: DataLoader) -> float:
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for x, y in loader:
            x, y = x.to(DEVICE), y.to(DEVICE)
            pred = model(x).argmax(dim=1)
            correct += (pred == y).sum().item()
            total += y.size(0)
    return 100.0 * correct / total
 
 
def main() -> None:
    train_loader, test_loader = get_loaders()
    model = SmallCNN().to(DEVICE)
    opt = optim.Adam(model.parameters(), lr=LR)
    loss_fn = nn.CrossEntropyLoss()
 
    params = sum(p.numel() for p in model.parameters())
    print(f"device={DEVICE}  params={params:,}")
 
    for epoch in range(1, EPOCHS + 1):
        model.train()
        running = 0.0
        for step, (x, y) in enumerate(train_loader, 1):
            x, y = x.to(DEVICE), y.to(DEVICE)
            opt.zero_grad()
            loss = loss_fn(model(x), y)
            loss.backward()
            opt.step()
            running += loss.item()
        acc = evaluate(model, test_loader)
        print(f"epoch {epoch}  train_loss={running/step:.3f}  test_acc={acc:.2f}%")
 
 
if __name__ == "__main__":
    main()