Bài 31: Xây dựng mô hình phân loại ảnh CIFAR-10 từ đầu

Sau bài học, bạn sẽ:

• Load CIFAR-10 bằng torchvision.datasets với normalize đúng stats và augmentation tiêu chuẩn.
• Ráp một CNN dạng VGG nhỏ: 3 block Conv–BatchNorm–ReLU + MaxPool, classifier head Flatten + Dropout + Linear.
• Tính output shape qua từng block để suy ra in_features của FC đầu tiên.
• Setup optimizer AdamW + scheduler CosineAnnealingLR, loss CrossEntropyLoss.
• Viết training loop có validation, save best checkpoint theo val accuracy.
• Biết debugging checklist (overfit 1 batch, in shape) và common bug (quên model.train()/eval(), sai norm stats, double softmax).
• Đạt val accuracy ~85–88% sau 30 epoch trên Colab T4 và biết hướng đi tiếp để cải thiện.

Bài này nối tiếp Bài 30 — Flatten và Fully Connected và chuẩn bị cho Bài 32 — Transfer Learning, nơi bạn sẽ thay backbone tự build bằng ResNet pretrain trên ImageNet.

CIFAR-10 (Krizhevsky, 2009) là dataset benchmark cổ điển cho image classification, lấy từ subset của 80 Million Tiny Images:

• 60 000 ảnh màu kích thước \( 32 \times 32 \times 3 \).
• 50 000 train + 10 000 test.
• 10 class cân bằng (6 000 ảnh/class): airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck.
• Khó hơn MNIST (ảnh tự nhiên có background, biến dạng, góc nhìn) nhưng vẫn nhỏ và train được trên một GPU consumer trong vài phút.

Vai trò trong roadmap deep learning: CIFAR-10 là môi trường thử cho hầu như mọi kiến trúc CNN mới (NIN, ResNet, DenseNet, Wide ResNet, ViT) trước khi scale lên ImageNet. Số liệu tham khảo:

• Baseline random: 10%.
• CNN nhỏ (bài này): ~85–88%.
• ResNet-110 với augmentation: ~93%.
• SOTA hiện tại (ViT, Wide ResNet + heavy augmentation): ~99.5%.

Dùng torchvision.datasets.CIFAR10 để download và parse. Lần đầu tiên gọi với download=True sẽ kéo về ~170MB.

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Stats trung bình / std mỗi kênh tính trên train set CIFAR-10
CIFAR_MEAN = (0.4914, 0.4822, 0.4465)
CIFAR_STD  = (0.2470, 0.2435, 0.2616)

transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),       # crop ngẫu nhiên sau khi pad
    transforms.RandomHorizontalFlip(),          # flip ngang xác suất 0.5
    transforms.ToTensor(),                      # [0,255] uint8 -> [0,1] float
    transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])

train_set = datasets.CIFAR10("./data", train=True,  download=True, transform=transform_train)
test_set  = datasets.CIFAR10("./data", train=False, download=True, transform=transform_test)

train_loader = DataLoader(train_set, batch_size=128, shuffle=True,  num_workers=2, pin_memory=True)
test_loader  = DataLoader(test_set,  batch_size=256, shuffle=False, num_workers=2, pin_memory=True)

print(len(train_set), len(test_set))   # 50000 10000
print(train_set.classes)               # ['airplane', 'automobile', ...]

Một vài lưu ý:

• Normalize stats: CIFAR_MEAN và CIFAR_STD là giá trị chuẩn từng kênh được tính trên train set. Dùng stats của ImageNet hoặc đoán bằng tay sẽ làm loss khó hội tụ.
• Augmentation chỉ ở train: transform_test không có random crop / flip — đảm bảo đánh giá deterministic.
• num_workers: 2–4 là đủ cho Colab. pin_memory=True cho transfer CPU → GPU nhanh hơn.
• Batch size: 128 train / 256 test phù hợp với T4 (16GB). Giảm còn 64 nếu OOM.

