Training loop

Một bài nhỏ để kiểm toàn bộ quy trình

Ta tự sinh 96 điểm hai đặc trưng; nhãn là 1 nếu $x_1+0.5x_2>0$, còn lại là 0. Quy tắc đã biết trước, nên nếu kết quả tệ thì lỗi nằm ở pipeline chứ không phải ở độ khó của dữ liệu. Chia train/validation trước; mean/std chỉ tính trên train và được dùng lại trên validation.

from pathlib import Path
import json
import subprocess
import sys
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader

torch.set_num_threads(1)
torch.manual_seed(21)
torch.use_deterministic_algorithms(True)
generator = torch.Generator().manual_seed(21)
raw_X = torch.randn(96, 2, generator=generator)
y = (raw_X[:, 0] + .5 * raw_X[:, 1] > 0).long()
order = torch.randperm(96, generator=generator)
train_id, valid_id = order[:72], order[72:]
mean = raw_X[train_id].mean(dim=0)
scale = raw_X[train_id].std(dim=0, unbiased=False).clamp_min(1e-6)
X = (raw_X - mean) / scale
train_data = TensorDataset(X[train_id], y[train_id])
valid_data = TensorDataset(X[valid_id], y[valid_id])
train_loader = DataLoader(train_data, batch_size=16, shuffle=True,
    generator=torch.Generator().manual_seed(21), num_workers=0)
valid_loader = DataLoader(valid_data, batch_size=16, shuffle=False, num_workers=0)
print(len(train_data), len(valid_data), len(train_loader))
print(tuple(next(iter(valid_loader))[0].shape))

Kết quả72 24 5
(16, 2)

Có 72 mẫu train, 24 validation và 5 batch mỗi epoch: bốn batch 16 mẫu, một batch 8 mẫu. Ba từ cần phân biệt:

• Batch: một nhóm mẫu dùng cho một lần tính gradient.
• Step: một lần optimizer cập nhật trọng số; mỗi batch một step.
• Epoch: một lần duyệt hết tập train.

Một epoch có $\lceil n/B \rceil$ step: 10 000 ảnh với batch 64 là 157 step; 1000 mẫu, batch 32, 5 epoch là $32 \cdot 5 = 160$ lần cập nhật.

TensorDataset gom các tensor cùng chiều đầu; dataset thực tế có thể tự đọc tệp ở __getitem__ (mục cuối bài). DataLoader ghép các mẫu thành batch và quyết định thứ tự lấy. Với shuffle=True, thứ tự được xáo lại ở đầu mỗi epoch. Nếu tập train xếp sẵn hết mèo rồi mới tới chó mà không xáo, mỗi đoạn batch chỉ thấy một lớp, và mô hình bị kéo qua lại giữa hai lớp.

Ở Windows, nhiều worker có thể tạo tiến trình con và yêu cầu bảo vệ code bằng if __name__ == '__main__'. Ví dụ đầu tiên dùng num_workers=0 để dễ chạy nối tiếp trong notebook. Việc tăng worker nên dựa trên đo nút nghẽn, không phải mặc định “càng nhiều càng nhanh”.

Loop là một chu trình trạng thái

Mỗi batch đi qua năm bước theo đúng thứ tự:

1. zero_grad xóa gradient cũ.
2. Forward dùng trọng số hiện tại để tạo logits.
3. Loss so logits với nhãn.
4. Backward tính gradient.
5. Step sửa trọng số.

Đảo step lên trước backward thì optimizer dùng gradient cũ hoặc không có gradient. Gọi zero_grad sau backward, trước step, thì xóa chính tín hiệu vừa tính. Nếu cần cắt gradient bằng clip_grad_norm_, đặt nó sau backward và trước step.

Optimizer được tạo từ model.parameters(), tức các trọng số và bias sẽ được chỉnh; nó quyết định chỉnh theo cách nào (SGD, Adam ở bài 3.5).

Vì sao phải xóa? Backward không ghi đè .grad mà cộng gradient mới vào đó. Với hàm $\mathbf{x}^\top\mathbf{w}$, gradient theo $\mathbf{w}$ chính là $\mathbf{x}$, nên dễ nhìn thấy điều này:

