PyTorch: lưu checkpoint và huấn luyện tiếp sau khi bị gián đoạn

Huấn luyện một mô hình ảnh công nghiệp có thể mất vài giờ đến vài ngày. Mất điện, máy tự khởi động lại để cập nhật, hết phiên Colab... và bạn phải chạy lại từ đầu, nếu không lưu checkpoint đúng cách. Bài này hướng dẫn lưu và nạp checkpoint trong PyTorch sao cho huấn luyện tiếp cho ra đúng kết quả như chưa từng bị gián đoạn, cùng những lỗi phổ biến khi dùng torch.load từ PyTorch 2.6.

Nếu chưa quen state_dict và vòng lặp huấn luyện, hãy đọc PyTorch cơ bản: tensor, autograd và vòng lặp huấn luyện. Code trong bài chạy trên CPU với PyTorch 2.x.

Checkpoint cần chứa những gì?

Một checkpoint chỉ dùng để suy luận thì chỉ cần trọng số. Còn checkpoint để huấn luyện tiếp phải có đủ cả bốn.

Chuẩn bị: một bài toán nhỏ, có thể lặp lại

Để kiểm chứng "huấn luyện tiếp cho kết quả y hệt", ta dùng dữ liệu cố định, khởi tạo trọng số cố định và thứ tự batch cố định:

import torch
import torch.nn as nn

torch.manual_seed(0)
X = torch.randn(256, 8)
y = (X[:, 0] + X[:, 1] > 0).long()            # bài toán OK/NG giả lập


def tao_mo_hinh():
    torch.manual_seed(1)                       # khởi tạo trọng số giống nhau mỗi lần
    mh = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 2))
    toi_uu = torch.optim.SGD(mh.parameters(), lr=0.1, momentum=0.9)
    lich = torch.optim.lr_scheduler.StepLR(toi_uu, step_size=3, gamma=0.5)
    return mh, toi_uu, lich


def mot_epoch(mh, toi_uu, lich):
    mh.train()
    for i in range(0, len(X), 64):             # 4 batch, thứ tự cố định
        toi_uu.zero_grad()
        loss = nn.functional.cross_entropy(mh(X[i:i + 64]), y[i:i + 64])
        loss.backward()
        toi_uu.step()
    lich.step()
    return loss.item()

StepLR(step_size=3, gamma=0.5) giảm learning rate một nửa sau mỗi 3 epoch: 0,1 → 0,05 → 0,025.

Hàm lưu và nạp

def luu_checkpoint(duong_dan, mh, toi_uu, lich, epoch, loss_tot_nhat):
    torch.save({
        "epoch": epoch,
        "model": mh.state_dict(),
        "optimizer": toi_uu.state_dict(),
        "scheduler": lich.state_dict(),
        "loss_tot_nhat": loss_tot_nhat,
    }, duong_dan)


def nap_checkpoint(duong_dan, mh, toi_uu, lich):
    ck = torch.load(duong_dan, map_location="cpu")    # weights_only=True là mặc định từ PyTorch 2.6
    mh.load_state_dict(ck["model"])
    toi_uu.load_state_dict(ck["optimizer"])
    lich.load_state_dict(ck["scheduler"])
    return ck["epoch"] + 1, ck["loss_tot_nhat"]       # epoch bắt đầu chạy tiếp

map_location="cpu" cho phép nạp checkpoint lưu trên máy có GPU vào máy chỉ có CPU (ví dụ máy tính công nghiệp tại trạm). Nếu huấn luyện tiếp trên GPU, hãy mh.to("cuda") trước khi tạo optimizer, rồi mới nạp state.

Kiểm chứng: gián đoạn rồi chạy tiếp có giống chạy liền?

# Chạy A: huấn luyện liền 6 epoch
mh, toi_uu, lich = tao_mo_hinh()
for epoch in range(6):
    loss = mot_epoch(mh, toi_uu, lich)
print(f"A (liền mạch)   : loss cuối = {loss:.6f}, lr = {lich.get_last_lr()[0]}")

# Chạy B: 3 epoch, lưu checkpoint, "mất điện"
mh, toi_uu, lich = tao_mo_hinh()
for epoch in range(3):
    loss = mot_epoch(mh, toi_uu, lich)
luu_checkpoint("last.pt", mh, toi_uu, lich, epoch, loss)
del mh, toi_uu, lich

# ... khởi động lại máy, tạo đối tượng mới và nạp checkpoint
mh, toi_uu, lich = tao_mo_hinh()
bat_dau, _ = nap_checkpoint("last.pt", mh, toi_uu, lich)
for epoch in range(bat_dau, 6):
    loss = mot_epoch(mh, toi_uu, lich)
print(f"B (có gián đoạn): loss cuối = {loss:.6f}, lr = {lich.get_last_lr()[0]}")