CIFAR-10 chỉ có 5 000 ảnh/class — không nhiều so với độ phức tạp của 10 class ảnh tự nhiên. Một CNN ~1M param dễ memorize gần như toàn bộ train set: train accuracy ~99% nhưng val accuracy mắc kẹt ở ~75%. Augmentation (xem bài 26) tạo các "biến thể" ngẫu nhiên của mỗi ảnh ở từng epoch, ép model học invariant với những biến đổi tự nhiên:

• RandomCrop(32, padding=4): pad 4 pixel mỗi cạnh (ảnh thành \( 40 \times 40 \)) rồi crop ngẫu nhiên một ô \( 32 \times 32 \). Tạo dịch chuyển ±4 pixel — vật thể vẫn cùng class.
• RandomHorizontalFlip(): flip ngang xác suất 0.5. Áp dụng cho hầu hết class CIFAR-10 (mèo, xe, chim... lật ngang vẫn cùng class). Lưu ý: không phù hợp cho ảnh chữ viết.

Hai phép này tuy đơn giản nhưng đẩy val accuracy thêm 5–10 điểm so với khi không có. Các phép mạnh hơn (Cutout, AutoAugment, RandAugment, Mixup) sẽ đẩy tiếp lên 92%+ nhưng phức tạp hơn, để dành cho bài sau.

Mô hình theo phong cách VGG (Simonyan & Zisserman, 2014, arXiv:1409.1556) thu nhỏ: chỉ dùng Conv \( 3 \times 3 \) padding 1, MaxPool \( 2 \times 2 \), và 3 block với số kênh tăng dần.

Backbone:

• Block 1: Conv(3 → 32) → BN → ReLU → Conv(32 → 32) → BN → ReLU → MaxPool(2).
• Block 2: Conv(32 → 64) → BN → ReLU → Conv(64 → 64) → BN → ReLU → MaxPool(2).
• Block 3: Conv(64 → 128) → BN → ReLU → Conv(128 → 128) → BN → ReLU → MaxPool(2).

Classifier head:

• Flatten → Dropout(0.5) → Linear(128·4·4 → 256) → ReLU → Dropout(0.5) → Linear(256 → 10).

Vì sao chọn pattern này:

• Hai Conv liên tiếp trước Pool: tăng receptive field mà vẫn rẻ về param (thay vì 1 Conv \( 5 \times 5 \)).
• BatchNorm sau mỗi Conv (bài 25): ổn định gradient, cho phép dùng LR cao hơn, giảm sự phụ thuộc vào weight initialization.
• Dropout chỉ trong classifier head (bài 23): backbone đã có BN làm regularize nhẹ; Dropout trong Conv ít hiệu quả.

Tra theo công thức bài 28: Conv \( 3 \times 3 \) padding 1 stride 1 giữ nguyên spatial size; MaxPool \( 2 \times 2 \) stride 2 chia đôi.

Backbone output \( (128, 4, 4) \). Flatten ra \( 128 \cdot 4 \cdot 4 = 2048 \) feature — đúng in_features của FC đầu tiên.

Viết model dưới dạng nn.Module cho rõ cấu trúc:

import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            # ----- Block 1 -----
            nn.Conv2d(3,  32, 3, padding=1), nn.BatchNorm2d(32),  nn.ReLU(inplace=True),
            nn.Conv2d(32, 32, 3, padding=1), nn.BatchNorm2d(32),  nn.ReLU(inplace=True),
            nn.MaxPool2d(2),                                 # 32 -> 16

            # ----- Block 2 -----
            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64),  nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64),  nn.ReLU(inplace=True),
            nn.MaxPool2d(2),                                 # 16 -> 8

            # ----- Block 3 -----
            nn.Conv2d(64,  128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True),
            nn.MaxPool2d(2),                                 # 8 -> 4
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(0.5),
            nn.Linear(128 * 4 * 4, 256), nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes),
        )

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

