PyTorch: 7 bước huấn luyện một mạng phân loại ảnh

Lý thuyết về neuron, loss, gradient descent và backprop đã đủ. Bài này chuyển sang thực hành: dùng PyTorch, framework phổ biến nhất trong nghiên cứu và ngày càng nhiều trong công nghiệp, để xây dựng trọn vẹn một pipeline huấn luyện phân loại ảnh CIFAR-10. Khung 7 bước ở đây dùng lại được cho CNN, transfer learning, detection…

Số trang (ví dụ "trang 81") trong bài chỉ tới slide của bài giảng Deep Learning – Computer Vision – NLP (Thang Nguyen, Viet Nguyen); bài này giải thích chi tiết, kiểm chứng và đính chính các slide đó.

Yêu cầu: biết Python cơ bản (xem chủ đề Python và chặng 1 của lộ trình), cài pip install torch torchvision (chọn bản có CUDA tại pytorch.org nếu máy có GPU NVIDIA), hoặc dùng Google Colab.

Bản đồ 7 bước

Toàn bộ nằm trên nền torch.autograd: dữ liệu đi qua Data → Dataset → DataLoader, vào Model → Loss, rồi Optimizer cập nhật trọng số (trang 76).

Tensor: "mảng NumPy chạy được trên GPU"

Tensor trong PyTorch tương đương ndarray của NumPy, cộng thêm hai siêu năng lực: chạy trên GPU và tự tính gradient (autograd).

import torch

tensor = torch.tensor([[[1, 2, 3],
                        [3, 6, 9],
                        [2, 4, 5]]])
print(tensor.shape)    # torch.Size([1, 3, 3])  ← dim 0 có 1 phần tử, dim 1 và dim 2 mỗi chiều 3
print(tensor.ndim)     # 3
print(tensor.dtype)    # torch.int64
print(tensor.device)   # cpu

# Các thao tác hay dùng
x = torch.rand(2, 3, 32, 32)            # batch 2 ảnh màu 32x32
print(x.view(2, -1).shape)              # (2, 3072): duỗi phẳng, -1 = tự suy ra
print(x.permute(0, 2, 3, 1).shape)      # (2, 32, 32, 3): đổi thứ tự trục NCHW -> NHWC
print(x[0].unsqueeze(0).shape)          # (1, 3, 32, 32): thêm trục batch
device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.to(device)                         # chuyển sang GPU (nếu có)

Cách đọc "dim": dim=0 là cặp ngoặc ngoài cùng, dim=1 là cấp ngoặc tiếp theo… Phần tử cuối của shape luôn ứng với cặp ngoặc trong cùng.

Từ ảnh đến tensor

Hai cách đọc ảnh ở trang 67 cho ra tensor khác nhau, rất đáng để ý:

import cv2
import torch
from PIL import Image
from torchvision.transforms import ToTensor

# Cách 1: Pillow + ToTensor (khuyên dùng với torchvision)
img = ToTensor()(Image.open("cat.jpg").convert("RGB"))     # (3, H, W), float [0, 1]

# Cách 2: OpenCV -> phải tự đổi màu, đổi trục và chuẩn hoá
bgr = cv2.imread("cat.jpg")                                 # (H, W, 3) uint8, BGR
rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
img2 = torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0  # (3, H, W) float [0, 1]
print(torch.allclose(img, img2, atol=1/255))                # gần như bằng nhau

Tiền xử lý dữ liệu: zero-center và normalize

Như minh hoạ ở trang 62: dữ liệu gốc lệch tâm và trải dài không đều; trừ trung bình (zero-centered) đưa nó về quanh gốc toạ độ, chia độ lệch chuẩn (normalized) đưa các chiều về cùng thang đo. Với ảnh, ta chuẩn hoá từng kênh màu:

Dữ liệu chuẩn hoá giúp mặt loss "tròn" hơn, gradient descent hội tụ nhanh và ổn định hơn. Hai bộ giá trị phổ biến:

• CIFAR-10: mean ≈ (0.4914, 0.4822, 0.4465), std ≈ (0.2470, 0.2435, 0.2616).
• ImageNet (dùng khi fine-tune mô hình pretrained): mean (0.485, 0.456, 0.406), std (0.229, 0.224, 0.225).

