从 batch 到训练循环：能跑，还不等于在学习

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

def make_dataset(count, seed):
    generator = torch.Generator().manual_seed(seed)
    features = torch.randn(count, 2, generator=generator)
    labels = (features[:, 0] + 0.5 * features[:, 1] > 0).long()
    return TensorDataset(features, labels)

torch.manual_seed(17)
train_data = make_dataset(80, seed=5)
validation_data = make_dataset(47, seed=6)
train_loader = DataLoader(
    train_data, batch_size=16, shuffle=True,
    generator=torch.Generator().manual_seed(7),
)
validation_loader = DataLoader(validation_data, batch_size=16, shuffle=False)
model = nn.Linear(2, 2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.15, momentum=0.9)

features, labels = next(iter(validation_loader))
assert features.shape == (16, 2)
assert labels.shape == (16,) and labels.dtype == torch.long
assert model(features).shape == (16, 2)

取出特征和标签
  → 清除上次的梯度
  → model(features)：得到 logits
  → loss：按当前 batch 求平均
  → backward：计算并累积梯度
  → step：修改参数
  → 累计本轮的 loss 总量与样本数

def train_epoch(model, loader, optimizer):
    model.train()
    loss_total = 0.0
    sample_count = 0
    for features, labels in loader:
        optimizer.zero_grad(set_to_none=True)
        logits = model(features)
        loss = nn.functional.cross_entropy(logits, labels)
        loss.backward()
        optimizer.step()
        count = labels.numel()
        loss_total += loss.detach().item() * count
        sample_count += count
    if sample_count == 0:
        raise ValueError("Training loader is empty")
    return loss_total / sample_count

def evaluate(model, loader):
    previous_mode = model.training
    model.eval()
    loss_total = 0.0
    correct_count = 0
    sample_count = 0
    try:
        with torch.no_grad():
            for features, labels in loader:
                logits = model(features)
                loss_total += nn.functional.cross_entropy(
                    logits, labels, reduction="sum"
                ).item()
                correct_count += (logits.argmax(dim=-1) == labels).sum().item()
                sample_count += labels.numel()
    finally:
        model.train(previous_mode)
    if sample_count == 0:
        raise ValueError("Evaluation loader is empty")
    return {"loss": loss_total / sample_count, "accuracy": correct_count / sample_count}

initial_metrics = evaluate(model, validation_loader)
history = []
for epoch in range(25):
    train_loss = train_epoch(model, train_loader, optimizer)
    validation_metrics = evaluate(model, validation_loader)
    history.append({"epoch": epoch + 1, "train_loss": train_loss, **validation_metrics})

assert history[-1]["loss"] < initial_metrics["loss"]
assert 0.0 <= history[-1]["accuracy"] <= 1.0

class TinyClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(2, 4), nn.Linear(4, 2)])
        self.register_buffer("input_scale", torch.ones(2))

    def forward(self, features):
        hidden = torch.relu(self.layers[0](features / self.input_scale))
        return self.layers[1](hidden)

registered_model = TinyClassifier()
assert "layers.0.weight" in dict(registered_model.named_parameters())
assert "input_scale" not in dict(registered_model.named_parameters())
assert "input_scale" in registered_model.state_dict()
assert registered_model(torch.ones(3, 2)).shape == (3, 2)

import copy
import io

snapshot = copy.deepcopy({
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "epoch": len(history),
})
buffer = io.BytesIO()
torch.save(snapshot, buffer)
buffer.seek(0)
loaded = torch.load(buffer, weights_only=True)
restored_model = nn.Linear(2, 2)
restored_model.load_state_dict(loaded["model"])
restored_optimizer = torch.optim.SGD(restored_model.parameters(), lr=0.15, momentum=0.9)
restored_optimizer.load_state_dict(loaded["optimizer"])
with torch.no_grad():
    torch.testing.assert_close(model(features), restored_model(features))
assert loaded["epoch"] == 25

def write_tensorboard(history, log_dir):
    from torch.utils.tensorboard import SummaryWriter
    with SummaryWriter(log_dir=log_dir) as writer:
        for entry in history:
            writer.add_scalar("loss/train", entry["train_loss"], entry["epoch"])
            writer.add_scalar("loss/validation", entry["loss"], entry["epoch"])
            writer.add_scalar("accuracy/validation", entry["accuracy"], entry["epoch"])

def device_step(model, optimizer, features, labels):
    device = next(model.parameters()).device
    features = features.to(device, non_blocking=True)
    labels = labels.to(device, non_blocking=True)
    model.train()
    optimizer.zero_grad(set_to_none=True)
    loss = nn.functional.cross_entropy(model(features), labels)
    loss.backward()
    optimizer.step()
    return loss.detach()