Một vài chi tiết:

• inplace=True ở ReLU: tiết kiệm bộ nhớ. An toàn ở đây vì input của ReLU không cần cho backprop.
• Layer cuối không apply Softmax — sẽ dùng CrossEntropyLoss trực tiếp trên logit (bài 30 mục 13).
• Đặt Dropout trước Linear đầu tiên: drop ngẫu nhiên feature từ backbone trước khi vào FC.

model = SimpleCNN()

def count_params(m):
    return sum(p.numel() for p in m.parameters() if p.requires_grad)

print(f"Total params:  {count_params(model):,}")
print(f"  features  :  {count_params(model.features):,}")
print(f"  classifier:  {count_params(model.classifier):,}")

Kết quả xấp xỉ:

Total params:  812,234
  features  :  291,072
  classifier:  524,810

Có thể in summary chi tiết hơn bằng torchinfo:

from torchinfo import summary
summary(model, input_size=(1, 3, 32, 32))

Số param ~800K — đủ để học CIFAR-10 mà vẫn fit trong vài giây/epoch trên T4.

Trước khi train full, kiểm tra model có khả năng overfit một mini-batch nhỏ. Nếu model không thể đẩy loss của 1 batch cố định về gần 0 trong vài chục step, có bug ở model hoặc loss.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

# Lấy 1 batch và lặp lại 100 step
X, y = next(iter(train_loader))
X, y = X.to(device), y.to(device)

model.train()
for step in range(100):
    optimizer.zero_grad()
    logits = model(X)
    loss = criterion(logits, y)
    loss.backward()
    optimizer.step()
    if step % 10 == 0:
        acc = (logits.argmax(1) == y).float().mean().item()
        print(f"step {step:3d} | loss={loss.item():.4f} | acc={acc:.3f}")

Sau ~100 step, loss phải tiến gần 0 và accuracy trên batch đó ~100%. Nếu loss đứng yên hoặc đi lên: shape sai, label sai, hoặc đang apply softmax trong model (double softmax với CrossEntropyLoss).

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Device:", device)

model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()

# AdamW: Adam với weight decay tách rời (Loshchilov & Hutter 2017, arXiv:1711.05101)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)

# Cosine annealing giảm LR từ 1e-3 -> 0 trong T_max epoch
NUM_EPOCHS = 30
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=NUM_EPOCHS)

Lựa chọn:

• AdamW thay Adam thường: tách weight decay khỏi gradient update, regularize đúng hơn — đặc biệt khi có BN.
• CosineAnnealingLR (bài 21): LR đi theo nửa chu kỳ cosine từ lr về 0 sau T_max epoch. Smoother hơn step decay, không cần tune milestone.
• weight_decay=1e-4: phổ biến cho CNN có BN.

def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()                       # bật BN/Dropout chế độ train
    total_loss, correct, total = 0.0, 0, 0
    for X, y in loader:
        X, y = X.to(device, non_blocking=True), y.to(device, non_blocking=True)

        optimizer.zero_grad()
        logits = model(X)
        loss = criterion(logits, y)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * X.size(0)
        correct    += (logits.argmax(1) == y).sum().item()
        total      += y.size(0)

    return total_loss / total, correct / total

@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()                        # tắt Dropout, dùng running stats của BN
    total_loss, correct, total = 0.0, 0, 0
    for X, y in loader:
        X, y = X.to(device, non_blocking=True), y.to(device, non_blocking=True)
        logits = model(X)
        loss = criterion(logits, y)
        total_loss += loss.item() * X.size(0)
        correct    += (logits.argmax(1) == y).sum().item()
        total      += y.size(0)
    return total_loss / total, correct / total

Vòng huấn luyện chính:

history = {"train_loss": [], "train_acc": [], "val_loss": [], "val_acc": []}

