训练循环、DataLoader 与断点恢复

import io
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader

torch.set_num_threads(1)
torch.use_deterministic_algorithms(True)
x = torch.linspace(-1, 1, 24).reshape(12, 2)
y = (x.sum(dim=1) > 0).long()
data = TensorDataset(x, y)

def create():
    torch.manual_seed(17)  # 模型初始化与Dropout使用全局随机状态
    model = nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Dropout(.2), nn.Linear(8, 2))
    optimizer = torch.optim.AdamW(model.parameters(), lr=.05)
    generator = torch.Generator().manual_seed(29)  # 数据顺序使用独立发生器
    return model, optimizer, generator

def train(model, optimizer, generator, epochs):
    loader = DataLoader(data, batch_size=4, shuffle=True, generator=generator, num_workers=0)
    for _ in range(epochs):
        model.train()
        for bx, by in loader:
            optimizer.zero_grad(set_to_none=True)
            loss = nn.functional.cross_entropy(model(bx), by)
            assert torch.isfinite(loss)
            loss.backward()
            optimizer.step()

full, full_opt, gen = create()
full.eval()
with torch.no_grad():
    initial = nn.functional.cross_entropy(full(x), y).item()
train(full, full_opt, gen, 6)
part, part_opt, part_gen = create()
train(part, part_opt, part_gen, 3)
buffer = io.BytesIO()
torch.save({"model": part.state_dict(), "optimizer": part_opt.state_dict(),
            "rng": torch.get_rng_state(), "generator": part_gen.get_state(), "epoch": 3}, buffer)
buffer.seek(0)
saved = torch.load(buffer, weights_only=True)
resumed, resumed_opt, resumed_gen = create()
resumed.load_state_dict(saved["model"])
resumed_opt.load_state_dict(saved["optimizer"])
resumed_gen.set_state(saved["generator"])
torch.set_rng_state(saved["rng"])  # 必须在重建对象之后恢复
train(resumed, resumed_opt, resumed_gen, 6 - saved["epoch"])
assert all(torch.equal(a, b) for a, b in zip(full.parameters(), resumed.parameters()))
assert all(torch.equal(full_opt.state[a]["exp_avg"], resumed_opt.state[b]["exp_avg"])
           for a, b in zip(full.parameters(), resumed.parameters()))
assert len(list(DataLoader(data, batch_size=5, drop_last=False))) == 3
assert len(list(DataLoader(data, batch_size=5, drop_last=True))) == 2
resumed.eval()
with torch.no_grad():
    first, second = resumed(x), resumed(x)
    final = nn.functional.cross_entropy(first, y).item()
assert torch.equal(first, second)  # eval使Dropout关闭
assert final < initial
print(f"CPU checkpoint checks passed; loss {initial:.4f} -> {final:.4f}")