w = torch.zeros(2, requires_grad=True)
for xi in [torch.tensor([1., 2.]), torch.tensor([3., -1.])]:
    (xi @ w).backward()
    print(w.grad.tolist())

Kết quả[1.0, 2.0]
[4.0, 1.0]

Lần thứ hai .grad là $[1,2]+[3,-1]=[4,1]$, không phải $[3,-1]$. Thiếu zero_grad, mỗi step dùng tổng gradient của mọi batch đã qua, nên bước cập nhật lớn dần. Cộng dồn có chủ ý lại là một kỹ thuật: gradient accumulation (bài 3.7) gom gradient vài batch nhỏ rồi mới step một lần.

Model có hai lớp linear xen ReLU. Loss nhận logits (B,2) và nhãn long (B,). Ta gom code train một epoch thành một hàm. Loss báo cáo nhân batch size trước khi cộng, để batch cuối nhỏ không được tính ngang trọng lượng với batch đầy:

def make_model():
    return nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(8, 2))

model = make_model()
optimizer = torch.optim.Adam(model.parameters(), lr=.03)
criterion = nn.CrossEntropyLoss()

def train_epoch(model, loader, optimizer):
    model.train()
    loss_sum = 0.
    seen = 0
    for xb, yb in loader:
        optimizer.zero_grad(set_to_none=True)
        logits = model(xb)
        assert logits.shape == (len(yb), 2)
        loss = criterion(logits, yb)
        assert torch.isfinite(loss)
        loss.backward()
        optimizer.step()
        loss_sum += loss.item() * len(yb)
        seen += len(yb)
    return loss_sum / seen

def evaluate(model, loader):
    model.eval()
    loss_sum = 0.
    correct = 0
    seen = 0
    with torch.no_grad():
        for xb, yb in loader:
            logits = model(xb)
            loss_sum += criterion(logits, yb).item() * len(yb)
            correct += (logits.argmax(dim=1) == yb).sum().item()
            seen += len(yb)
    return loss_sum / seen, correct / seen

print('initial validation:', tuple(round(v, 4) for v in evaluate(model, valid_loader)))

Kết quảinitial validation: (0.7394, 0.4583)

Trước khi học, accuracy $0.4583=11/24$, gần mức đoán bừa của hai lớp, và loss $0.7394$ gần $\ln 2\approx0.693$, loss của một model chia đều xác suất cho hai lớp. Đó là điểm xuất phát đúng. Với $K$ lớp, loss ban đầu lệch xa $\ln K$ là dấu hiệu nên kiểm khởi tạo hoặc nhãn trước khi train.

Hai assert trong train_epoch là hai phép kiểm rẻ: logits đúng shape (B,2) và loss hữu hạn. Loss thành nan thường do learning rate quá lớn làm trọng số phân kỳ, hoặc do phép log, phép chia không an toàn trong loss tự viết; assert dừng loop ngay ở batch đầu tiên gặp lỗi, trước khi nan lan vào trọng số.

model.train() và model.eval() điều khiển hành vi của các layer như dropout, BatchNorm (bài 3.6). no_grad() điều khiển việc ghi graph. Eval không tự tắt gradient; no_grad không tự tắt dropout. Mạng nhỏ này chưa có hai layer đó, nhưng vẫn giữ hai lời gọi để loop dùng được khi kiến trúc thay đổi.

Training loss trong hàm là trung bình loss quan sát trong khi trọng số thay đổi qua các batch. Nó không nhất thiết bằng loss đánh giá lại toàn bộ train ở checkpoint cuối epoch. Nếu cần so train/validation chặt chẽ, gọi cùng evaluate trên cả hai loader sau epoch, dùng cùng chế độ layer.

Cộng loss qua các batch cũng phải tính theo mẫu. Giả sử batch đầu có 16 mẫu, loss mean 0.2; batch cuối có 8 mẫu, loss mean 0.8. Trung bình hai số bằng 0.5, nhưng trung bình theo 24 mẫu là $(16\cdot0.2+8\cdot0.8)/24=0.4$. Đó là lý do code nhân len(yb) trước khi cộng.