for epoch in range(NUM_EPOCHS):
    train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
    val_loss,   val_acc   = evaluate(model, test_loader,  criterion, device)
    scheduler.step()                    # gọi sau 1 epoch (default)

    history["train_loss"].append(train_loss)
    history["train_acc"].append(train_acc)
    history["val_loss"].append(val_loss)
    history["val_acc"].append(val_acc)

    lr_now = scheduler.get_last_lr()[0]
    print(f"Epoch {epoch+1:02d}/{NUM_EPOCHS} | "
          f"lr={lr_now:.5f} | "
          f"train_loss={train_loss:.4f} acc={train_acc:.4f} | "
          f"val_loss={val_loss:.4f} acc={val_acc:.4f}")

Một số lưu ý:

• model.train() phải gọi đầu mỗi epoch — không phải gọi 1 lần đầu tiên. evaluate() phải gọi model.eval() để BN dùng running stats và Dropout tắt.
• @torch.no_grad(): vô hiệu autograd trong evaluate, tiết kiệm bộ nhớ và nhanh hơn.
• scheduler.step() đặt sau khi train xong 1 epoch (đối với CosineAnnealingLR kiểu epoch-based).

Lưu state_dict mỗi khi val accuracy cải thiện — đây chính là pattern checkpoint của bài 24.

best_acc = 0.0
ckpt_path = "best_cifar10_cnn.pt"

for epoch in range(NUM_EPOCHS):
    train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
    val_loss,   val_acc   = evaluate(model, test_loader,  criterion, device)
    scheduler.step()

    if val_acc > best_acc:
        best_acc = val_acc
        torch.save({
            "epoch": epoch,
            "model_state": model.state_dict(),
            "optimizer_state": optimizer.state_dict(),
            "val_acc": val_acc,
        }, ckpt_path)
        print(f"  saved checkpoint at epoch {epoch+1} (val_acc={val_acc:.4f})")

print(f"Best val_acc: {best_acc:.4f}")

Khi cần load lại để inference hoặc fine-tune:

ckpt = torch.load(ckpt_path, map_location=device)
model.load_state_dict(ckpt["model_state"])
print("Loaded epoch", ckpt["epoch"], "val_acc =", ckpt["val_acc"])

Vì sao chỉ save state_dict thay vì cả model: state_dict portable hơn (không gắn với class definition), nhỏ hơn, ít issue khi refactor code.

Tham khảo log thực tế khi chạy 30 epoch trên Colab T4 (kết quả của bạn có thể lệch 1–2 điểm):

Pattern bình thường:

• Loss giảm đều, accuracy tăng. Train accuracy > val accuracy 2–5 điểm là healthy.
• Nếu val accuracy < 50% sau 5 epoch: kiểm tra lại bug (xem mục 14–15).
• Train acc xa hơn val acc >10 điểm: overfitting — tăng augmentation hoặc thêm regularization.

Khi pipeline không hoạt động đúng, đi qua checklist sau theo thứ tự:

1. In shape qua từng module — forward một tensor dummy và print shape, so với bảng ở mục 6.
2. Overfit 1 batch (mục 9) — loss phải về gần 0 trong <200 step. Nếu không, lỗi ở model/loss.
3. Loss epoch đầu phải giảm — sau 1 epoch full, train_loss phải nhỏ hơn lúc khởi tạo (\( \log 10 \approx 2.30 \) cho 10 class).
4. Train accuracy > baseline random (10%) — sau 1 epoch phải đạt ≥30%. Nếu vẫn ~10%, label hoặc loss có vấn đề.
5. Visualize 1 batch sau augmentation — đảm bảo ảnh trông vẫn nhận biết được class.
6. Check device — X.device, next(model.parameters()).device phải là cùng cuda:0.

Visualize batch sau augmentation:

import matplotlib.pyplot as plt
import numpy as np

