Việc xây dựng một mô hình học sâu đơn giản như mạng perceptron đa tầng (MLP) trên tập dữ liệu MNIST là bước nền tảng giúp người mới làm quen với luồng xử lý trong PyTorch: từ tải dữ liệu, thiết kế kiến trúc, huấn luyện đến đánh giá. Dưới đây là triển khai tối ưu hóa, rõ ràng và dễ mở rộng — không phụ thuộc vào các hướng dẫn ngoài, chỉ sử dụng thư viện chuẩn.

Tải và tiền xử lý dữ liệu MNIST

Chúng ta sử dụng torchvision.datasets.MNIST với chuẩn hóa dựa trên thống kê toàn bộ tập huấn luyện (trung bình 0.1307, độ lệch chuẩn 0.3081), đảm bảo đầu vào ổn định cho mạng:

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

BATCH = 128
EPOCHS = 8
LR = 1e-3

preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_set = datasets.MNIST("./data", train=True, download=True, transform=preprocess)
test_set = datasets.MNIST("./data", train=False, transform=preprocess)

train_loader = DataLoader(train_set, batch_size=BATCH, shuffle=True)
test_loader = DataLoader(test_set, batch_size=BATCH, shuffle=False)

Thiết kế kiến trúc MLP cải tiến

Mô hình được viết lại với cấu trúc rõ ràng hơn: sử dụng nn.Sequential để tăng tính minh bạch, thêm dropout nhằm giảm overfitting, và áp dụng hàm kích hoạt LeakyReLU thay vì ReLU truyền thống để tránh vấn đề "neuron chết":

class DigitClassifier(nn.Module):
    def __init__(self, input_dim=784, hidden_dims=[512, 256], num_classes=10, dropout_rate=0.2):
        super().__init__()
        layers = []
        prev_dim = input_dim
        
        for dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, dim),
                nn.LeakyReLU(negative_slope=0.1),
                nn.Dropout(dropout_rate)
            ])
            prev_dim = dim
            
        layers.append(nn.Linear(prev_dim, num_classes))
        self.network = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.network(x.flatten(1))

model = DigitClassifier()

Huấn luyện có kiểm soát

Quá trình huấn luyện tích hợp báo cáo chi tiết theo epoch, đồng thời lưu lại tổn thất trung bình mỗi epoch để theo dõi xu hướng hội tụ:

criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-4)

for epoch in range(EPOCHS):
    model.train()
    total_loss = 0.0
    
    for imgs, labels in train_loader:
        optimizer.zero_grad()
        preds = model(imgs)
        loss = criterion(preds, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    
    avg_loss = total_loss / len(train_loader)
    print(f"[Epoch {epoch+1}/{EPOCHS}] Avg Loss: {avg_loss:.4f}")

Đánh giá độ chính xác trên tập kiểm thử

Sau huấn luyện, mô hình được chuyển sang chế độ eval(), tắt các lớp như Dropout, và tính toán độ chính xác từng mẫu một cách hiệu quả:

model.eval()
correct, total = 0, 0

with torch.no_grad():
    for imgs, labels in test_loader:
        outputs = model(imgs)
        _, preds = outputs.max(1)
        correct += preds.eq(labels).sum().item()
        total += labels.size(0)

accuracy = 100. * correct / total
print(f"Final Test Accuracy: {accuracy:.2f}%")

Tùy chỉnh dữ liệu từ thư mục cục bộ

Nếu muốn huấn luyện trên ảnh cá nhân (ví dụ: ảnh chữ số chụp bằng điện thoại), bạn có thể tạo dataset tùy chỉnh với logic phân loại rõ ràng theo tên tệp — ví dụ "digit_3_001.png" → nhãn 3:

import os
from pathlib import Path
from PIL import Image

class LocalDigitDataset(torch.utils.data.Dataset):
    def __init__(self, data_dir, transform=None):
        self.data_dir = Path(data_dir)
        self.transform = transform
        self.samples = []
        
        for img_path in self.data_dir.glob("*.png"):
            try:
                label = int(img_path.stem.split('_')[1])
                self.samples.append((img_path, label))
            except (IndexError, ValueError):
                continue
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        img_path, label = self.samples[idx]
        img = Image.open(img_path).convert("L")
        if self.transform:
            img = self.transform(img)
        return img, label

# Sử dụng tương tự như MNIST
custom_transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

custom_train = LocalDigitDataset("./custom_digits/train", custom_transform)
custom_loader = DataLoader(custom_train, batch_size=32, shuffle=True)