Cách này đúng với cross-entropy không trọng số của bài. Khi dùng class weight, ignore_index hoặc loss trung bình theo pixel, mẫu số đúng là tổng trọng số hoặc số phần tử thật sự được tính, nên cần đọc reduction của loss trước khi cộng. CrossEntropyLoss mặc định bỏ qua nhãn $-100$ (ignore_index) và chia cho số nhãn còn lại. Đổi reduction='mean' sang 'sum' với batch 64 mà giữ learning rate của SGD là nhân bước đi lên 64 lần.

Chọn checkpoint theo validation

Nếu chỉ lưu epoch cuối, ta có thể bỏ mất model đã tổng quát hóa tốt hơn ở giữa quá trình. Ta chọn loss validation nhỏ nhất; metric lựa chọn trong bài thật có thể là F1, RMSE hoặc metric của scorer. Chốt hướng “nhỏ hơn tốt” hay “lớn hơn tốt” trước khi viết điều kiện so sánh.

Lưu bản tốt nhất phải lưu bản sao. Các tensor trong model.state_dict() dùng chung bộ nhớ với tham số; chỉ gán best_state = model.state_dict() rồi tiếp tục học thì “bản tốt nhất” đổi theo model. Dùng clone hoặc deepcopy, hoặc ghi tệp ngay khi gặp bản tốt hơn.

best_loss = float('inf')
best_epoch = -1
best_state = None
history = []
for epoch in range(35):
    training_loss = train_epoch(model, train_loader, optimizer)
    valid_loss, valid_accuracy = evaluate(model, valid_loader)
    history.append((epoch + 1, training_loss, valid_loss, valid_accuracy))
    if valid_loss < best_loss:
        best_loss = valid_loss
        best_epoch = epoch + 1
        best_state = {name: value.detach().cpu().clone()
                      for name, value in model.state_dict().items()}
    if epoch in [0, 9, 19, 34]:
        print(epoch + 1, round(training_loss, 4),
              round(valid_loss, 4), round(valid_accuracy, 4))

perfect = [e for e, _, _, accuracy in history if accuracy == 1.0]
print('accuracy 1.0:', len(perfect), 'epochs, first', perfect[0])
model.load_state_dict(best_state)
model.eval()
print('best epoch:', best_epoch)
print('best validation:', tuple(round(v, 4) for v in evaluate(model, valid_loader)))

Kết quả1 0.6312 0.6167 0.625
10 0.0687 0.0751 0.9583
20 0.039 0.0396 1.0
35 0.0197 0.0429 1.0
accuracy 1.0: 21 epochs, first 7
best epoch: 32
best validation: (0.03, 1.0)

Mỗi dòng epoch có train loss, validation loss và validation accuracy. Train loss giảm gần như đều từ 0.6312 xuống 0.0197, còn validation loss thấp nhất ở epoch 32 (0.03) rồi lên 0.0429 ở epoch 35: epoch cuối không phải epoch tốt nhất.

Validation chỉ 24 mẫu nên accuracy đi theo nấc $1/24$; 21 trong 35 epoch cùng đạt 1.0, accuracy không phân biệt được chúng, còn loss thì phân biệt được. Điểm (0.03, 1.0) đã được dùng để chọn epoch, nên nó hơi lạc quan so với dữ liệu mới.

Lưu đủ để suy luận lại

Đưa đầu vào mới qua model đã học để lấy dự đoán gọi là suy luận (inference), khác với huấn luyện là lúc trọng số còn được chỉnh. Để suy luận lại ở một chương trình khác, ta lưu một gói (packet) gồm state_dict, mean/std và vài đầu vào kiểm tra kèm logits mong đợi, vào thư mục receipts/phase-3 cạnh notebook. Muốn tiếp tục training chính xác thì còn cần trạng thái optimizer, scheduler, số epoch và trạng thái bộ sinh số ngẫu nhiên.