X, y = next(iter(train_loader))
mean = np.array(CIFAR_MEAN).reshape(3, 1, 1)
std  = np.array(CIFAR_STD).reshape(3, 1, 1)

fig, axes = plt.subplots(2, 8, figsize=(16, 4))
for i, ax in enumerate(axes.flat):
    img = X[i].cpu().numpy() * std + mean       # unnormalize
    img = np.clip(img.transpose(1, 2, 0), 0, 1) # CHW -> HWC
    ax.imshow(img)
    ax.set_title(train_set.classes[y[i]])
    ax.axis("off")
plt.tight_layout()
plt.show()

• Quên model.train() / model.eval(): BN sẽ dùng batch stats khi đang evaluate, Dropout vẫn drop khi inference → val accuracy lảo đảo theo batch.
• Sai normalize stats: dùng stats của ImageNet (mean 0.485, 0.456, 0.406) cho CIFAR-10 sẽ làm input shift, loss khó hội tụ và accuracy thấp hơn 3–5 điểm.
• Augmentation trên test set: nếu để transform_test = transform_train, val accuracy sẽ biến động giữa các run và thấp hơn thực.
• Quên optimizer.zero_grad(): gradient của step trước cộng dồn — loss tăng vọt sau vài step.
• Sai in_features của FC: thường vì copy code MNIST (\( 7 \times 7 \)) sang CIFAR (\( 4 \times 4 \) với 3 lần pool). Lỗi mat1 and mat2 shapes cannot be multiplied.
• Apply Softmax + CrossEntropyLoss: double softmax — bài 30 mục 13.
• Batch size quá lớn: CUDA out of memory. Giảm batch size hoặc dùng gradient accumulation.
• Quên .to(device): lỗi Expected all tensors to be on the same device. Phải .to(device) cả model và mỗi batch.

Sau khi train xong, dự đoán cho một ảnh PIL bất kỳ:

from PIL import Image
import torch.nn.functional as F

classes = train_set.classes   # 10 class CIFAR-10

# Load checkpoint best
ckpt = torch.load("best_cifar10_cnn.pt", map_location=device)
model.load_state_dict(ckpt["model_state"])
model.eval()

img = Image.open("my_cat.jpg").convert("RGB").resize((32, 32))
x = transform_test(img).unsqueeze(0).to(device)   # (1, 3, 32, 32)

with torch.no_grad():
    logits = model(x)
    probs  = F.softmax(logits, dim=1)[0]
    top3   = probs.topk(3)

for prob, idx in zip(top3.values, top3.indices):
    print(f"{classes[idx]:12s} {prob.item():.3f}")

Một số lưu ý:

• Phải dùng transform_test (có Normalize) — không phải transform_train.
• Resize về \( 32 \times 32 \) trước khi vào model — model train cho size này.
• Apply softmax CHỈ ở đây để hiển thị xác suất, không phải trong forward của model.
• Ảnh ngoài CIFAR-10 (kích thước thật \( 224 \times 224 \), ảnh điện thoại) sẽ accuracy thấp — model được train trên ảnh thumbnail thấp resolution.

Thời gian train tham khảo:

Nếu chỉ có CPU, có thể giảm NUM_EPOCHS = 10 để test pipeline trước rồi mới scale. Trên Colab, đừng quên đổi runtime sang GPU (Runtime → Change runtime type → T4 GPU).

Khi muốn đẩy val accuracy từ ~87% lên cao hơn, thử theo thứ tự (effort tăng dần):

1. Train lâu hơn: 60–100 epoch với scheduler tương ứng, dễ thêm 1–2 điểm.
2. Augmentation mạnh hơn: thêm transforms.RandAugment() hoặc Cutout — đẩy thêm 2–3 điểm.
3. Mixup / CutMix: trộn cặp ảnh + label, regularize mạnh hơn nhiều.
4. Deeper / wider network: thêm 1 block Conv (4 block thay 3), hoặc tăng kênh lên 64–128–256–512. Cẩn thận overfit.
5. Label smoothing: CrossEntropyLoss(label_smoothing=0.1), thường +0.5–1 điểm.
6. Pretrained backbone: bài 32 (transfer learning) — ResNet pretrain trên ImageNet thường đạt 95%+ trên CIFAR-10 mà không cần train từ đầu.