from torchvision import datasets, transforms

raw = datasets.CIFAR10("./data", train=True, download=True, transform=transforms.ToTensor())
x = torch.stack([img for img, _ in raw])        # (50000, 3, 32, 32), tốn ~600MB RAM
print(x.mean(dim=(0, 2, 3)), x.std(dim=(0, 2, 3)))   # tính theo từng kênh

Bước 1: Dataset và DataLoader

Cách 1: dataset có sẵn trong torchvision

torchvision.datasets có sẵn MNIST, CIFAR-10/100, ImageNet, COCO, VOC… Mỗi phần tử trả về một cặp (ảnh, nhãn):

from torchvision.datasets import CIFAR10

train_dataset = CIFAR10(root="./cifar", train=True, download=True)
test_dataset = CIFAR10(root="./cifar", train=False, download=True)
image, label = train_dataset[1009]       # tương đương train_dataset.__getitem__(1009)
print(image.size, label)                  # (32, 32) 2  ← ảnh PIL 32x32, nhãn 2 = 'bird'
print(train_dataset.classes)              # ['airplane', 'automobile', 'bird', ...]
print(train_dataset.class_to_idx)         # {'airplane': 0, 'automobile': 1, ...}

Cách 2: ImageFolder, khi ảnh đã chia theo thư mục

Nếu dữ liệu được sắp xếp mỗi lớp một thư mục, ImageFolder tự gán nhãn theo tên thư mục (sắp xếp theo thứ tự chữ cái):

butterfly/
├── train/
│   ├── ADONIS/        1.jpg, 2.jpg, ...
│   └── ATLAS MOTH/    1.jpg, 2.jpg, ...
├── valid/ ...
└── test/  ...

from torchvision.datasets import ImageFolder

train_dataset = ImageFolder(root="butterfly/train")
val_dataset = ImageFolder(root="butterfly/valid")
image, label = train_dataset[1009]
print(image.size, label, train_dataset.classes[label])

Cách 3: tự viết Dataset

Dữ liệu thực tế hiếm khi đúng khuôn. Một custom dataset chỉ cần (trang 71): kế thừa torch.utils.data.Dataset, cài đặt __len__ (số mẫu) và __getitem__ (trả về mẫu thứ idx). Ví dụ đọc ảnh lỗi sản phẩm với nhãn trong file CSV:

import os
import pandas as pd
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms


class DefectDataset(Dataset):
    """Đọc ảnh trong thư mục images/ và nhãn từ file CSV (cột filename, label)."""

    def __init__(self, root, csv_file, transform=None):
        self.root = root
        self.df = pd.read_csv(os.path.join(root, csv_file))
        self.classes = sorted(self.df["label"].unique())
        self.class_to_idx = {c: i for i, c in enumerate(self.classes)}
        self.transform = transform

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        image = Image.open(os.path.join(self.root, "images", row["filename"])).convert("RGB")
        label = self.class_to_idx[row["label"]]
        if self.transform:
            image = self.transform(image)
        return image, label


tf = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
ds = DefectDataset("data/defects", "labels.csv", transform=tf)
loader = DataLoader(ds, batch_size=4, shuffle=True)
images, labels = next(iter(loader))
print(images.shape)    # torch.Size([4, 3, 224, 224])

DataLoader: gom batch, xáo trộn, đọc song song

from torch.utils.data import DataLoader

train_loader = DataLoader(train_set, batch_size=16, shuffle=True,
                          num_workers=8, drop_last=True, pin_memory=True)
test_loader = DataLoader(test_set, batch_size=4, shuffle=False, num_workers=8)

for images, labels in train_loader:
    print(images.shape, labels.shape)    # torch.Size([16, 3, 32, 32]) torch.Size([16])
    break

• shuffle=True cho tập train (mỗi epoch thứ tự khác nhau), False cho val/test.
• num_workers: số tiến trình đọc dữ liệu song song. Trên Windows, code phải nằm trong khối if __name__ == "__main__":, nếu không sẽ lỗi.
• drop_last=True: bỏ batch cuối bị thiếu (hữu ích khi có BatchNorm).
• pin_memory=True: tăng tốc chép dữ liệu sang GPU.