torch.load một tệp lạ có thể chạy code tùy ý, nên chỉ tải checkpoint đáng tin cậy; weights_only=True chỉ cho phép tensor và kiểu dữ liệu cơ bản, đủ cho packet này.

artifact_dir = Path('receipts/phase-3')
artifact_dir.mkdir(parents=True, exist_ok=True)
checkpoint_path = artifact_dir / '3-4-best.pt'
query = torch.tensor([[1., 0.], [-1., 0.]])
with torch.no_grad():
    expected = model((query - mean) / scale)
torch.save({'state': best_state, 'mean': mean, 'scale': scale,
            'query': query, 'expected': expected, 'epoch': best_epoch}, checkpoint_path)
fresh = make_model()
packet = torch.load(checkpoint_path, map_location='cpu', weights_only=True)
fresh.load_state_dict(packet['state'])
fresh.eval()
with torch.no_grad():
    reloaded = fresh((packet['query'] - packet['mean']) / packet['scale'])
print('same logits:', torch.allclose(reloaded, expected, atol=1e-7))
print('classes:', reloaded.argmax(dim=1).tolist())

Kết quảsame logits: True
classes: [1, 0]

True kiểm khôi phục trong cùng tiến trình; hai lớp [1, 0] khớp quy tắc ban đầu: $(1,0)$ cho $x_1+0.5x_2=1>0$, $(-1,0)$ cho $-1<0$.

Phép kiểm trong cùng tiến trình chưa cho biết notebook có dựa vào biến cũ hay không, nên ta tạo một script nhỏ chạy bằng Python mới. Script tự tạo kiến trúc, đọc checkpoint và chỉ dùng thông tin trong packet:

inference_source = '''import sys
import torch
from torch import nn
torch.set_num_threads(1)
packet = torch.load(sys.argv[1], map_location="cpu", weights_only=True)
model = nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(8, 2))
model.load_state_dict(packet["state"])
model.eval()
with torch.no_grad():
    x = (packet["query"] - packet["mean"]) / packet["scale"]
    logits = model(x)
print("new process match:", torch.allclose(logits, packet["expected"], atol=1e-7))
print("new process classes:", logits.argmax(dim=1).tolist())
'''
inference_path = artifact_dir / '3-4-infer.py'
inference_path.write_text(inference_source, encoding='utf-8')
result = subprocess.check_output([sys.executable, '-B', '-X', 'utf8',
    str(inference_path), str(checkpoint_path)], text=True)
print(result.rstrip())

Kết quảnew process match: True
new process classes: [1, 0]

Phép kiểm này giống chạy lời giải CP trên input mới qua stdin: kết quả không được phụ thuộc vào biến còn sót trong môi trường thử. Ở bài thi thật, script suy luận còn phải theo đúng quy định nộp bài của đề (format file, thứ tự ID, tệp đề cho), vì vậy phép kiểm ở phiên mới phải chạy đúng script sẽ nộp.

Tái lập

Seed của khởi tạo model, chia dữ liệu và DataLoader cần được cố định. Bài đặt seed, dùng CPU một thread và chế độ deterministic, nên chạy lại trong cùng môi trường cho cùng số. PyTorch không hứa kết quả giống từng bit giữa mọi phiên bản, nền tảng và GPU; vì vậy ghi lại phiên bản thư viện, cách chia và cấu hình cùng điểm số.

Tái lập được một lần chạy chưa có nghĩa là một lần chạy đủ để so hai cấu hình. Ngay trong log trên, validation loss đi từ 0.0396 (epoch 20) xuống 0.03 (epoch 32) rồi lên 0.0429 (epoch 35). Chênh lệch giữa hai cấu hình chỉ đáng tin khi lớn hơn rõ mức dao động giữa các seed.

Tự viết Dataset và collate_fn

TensorDataset cần cả bộ dữ liệu nằm sẵn trong một tensor. Thư mục hàng nghìn ảnh, mỗi ảnh một cỡ, thì nên viết một lớp con của Dataset với hai hàm ở bài 0.1: __len__ trả số mẫu, __getitem__(i) đọc một mẫu từ đĩa khi cần.