Thứ tự này tuân theo "bottleneck" thực tế: với một CNN nhỏ chưa augmentation đủ, thêm Mixup hữu ích hơn là tăng deeper.

• Tự implement AlexNet, VGG-16, ResNet-18 từ paper — tốt cho việc đọc hiểu kiến trúc. Có thể dùng torchvision.models làm reference.
• Hyperparameter sweep: thử LR \( 5 \times 10^{-4}, 1 \times 10^{-3}, 3 \times 10^{-3} \); batch size 64/128/256; weight decay \( 10^{-5}, 10^{-4}, 10^{-3} \). Ghi lại val accuracy trong bảng.
• Visualize filter của layer Conv đầu tiên: model.features[0].weight có shape (32, 3, 3, 3). Plot 32 filter dạng patch màu để thấy model học edge / color blob.
• Confusion matrix: tính ma trận \( 10 \times 10 \) trên test set để biết class nào hay nhầm với class nào (cat ↔ dog là cặp khó cổ điển).
• Test-time augmentation (TTA): predict trên ảnh gốc + ảnh flip, lấy trung bình logit — thường +0.5–1 điểm.

1. Chạy full pipeline (load + model + train 30 epoch + evaluate). Báo cáo val accuracy tốt nhất.
2. Print tổng số param của model và phân bổ giữa features vs classifier. Giải thích vì sao classifier head có nhiều param hơn dù chỉ 2 Linear.
3. Thay đổi 1 hyperparameter (chọn 1: LR, batch size, dropout rate, weight decay) thành 2 giá trị khác. Train mỗi giá trị, ghi lại bảng val accuracy. So sánh.
4. Visualize 1 batch ảnh đã augmentation (dùng code mục 14). Ảnh có khác nhau giữa 2 epoch không? Tại sao?
5. Thêm label_smoothing=0.1 vào CrossEntropyLoss. Train lại 30 epoch. So sánh val accuracy với version không label smoothing.
6. Tắt augmentation (set transform_train = transform_test). Train lại. Quan sát gap train_acc - val_acc — overfitting xảy ra rõ hơn.
7. Sau khi train xong, load best checkpoint và tính confusion matrix \( 10 \times 10 \) bằng sklearn.metrics.confusion_matrix. Xác định 2 cặp class hay nhầm nhất.

• CIFAR-10 = 60 000 ảnh \( 32 \times 32 \times 3 \), 10 class cân bằng, 50K train / 10K test.
• Load bằng torchvision.datasets.CIFAR10; dùng đúng CIFAR_MEAN / CIFAR_STD khi Normalize.
• Augmentation cơ bản: RandomCrop(32, padding=4) + RandomHorizontalFlip() chỉ ở train.
• Kiến trúc VGG nhỏ: 3 block Conv–BN–ReLU × 2 + MaxPool, classifier Flatten + Dropout + FC. Backbone \( (3, 32, 32) \to (128, 4, 4) \), flatten 2048 → 256 → 10.
• Sanity check: overfit 1 batch trước khi train full.
• Training: AdamW (lr=1e-3, wd=1e-4) + CosineAnnealingLR(T_max=30) + CrossEntropyLoss.
• Loop phải gọi model.train()/model.eval(), optimizer.zero_grad(), @torch.no_grad() ở evaluate.
• Save best checkpoint theo val accuracy; load lại bằng state_dict.
• Expected: ~85–88% sau 30 epoch trên T4 trong ~10–15 phút.
• Muốn lên cao hơn: augmentation mạnh hơn, label smoothing, deeper net, hoặc transfer learning (bài 32).