Bước 2: Định nghĩa mô hình

Mọi mô hình kế thừa nn.Module: khai báo các layer trong __init__, mô tả luồng dữ liệu trong forward. Hai cách viết ở trang 73 là tương đương; cách dùng nn.Sequential gọn hơn:

import torch
import torch.nn as nn


class SimpleNeuralNetwork(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.flatten = nn.Flatten()
        self.layers = nn.Sequential(
            nn.Linear(3 * 32 * 32, 512), nn.ReLU(),
            nn.Linear(512, 1024), nn.ReLU(),
            nn.Linear(1024, 512), nn.ReLU(),
            nn.Linear(512, num_classes),        # đầu ra là logits, KHÔNG thêm softmax
        )

    def forward(self, x):
        return self.layers(self.flatten(x))


model = SimpleNeuralNetwork()
random_image = torch.rand(1, 3, 32, 32)          # [B, C, H, W]
print(model(random_image).shape)                  # torch.Size([1, 10])
print(sum(p.numel() for p in model.parameters()))  # 2628618 tham số

Gần 2,63 triệu tham số, phần lớn nằm ở lớp đầu (3072×512). Con số này sẽ là điểm so sánh khi ta chuyển sang CNN.

Bước 3 và 4: Loss function và Optimizer

import torch.nn as nn
from torch.optim import SGD

criterion = nn.CrossEntropyLoss()
optimizer = SGD(model.parameters(), lr=0.001, momentum=0.9)

Lý do chọn và các lựa chọn khác đã phân tích ở bài Loss function và Optimizer.

Bước 5, 6, 7: Train, Validate, Test (script hoàn chỉnh)

Ba tập dữ liệu có vai trò khác nhau, nhầm lẫn là lỗi phương pháp nghiêm trọng:

• Train: dùng để cập nhật trọng số.
• Validation: theo dõi trong lúc huấn luyện, chọn siêu tham số, chọn checkpoint tốt nhất, phát hiện overfitting. CIFAR-10 không có sẵn tập validation, nên ta tách 5.000 ảnh từ tập train.
• Test: chỉ chạy một lần ở cuối để báo cáo kết quả khách quan.

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

MEAN, STD = (0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)


class SimpleNeuralNetwork(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.flatten = nn.Flatten()
        self.layers = nn.Sequential(
            nn.Linear(3 * 32 * 32, 512), nn.ReLU(),
            nn.Linear(512, 1024), nn.ReLU(),
            nn.Linear(1024, 512), nn.ReLU(),
            nn.Linear(512, num_classes),
        )

    def forward(self, x):
        return self.layers(self.flatten(x))


def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()                                   # bật chế độ huấn luyện (Dropout, BatchNorm)
    total_loss, correct, n = 0.0, 0, 0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)                     # forward
        loss = criterion(outputs, labels)
        optimizer.zero_grad()                       # xoá gradient cũ
        loss.backward()                             # backward
        optimizer.step()                            # cập nhật trọng số
        total_loss += loss.item() * images.size(0)
        correct += (outputs.argmax(1) == labels).sum().item()
        n += images.size(0)
    return total_loss / n, correct / n


@torch.no_grad()                                     # không cần tính gradient khi đánh giá
def evaluate(model, loader, criterion, device):
    model.eval()                                    # tắt Dropout, BatchNorm dùng thống kê đã lưu
    total_loss, correct, n = 0.0, 0, 0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        total_loss += criterion(outputs, labels).item() * images.size(0)
        correct += (outputs.argmax(1) == labels).sum().item()
        n += images.size(0)
    return total_loss / n, correct / n


def main():
    torch.manual_seed(42)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # ---- Bước 1: dataset + dataloader ----
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(MEAN, STD)])
    full_train = datasets.CIFAR10("./data", train=True, download=True, transform=transform)
    test_set = datasets.CIFAR10("./data", train=False, download=True, transform=transform)
    train_set, val_set = random_split(full_train, [45000, 5000],
                                      generator=torch.Generator().manual_seed(42))
    train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=2, drop_last=True)
    val_loader = DataLoader(val_set, batch_size=256, shuffle=False, num_workers=2)
    test_loader = DataLoader(test_set, batch_size=256, shuffle=False, num_workers=2)

    # ---- Bước 2, 3, 4: model, loss, optimizer ----
    model = SimpleNeuralNetwork(num_classes=10).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

    # ---- Bước 5, 6: train + validate ----
    best_acc = 0.0
    for epoch in range(20):
        tr_loss, tr_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
        va_loss, va_acc = evaluate(model, val_loader, criterion, device)
        print(f"epoch {epoch + 1:2d} | train loss {tr_loss:.3f} acc {tr_acc:.3f} "
              f"| val loss {va_loss:.3f} acc {va_acc:.3f}")
        if va_acc > best_acc:                       # lưu checkpoint tốt nhất theo validation
            best_acc = va_acc
            torch.save(model.state_dict(), "best_model.pt")

    # ---- Bước 7: test với checkpoint tốt nhất ----
    model.load_state_dict(torch.load("best_model.pt", map_location=device))
    te_loss, te_acc = evaluate(model, test_loader, criterion, device)
    print(f"TEST loss {te_loss:.3f} acc {te_acc:.3f}")