import numpy as np
import pandas as pd
from PIL import Image
from tempfile import TemporaryDirectory
from torch.utils.data import Dataset

class FolderDataset(Dataset):
    def __init__(self, root, table, size=8):
        self.root, self.table, self.size = Path(root), table.reset_index(drop=True), size

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

    def __getitem__(self, i):
        row = self.table.iloc[i]
        image = Image.open(self.root / row["file"]).convert("L").resize((self.size, self.size))
        x = torch.from_numpy(np.asarray(image, dtype=np.float32) / 255).unsqueeze(0)   # (1, H, W)
        return x, int(row["label"])

with TemporaryDirectory() as directory:
    root = Path(directory)
    rng = np.random.default_rng(0)
    rows = []
    for i in range(10):
        w, h = rng.integers(10, 30, size=2)                    # mỗi ảnh một kích thước
        label = i % 2                                          # nhãn 1: ảnh sáng hơn
        pixels = rng.integers(0, 120, size=(h, w)) + 120 * label
        Image.fromarray(pixels.astype(np.uint8)).save(root / f"{i}.png")
        rows.append({"file": f"{i}.png", "label": label})
    data = FolderDataset(root, pd.DataFrame(rows))
    loader = DataLoader(data, batch_size=4, shuffle=False, num_workers=0)
    xb, yb = next(iter(loader))
    print(len(data), len(loader), tuple(data[0][0].shape), tuple(xb.shape), yb.tolist())
    print(round(float(xb[yb == 0].mean()), 3), round(float(xb[yb == 1].mean()), 3))

Kết quả10 3 (1, 8, 8) (4, 1, 8, 8) [0, 1, 0, 1]
0.236 0.709

__getitem__ đưa mọi ảnh về $(1, 8, 8)$, nên DataLoader xếp bốn mẫu thành batch $(4, 1, 8, 8)$; 10 mẫu, batch 4 cho 3 batch. Độ sáng 0.236 và 0.709 khớp cách sinh: ảnh nhãn 1 sáng hơn 120 mức. Đổi cỡ, chuẩn hóa, augmentation (chương 5) đều đặt trong __getitem__.

Văn bản và âm thanh mỗi mẫu một độ dài. DataLoader mặc định ghép bằng torch.stack nên báo lỗi; collate_fn tự viết đệm chuỗi ngắn bằng 0 và giữ độ dài thật.

from torch.nn.utils.rnn import pad_sequence

pairs = [(torch.tensor([5, 2, 9]), 1), (torch.tensor([7]), 0), (torch.tensor([3, 3, 1, 4]), 1)]
try:
    next(iter(DataLoader(pairs, batch_size=3)))
except RuntimeError as error:
    print("không có collate_fn:", str(error).split(" at entry")[0])

def pad_collate(batch):
    seqs, labels = zip(*batch)
    lengths = torch.tensor([len(s) for s in seqs])
    padded = pad_sequence(seqs, batch_first=True, padding_value=0)
    return padded, lengths, torch.tensor(labels)

for tokens, lengths, labels in DataLoader(pairs, batch_size=3, collate_fn=pad_collate):
    print(tokens)
    print(lengths.tolist(), labels.tolist())

Kết quảkhông có collate_fn: stack expects each tensor to be equal size, but got [3]
tensor([[5, 2, 9, 0],
        [7, 0, 0, 0],
        [3, 3, 1, 4]])
[3, 1, 4] [1, 0, 1]

collate_fn nhận list các mẫu do __getitem__ trả về và trả về một batch. Số 0 là token đệm; tính loss trên cả ma trận mà không dùng lengths hay mặt nạ thì model học cả phần đệm. Bài 3.8 và chương 6 dùng khuôn này kèm mask.

collate_fn cũng là chỗ gọn để bỏ mẫu hỏng: nếu __getitem__ trả None cho ảnh không đọc được, lọc các None trước khi gom batch; tốt hơn nữa là lọc danh sách tệp hỏng trước khi huấn luyện.