A (liền mạch)   : loss cuối = 0.192472, lr = 0.025
B (có gián đoạn): loss cuối = 0.192472, lr = 0.025

Hai kết quả trùng nhau đến chữ số thứ sáu. Checkpoint đã lưu đầy đủ trạng thái.

Điều gì xảy ra nếu chỉ lưu trọng số?

# Chạy C: SAI, chỉ lưu trọng số mô hình, quên optimizer và scheduler
mh, toi_uu, lich = tao_mo_hinh()
for epoch in range(3):
    mot_epoch(mh, toi_uu, lich)
torch.save(mh.state_dict(), "chi_trong_so.pt")

mh, toi_uu, lich = tao_mo_hinh()
mh.load_state_dict(torch.load("chi_trong_so.pt"))
for epoch in range(3, 6):
    loss = mot_epoch(mh, toi_uu, lich)
print(f"C (thiếu optimizer): loss cuối = {loss:.6f}, lr = {lich.get_last_lr()[0]}")

C (thiếu optimizer): loss cuối = 0.177919, lr = 0.05

Kết quả đã khác: learning rate quay về lịch ban đầu (0,1 rồi giảm còn 0,05 thay vì 0,025) và vận tốc momentum bị mất. Ở ví dụ nhỏ này loss tình cờ thấp hơn, nhưng đó không phải điều đáng mừng: quá trình huấn luyện đã khác với cấu hình bạn thiết kế, không lặp lại được, và với mô hình lớn dùng Adam, learning rate bị "reset" lên cao có thể phá hỏng trọng số đã học tốt.

torch.load và weights_only

Từ PyTorch 2.6, torch.load mặc định weights_only=True: chỉ cho phép nạp tensor, số, chuỗi, list, dict và một số kiểu an toàn. Lý do: file .pt dùng định dạng pickle, và một file pickle độc hại có thể chạy mã tùy ý khi nạp. Checkpoint như trên (dict gồm state_dict, số nguyên, số thực) nạp bình thường. Nhưng một lỗi rất hay gặp là vô tình lưu kiểu NumPy:

import numpy as np

torch.save({"model": mh.state_dict(), "nguong": np.float64(0.42)}, "co_numpy.pt")
try:
    torch.load("co_numpy.pt")
except Exception as ex:
    print(type(ex).__name__, "-", str(ex).split(".")[0])

torch.save({"model": mh.state_dict(), "nguong": float(np.float64(0.42))}, "sach.pt")
print("nạp được:", torch.load("sach.pt")["nguong"])

UnpicklingError - Weights only load failed
nạp được: 0.42

Các giá trị như ngưỡng, F1, độ chính xác tính bằng NumPy hoặc scikit-learn thường là np.float64. Hãy chuyển sang float() trước khi lưu. Với checkpoint cũ do chính bạn tạo, có thể dùng torch.load(..., weights_only=False); tuyệt đối không làm vậy với file tải từ nguồn lạ. Nếu cần giữ weights_only=True mà vẫn có kiểu riêng, PyTorch cho phép khai báo kiểu an toàn bằng torch.serialization.add_safe_globals.

Quy trình lưu trong thực tế

1. last.pt: ghi đè sau mỗi epoch, để chạy tiếp khi bị gián đoạn.
2. best.pt: chỉ lưu khi chỉ số validation (ví dụ F1 lớp NG) tốt hơn mức tốt nhất trước đó.
3. Ghi an toàn: lưu ra last.pt.tmp rồi os.replace sang last.pt, tránh file hỏng nếu mất điện giữa lúc ghi.
4. Lưu kèm cấu hình: kích thước ảnh đầu vào, tên các lớp, phiên bản dữ liệu, dưới dạng dict thuần (chuỗi, số). Sau 3 tháng, bạn sẽ cảm ơn chính mình.
5. Dùng DataLoader có xáo trộn: để huấn luyện tiếp lặp lại hoàn toàn, lưu thêm trạng thái bộ sinh số ngẫu nhiên (torch.get_rng_state()). Trong thực tế, lặp lại xấp xỉ thường là đủ; quan trọng nhất là đúng optimizer và learning rate.

import os


def luu_an_toan(du_lieu, duong_dan):
    tam = duong_dan + ".tmp"
    torch.save(du_lieu, tam)
    os.replace(tam, duong_dan)       # thay thế nguyên tử trên cùng ổ đĩa

Khi mô hình đã xong, bước tiếp theo để triển khai là xuất sang ONNX và tăng tốc, xem tăng tốc suy luận: ONNX Runtime, OpenVINO, TensorRT. Thư viện như Ultralytics đã tự lưu last.pt và best.pt và có tùy chọn resume; hiểu cơ chế bên dưới giúp bạn xử lý khi có sự cố.

Nguồn tham khảo: PyTorch: Serialization semantics (weights_only), PyTorch Tutorials: Saving and loading models.