if __name__ == "__main__":      # bắt buộc trên Windows khi num_workers > 0
    main()

Kỳ vọng: loss epoch đầu bắt đầu quanh 2.30 (\(\ln 10\)), val accuracy của mạng MLP này dừng ở khoảng 50–55%, trong khi train accuracy tiếp tục tăng: dấu hiệu overfitting. MLP coi ảnh như một dãy số rời rạc, bỏ qua cấu trúc không gian. Một CNN nhỏ với số tham số ít hơn có thể vượt 80%: đó là lý do ra đời của CNN.

Hai chế độ train() và eval()

• model.train(): Dropout bật, BatchNorm dùng thống kê của batch hiện tại và cập nhật trung bình chạy.
• model.eval(): Dropout tắt, BatchNorm dùng trung bình chạy đã lưu. Quên gọi eval() khi đánh giá là lỗi rất hay gặp, cho kết quả dao động khó hiểu.
• torch.no_grad(): không xây đồ thị autograd, tiết kiệm bộ nhớ và tăng tốc khi đánh giá/suy luận.

Dự đoán một ảnh mới (inference)

from PIL import Image

classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
tf = transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize(MEAN, STD)])

model.eval()
x = tf(Image.open("my_cat.jpg").convert("RGB")).unsqueeze(0).to(device)   # thêm trục batch
with torch.no_grad():
    probs = torch.softmax(model(x), dim=1)[0]
top3 = probs.topk(3)
for p, i in zip(top3.values, top3.indices):
    print(f"{classes[i]}: {p:.1%}")

Lưu ý ảnh khi suy luận phải qua đúng phép tiền xử lý như lúc huấn luyện (cùng kích thước, cùng mean/std).

Mẹo gỡ lỗi khi huấn luyện

• Overfit một batch: lấy 1 batch nhỏ, huấn luyện lặp lại; loss phải về gần 0. Nếu không, pipeline có lỗi.
• Kiểm tra dữ liệu bằng mắt: vẽ vài ảnh sau transform kèm nhãn trước khi huấn luyện.
• Cố định seed (torch.manual_seed) để so sánh thí nghiệm công bằng.
• Ghi log loss/accuracy mỗi epoch (TensorBoard, Weights & Biases) và vẽ đồ thị train vs val.

Tóm tắt

• Tensor ảnh trong PyTorch: C×H×W, batch N×C×H×W, float đã chuẩn hoá.
• Dataset trả về một mẫu; DataLoader gom batch, xáo trộn, đọc song song.
• Model kế thừa nn.Module; đầu ra là logits.
• Vòng lặp: train() → forward → loss → zero_grad → backward → step; đánh giá với eval() + no_grad().
• Train để học, validation để chọn, test để báo cáo một lần.

Đọc tiếp: Convolutional Neural Network: phép tích chập và pooling.

Tài liệu tham khảo

• PyTorch Tutorials: Learn the Basics, Datasets & DataLoaders, Training a Classifier (pytorch.org/tutorials).
• Tài liệu torchvision: Datasets và Transforms.
• Andrej Karpathy (2019). A Recipe for Training Neural Networks (karpathy.github.io).