CNN

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

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

# 1) Data + augmentation
mean, std = (0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)
train_tf = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean, std),
])
test_tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)])

train_ds = datasets.CIFAR10('./data', train=True,  download=True, transform=train_tf)
test_ds  = datasets.CIFAR10('./data', train=False, download=True, transform=test_tf)

train_dl = DataLoader(train_ds, batch_size=128, shuffle=True,  num_workers=4, pin_memory=True)
test_dl  = DataLoader(test_ds,  batch_size=256, shuffle=False, num_workers=4, pin_memory=True)

# 2) Model — small VGG-style
class ConvBlock(nn.Module):
    def __init__(self, in_c, out_c):
        super().__init__()
        self.c1 = nn.Conv2d(in_c, out_c, 3, padding=1, bias=False)
        self.b1 = nn.BatchNorm2d(out_c)
        self.c2 = nn.Conv2d(out_c, out_c, 3, padding=1, bias=False)
        self.b2 = nn.BatchNorm2d(out_c)
    def forward(self, x):
        x = F.relu(self.b1(self.c1(x)))
        x = F.relu(self.b2(self.c2(x)))
        return F.max_pool2d(x, 2)

class SmallVGG(nn.Module):
    def __init__(self, n_classes=10):
        super().__init__()
        self.block1 = ConvBlock(3,   32)
        self.block2 = ConvBlock(32,  64)
        self.block3 = ConvBlock(64, 128)
        self.head = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 256), nn.ReLU(), nn.Dropout(0.5),
            nn.Linear(256, n_classes),
        )
    def forward(self, x):
        return self.head(self.block3(self.block2(self.block1(x))))

model = SmallVGG().to(device)
opt   = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=30)
loss_fn = nn.CrossEntropyLoss()
scaler  = torch.amp.GradScaler('cuda')

# 3) Train
best = 0.0
for epoch in range(1, 31):
    model.train()
    for x, y in train_dl:
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        opt.zero_grad(set_to_none=True)
        with torch.amp.autocast('cuda', dtype=torch.float16):
            loss = loss_fn(model(x), y)
        scaler.scale(loss).backward()
        scaler.step(opt)
        scaler.update()
    sched.step()

    # 4) Eval
    model.eval()
    correct, total = 0, 0
    with torch.inference_mode():
        for x, y in test_dl:
            x, y = x.to(device), y.to(device)
            correct += (model(x).argmax(1) == y).sum().item()
            total   += y.size(0)
    acc = correct / total
    print(f'ep {epoch}  test_acc={acc:.4f}')
    if acc > best:
        best = acc
        torch.save({'model': model.state_dict(), 'acc': acc}, 'best.pt')

# 5) Transfer learning — alternative starting point
from torchvision.models import resnet18, ResNet18_Weights
base = resnet18(weights=ResNet18_Weights.DEFAULT)
base.fc = nn.Linear(base.fc.in_features, 10)
# Optionally freeze backbone:
# for p in base.parameters(): p.requires_grad = False
# for p in base.fc.parameters(): p.requires_grad = True

# 6) Predict
model.eval()
with torch.inference_mode():
    probs = F.softmax(model(x_one.unsqueeze(0).to(device)), dim=1)

import torch.nn as nn
cnn = nn.Sequential(
    nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
    nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1),
    nn.Flatten(), nn.Linear(64, 10),
)