AI 工程基础体系 · 第 10/100 篇。内容覆盖机器学习、深度学习与生成式 AI；模型、数据、评测、权限和成本会作为同一生产系统处理。

PyTorch 工程基础：Tensor、Autograd、Dataset、训练循环和检查点

PyTorch 训练程序可以抽象为一条数据流：

flowchart LR
    A[原始样本] --> B[Dataset]
    B --> C[DataLoader]
    C --> D[Tensor: device/dtype/shape]
    D --> E[模型 forward]
    E --> F[预测值]
    F --> G[损失函数]
    G --> H[Autograd 计算梯度]
    H --> I[Optimizer 更新参数]
    I --> J[评测与日志]
    I --> K[Checkpoint 保存状态]
    K --> L[恢复训练或部署]

这条链路中，Tensor 是数据和参数的基本载体，Autograd 负责从损失反向计算梯度，Dataset 和 DataLoader 负责把样本组织成批次，训练循环负责按顺序执行前向、反向和参数更新，检查点则保存足以恢复实验状态的数据。

它们不是互相独立的 API。形状或数据类型错误会在模型前向阶段暴露，梯度生命周期错误会导致参数不更新，数据切分错误会使验证指标失真，检查点状态不完整则会使“恢复训练”实际上变成一次新的训练。

一、Tensor：带有形状、类型和设备的数值对象

1. Tensor 的四个核心属性

PyTorch 的 torch.Tensor 是多维数值数组。工程上最重要的属性通常是：

  • shape：各维大小；
  • dtype：元素类型，例如 torch.float32、torch.float16、torch.int64；
  • device：所在设备，例如 CPU 或某个 CUDA GPU；
  • requires_grad：是否需要对该 Tensor 建立梯度关系。
import torch

x = torch.randn(4, 3, dtype=torch.float32, device="cpu", requires_grad=True)

print(x.shape)          # torch.Size([4, 3])
print(x.dtype)          # torch.float32
print(x.device)         # cpu
print(x.requires_grad)  # True

这四个属性分别回答了不同问题：

  • shape 决定矩阵乘法、卷积和批处理是否成立；
  • dtype 决定可表示的数值范围、精度和部分算子的可用性；
  • device 决定运算发生在哪里；
  • requires_grad 决定 Autograd 是否追踪相关运算。

一个常见误解是认为 Tensor 只要形状一致就能参与运算。实际上，参与同一个算子的 Tensor 通常还必须满足设备兼容和数据类型兼容要求。例如，CPU Tensor 不能直接与 CUDA Tensor 相加：

cpu_x = torch.ones(2)
gpu_x = torch.ones(2, device="cuda")

# RuntimeError：设备不一致
# cpu_x + gpu_x

模型参数和输入必须放在同一设备上：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = torch.nn.Linear(3, 2).to(device)
x = torch.randn(8, 3, device=device)
y = model(x)

.to(device) 会返回移动或转换后的 Tensor；对 nn.Module 使用时会就地修改其参数和缓冲区所在设备，但显式写成 model = model.to(device) 更容易表达意图。

2. 形状、批次和广播

假设输入为一个批次的三维特征：

X∈RB×DX \in \mathbb{R}^{B \times D}

其中：

  • BB 是批次大小；
  • DD 是特征维度。

一个线性层的参数为：

W∈RD×K,b∈RKW \in \mathbb{R}^{D \times K}, \quad b \in \mathbb{R}^{K}

前向计算为：

Y=XW+bY = XW + b

于是：

(B,D)(D,K)→(B,K)(B,D)(D,K) \rightarrow (B,K)

偏置 bb 的形状是 (K,)，它通过广播参与每个样本的计算。

x = torch.randn(8, 3)       # [B, D] = [8, 3]
linear = torch.nn.Linear(3, 2)

y = linear(x)               # [8, 2]
print(y.shape)

广播不是任意复制。PyTorch 从最后一维开始比较两个形状，两个维度相等、其中一个为 1，或其中一方不存在时，才可能广播。例如：

a = torch.randn(8, 3)
b = torch.randn(3)

c = a + b                   # b 被视为对每个 batch 复用，结果 [8, 3]

但下面的形状不满足广播规则：

a = torch.randn(8, 3)
b = torch.randn(8, 2)

# RuntimeError：最后一维 3 和 2 不兼容
# a + b

不要把“能够广播”误认为“语义正确”。例如标签形状为 [B, 1]，预测形状为 [B] 时，某些逐元素损失可能广播成 [B, B]，程序不一定立即报错，但损失已经不再表示逐样本比较。训练前应显式检查：

assert prediction.shape == target.shape

3. reshape、view、切片与存储共享

reshape 和 view 都可以改变张量形状，但底层机制可能不同。view 要求元素在内存中具有兼容的连续布局；reshape 会在可能时返回视图，否则创建拷贝。

x = torch.arange(12).reshape(3, 4)

y = x.view(12)
z = x.reshape(12)

print(y.shape, z.shape)  # torch.Size([12]) torch.Size([12])

视图与原 Tensor 可能共享存储，因此修改视图可能影响原对象。切片通常也可能产生视图：

x = torch.zeros(4)
part = x[:2]
part[0] = 1

print(x)  # tensor([1., 0., 0., 0.])

这在数据预处理和 Autograd 中都重要。若需要独立数据，可以使用：

copy = x.clone()

但 clone() 会保留梯度关系；若需要创建一个不参与原计算图的新数据，通常使用：

copy = x.detach().clone()

4. 原地操作的风险

带下划线的方法，例如 add_()、relu_()，以及通过索引直接赋值，属于原地操作。原地修改可能破坏 Autograd 为反向传播保存的中间值：

x = torch.tensor([2.0], requires_grad=True)
y = x * x
# 某些情况下修改参与计算图的值会导致反向传播失败

典型错误是：

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation

工程上不要为了少一次内存分配就随意使用原地操作。只有在明确知道该 Tensor 不再需要旧值，并且已经验证反向传播正确时，才考虑原地版本。

二、Autograd：从损失到参数梯度

1. 计算图和链式法则

Autograd 会记录由 requires_grad=True 的 Tensor 参与的可微运算，并形成动态计算图。调用 loss.backward() 时，它从损失开始沿图反向应用链式法则。

考虑最简单的模型：

y^=wx\hat{y}=wx

损失为均方误差：

L=(y^−y)2=(wx−y)2L=(\hat{y}-y)^2=(wx-y)^2

对参数 ww 求导：

∂L∂w=2(wx−y)∂(wx−y)∂w=2(wx−y)x\frac{\partial L}{\partial w} = 2(wx-y)\frac{\partial(wx-y)}{\partial w} = 2(wx-y)x

运行代码：

x = torch.tensor(3.0)
target = torch.tensor(10.0)
w = torch.tensor(2.0, requires_grad=True)

prediction = w * x
loss = (prediction - target) ** 2

loss.backward()

print(prediction.item())  # 6.0
print(loss.item())        # 16.0
print(w.grad.item())      # -24.0

中间结果是：

  1. prediction = 2 × 3 = 6；
  2. 误差 prediction - target = -4；
  3. 损失 (-4)^2 = 16；
  4. 梯度 2 × (-4) × 3 = -24。

梯度为负，说明在当前点增大 ww 会降低损失，因此梯度下降更新：

wnew=w−η∂L∂ww_{\text{new}}=w-\eta\frac{\partial L}{\partial w}

若学习率 η=0.1\eta=0.1，则：

wnew=2−0.1×(−24)=4.4w_{\text{new}}=2-0.1\times(-24)=4.4

参数朝着增大方向移动，符合直觉。

2. 梯度默认累积

backward() 不会自动清空已有梯度，而是把新梯度累加到 .grad 中：

w = torch.tensor(2.0, requires_grad=True)

for _ in range(2):
    loss = (w * 3 - 10) ** 2
    loss.backward()

print(w.grad.item())  # -48.0，而不是 -24.0

因此标准训练循环必须在每次更新前清空梯度：

optimizer.zero_grad()
loss.backward()
optimizer.step()

也可以使用：

optimizer.zero_grad(set_to_none=True)

这通常能避免写入全零 Tensor，并让“没有梯度”和“梯度恰好为零”在状态上区分开来。使用 set_to_none=True 后，某些未参与当前计算的参数 .grad 会保持为 None，代码不能假设所有参数都有 Tensor 类型的梯度。

3. requires_grad、detach 与推理模式

模型参数通常设置 requires_grad=True。输入是否需要梯度取决于任务，例如对抗样本、输入优化或可解释性分析可能需要输入梯度，普通训练不需要。

detach() 会返回一个不再连接到原计算图的 Tensor：

x = torch.tensor(2.0, requires_grad=True)
y = x * 3
z = y.detach()

print(z.requires_grad)  # False

如果在模型内部过早 detach()，梯度无法继续传回前面的模块：

h = encoder(x)
h = h.detach()          # encoder 不再获得来自后续损失的梯度
out = decoder(h)

这可能是有意的冻结边界，也可能是导致模型“部分不学习”的隐蔽 bug。

评测和推理时应关闭梯度追踪：

model.eval()

with torch.inference_mode():
    prediction = model(x)

torch.no_grad() 和 torch.inference_mode() 都用于不记录梯度；后者还会启用更严格的推理优化和限制。若推理结果需要继续参与可微计算，不能使用 inference_mode()。

model.eval() 是另一件事：它切换 Dropout、BatchNorm 等模块的行为，但不会关闭 Autograd。反过来，torch.no_grad() 也不会把模型切换到评估模式。验证阶段通常需要两者同时使用：

model.eval()
with torch.inference_mode():
    ...

4. 冻结参数与 eval() 的区别

冻结参数意味着不为某些参数计算梯度：

for parameter in backbone.parameters():
    parameter.requires_grad = False

eval() 只改变模块行为。一个冻结的 BatchNorm 仍可能在训练模式下更新运行统计量；一个处于 eval() 的模块，其参数仍然可以有梯度。两者不能互相替代。

三、Dataset 和 DataLoader：把样本变成稳定批次

1. Dataset 的契约

Dataset 描述“第 ii 个样本如何取得”。最小接口是：

  • __len__()：样本总数；
  • __getitem__(index)：返回一个样本。
from torch.utils.data import Dataset

class RegressionDataset(Dataset):
    def __init__(self, x, y):
        assert len(x) == len(y)
        self.x = x
        self.y = y

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

    def __getitem__(self, index):
        return self.x[index], self.y[index]

Dataset 不负责模型、不负责反向传播，也不应该把一个完整训练周期塞进 __getitem__()。它的职责是稳定地提供样本及其必要的预处理。

IterableDataset 则适合日志流、超大数据集或无法随机索引的数据源。它通过 __iter__() 产生样本。使用多进程加载时，必须根据 worker 信息切分数据，否则每个 worker 可能重复读取同一数据流。

2. DataLoader 的批处理过程

DataLoader 在 Dataset 之上提供：

  • 批次组织；
  • 可选随机打乱；
  • 多进程加载；
  • collate_fn 将多个样本组合成一个 batch；
  • 可选固定长度或自定义采样策略。
from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=0,
    pin_memory=torch.cuda.is_available(),
)

默认 collate_fn 能把形状一致的 Tensor 堆叠起来。例如两个样本：

(x1, y1), (x2, y2)

会变成：

x_batch = torch.stack([x1, x2])
y_batch = torch.stack([y1, y2])

如果文本序列长度不同，默认堆叠会失败，需要自定义 padding：

def collate_batch(samples):
    sequences, labels = zip(*samples)
    lengths = torch.tensor([len(s) for s in sequences])

    max_len = max(lengths).item()
    padded = torch.zeros(len(sequences), max_len, dtype=torch.long)

    for i, seq in enumerate(sequences):
        padded[i, :len(seq)] = torch.tensor(seq, dtype=torch.long)

    return padded, lengths, torch.tensor(labels)

自定义批处理必须明确 padding 值、长度信息和 mask 的语义。Transformer 中如果把 padding token 当成真实 token，注意力和损失都会被污染。

3. 数据切分和泄漏

训练、验证和测试集承担不同职责：

  • 训练集用于更新参数；
  • 验证集用于选择超参数、阈值和检查过拟合；
  • 测试集只用于最终报告。

随机切分并不总是正确。时间序列应按时间切分，用户级任务应按用户切分，医学或企业数据可能需要按实体或机构切分。若同一用户的记录同时出现在训练和验证中，模型可能记住用户特征，验证分数会虚高。

归一化也必须遵守切分边界。正确流程是只用训练集估计均值和标准差：

μtrain=1N∑ixi,σtrain=1N∑i(xi−μtrain)2\mu_{\text{train}}=\frac{1}{N}\sum_i x_i,\qquad \sigma_{\text{train}}=\sqrt{\frac{1}{N}\sum_i(x_i-\mu_{\text{train}})^2}

然后将同一组统计量应用于验证和测试集。若用全量数据计算统计量，测试集信息已经泄漏到训练流程中。

四、训练循环：状态变化的完整顺序

1. 一个训练 step 做了什么

对一个批次，训练过程通常是：

  1. 从 DataLoader 取出 (x, y)；
  2. 将数据移动到目标设备；
  3. 清空上一批次的梯度；
  4. 执行前向计算 prediction = model(x)；
  5. 计算损失 loss = criterion(prediction, y)；
  6. 反向计算梯度 loss.backward()；
  7. 可选地裁剪梯度；
  8. 执行参数更新 optimizer.step()；
  9. 记录损失、吞吐量和异常状态。

参数更新以 SGD 为例：

θt+1=θt−ηtgt\theta_{t+1}=\theta_t-\eta_t g_t

其中：

  • θt\theta_t 是第 tt 步参数；
  • gt=∇θLtg_t=\nabla_\theta L_t 是当前批次梯度；
  • ηt\eta_t 是当前学习率。

Adam 还维护一阶和二阶矩估计，因此其优化器状态不只是学习率。恢复训练时如果只保存模型参数而不保存优化器状态，恢复后的轨迹通常不会与中断前连续。

2. 完整可运行示例

下面的代码使用合成二分类数据，避免依赖外部数据下载。它展示 Dataset、DataLoader、模型、训练、验证和检查点的基本生命周期。

import os
import random
from pathlib import Path

import numpy as np
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader, random_split


class BinaryDataset(Dataset):
    def __init__(self, n=2000, seed=1234):
        generator = torch.Generator().manual_seed(seed)

        self.x = torch.randn(n, 2, generator=generator)
        noise = 0.3 * torch.randn(n, generator=generator)
        score = self.x[:, 0] + 0.8 * self.x[:, 1] + noise
        self.y = (score > 0).long()

    def __len__(self):
        return self.x.shape[0]

    def __getitem__(self, index):
        return self.x[index], self.y[index]


class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 16),
            nn.ReLU(),
            nn.Linear(16, 2),
        )

    def forward(self, x):
        return self.net(x)


def seed_everything(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def evaluate(model, loader, device):
    model.eval()
    total_loss = 0.0
    total_correct = 0
    total_count = 0
    criterion = nn.CrossEntropyLoss()

    with torch.inference_mode():
        for x, y in loader:
            x = x.to(device)
            y = y.to(device)

            logits = model(x)
            loss = criterion(logits, y)

            batch_size = y.size(0)
            total_loss += loss.item() * batch_size
            total_correct += (logits.argmax(dim=1) == y).sum().item()
            total_count += batch_size

    return total_loss / total_count, total_correct / total_count


def save_checkpoint(path, epoch, model, optimizer, scheduler, best_val_loss):
    checkpoint = {
        "epoch": epoch,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict() if scheduler else None,
        "best_val_loss": best_val_loss,
        "torch_rng_state": torch.get_rng_state(),
        "cuda_rng_state_all": (
            torch.cuda.get_rng_state_all()
            if torch.cuda.is_available() else None
        ),
    }

    path = Path(path)
    temporary_path = path.with_suffix(path.suffix + ".tmp")
    torch.save(checkpoint, temporary_path)
    os.replace(temporary_path, path)


def main():
    seed_everything(42)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dataset = BinaryDataset()

    train_size = int(0.8 * len(dataset))
    val_size = len(dataset) - train_size
    split_generator = torch.Generator().manual_seed(42)

    train_set, val_set = random_split(
        dataset,
        [train_size, val_size],
        generator=split_generator,
    )

    train_loader = DataLoader(
        train_set,
        batch_size=64,
        shuffle=True,
        num_workers=0,
        pin_memory=(device.type == "cuda"),
    )
    val_loader = DataLoader(
        val_set,
        batch_size=256,
        shuffle=False,
        num_workers=0,
        pin_memory=(device.type == "cuda"),
    )

    model = Classifier().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=20,
    )

    best_val_loss = float("inf")
    checkpoint_path = "classifier.pt"

    for epoch in range(20):
        model.train()
        running_loss = 0.0
        seen = 0

        for x, y in train_loader:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)

            optimizer.zero_grad(set_to_none=True)

            logits = model(x)
            loss = criterion(logits, y)

            if not torch.isfinite(loss):
                raise FloatingPointError(
                    f"non-finite loss at epoch={epoch}: {loss.item()}"
                )

            loss.backward()

            # 防止极端情况下梯度范数过大。
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

            optimizer.step()

            batch_size = y.size(0)
            running_loss += loss.item() * batch_size
            seen += batch_size

        scheduler.step()

        train_loss = running_loss / seen
        val_loss, val_accuracy = evaluate(model, val_loader, device)

        print(
            f"epoch={epoch + 1:02d} "
            f"train_loss={train_loss:.4f} "
            f"val_loss={val_loss:.4f} "
            f"val_acc={val_accuracy:.4f} "
            f"lr={optimizer.param_groups[0]['lr']:.6g}"
        )

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            save_checkpoint(
                checkpoint_path,
                epoch=epoch,
                model=model,
                optimizer=optimizer,
                scheduler=scheduler,
                best_val_loss=best_val_loss,
            )

    # 恢复“最佳验证损失”对应的模型
    checkpoint = torch.load(
        checkpoint_path,
        map_location=device,
        weights_only=False,  # 某些旧版本没有该参数，需按实际 PyTorch 版本调整
    )
    model.load_state_dict(checkpoint["model"])
    model.eval()

    print("restored best validation loss:", checkpoint["best_val_loss"])


if __name__ == "__main__":
    main()

运行前需要安装 PyTorch 和 NumPy。CPU 上可以运行，CUDA 可用时会自动使用 GPU。输出中的损失和准确率不应被视为固定数字；随机种子、PyTorch 版本、硬件和算子实现都可能造成差异，但验证损失通常应明显低于初始状态。

这里有几个容易忽略的因果关系：

  • CrossEntropyLoss 接受的是未归一化 logits，内部包含适合数值计算的 log_softmax 逻辑；不要先手动 softmax 再传入，除非明确使用了不同损失；
  • 标签必须是类别索引，通常为 torch.long，形状为 [B]；
  • scheduler.step() 的调用频率必须与调度器设计一致。按 epoch 调度的调度器应每个 epoch 调用一次，按 step 调度的调度器则应每个优化器更新后调用；
  • 日志中的平均损失应按样本数加权，而不是简单平均每个 batch 的损失，因为最后一个 batch 可能更小。

3. 梯度累积和有效批次

显存不足时，可以把多个小批次累积后再更新一次参数。若累积步数为 KK，应将每个小批次损失除以 KK：

optimizer.zero_grad(set_to_none=True)

for step, (x, y) in enumerate(train_loader):
    x = x.to(device)
    y = y.to(device)

    logits = model(x)
    loss = criterion(logits, y) / accumulation_steps
    loss.backward()

    if (step + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

除以 KK 的原因是：

∇(1K∑k=1KLk)=1K∑k=1K∇Lk\nabla\left(\frac{1}{K}\sum_{k=1}^{K}L_k\right) = \frac{1}{K}\sum_{k=1}^{K}\nabla L_k

如果不除，梯度大约放大 KK 倍，等价于改变了优化步长。循环结束时若剩余 batch 数不足 accumulation_steps，还需要额外执行一次 optimizer.step()，否则这些梯度会被丢弃。

4. 混合精度与稳定性

混合精度通常让部分计算使用低精度 Tensor，同时保留必要的高精度状态。它涉及三个概念：

  • autocast：根据算子选择合适精度；
  • GradScaler：在低精度训练中放大损失，减少梯度下溢；
  • optimizer step：检测到非有限梯度时跳过不安全更新。

当前 PyTorch 文档推荐使用 torch.amp 命名空间；具体构造方式会随 PyTorch 版本演进，因此生产代码应以所部署版本的官方文档为准。一个 CUDA 训练片段可以写成：

use_amp = device.type == "cuda"
scaler = torch.amp.GradScaler("cuda", enabled=use_amp)

for x, y in train_loader:
    x = x.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)

    optimizer.zero_grad(set_to_none=True)

    with torch.autocast(
        device_type=device.type,
        dtype=torch.float16,
        enabled=use_amp,
    ):
        logits = model(x)
        loss = criterion(logits, y)

    scaler.scale(loss).backward()

    # 裁剪前必须先反缩放，否则裁剪的是放大后的梯度
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

    scaler.step(optimizer)
    scaler.update()

混合精度不是“所有 Tensor 都改成 float16”。参数主副本、归一化和损失计算可能仍使用 float32，具体由 autocast 和算子实现决定。若出现 NaN、inf 或训练突然不收敛，应检查学习率、输入范围、损失缩放、梯度范数和数据中是否存在非有限值。

五、训练状态：train()、eval()、参数和优化器不是一回事

一个模型至少有三类不同状态：

  1. 参数状态：model.parameters() 中的权重和偏置；
  2. 模块缓冲区：例如 BatchNorm 的运行均值和方差；
  3. 优化器状态：例如 Adam 的一阶、二阶矩估计；
  4. 执行模式：训练模式或评估模式；
  5. 随机状态：Python、NumPy、PyTorch 和 CUDA 的随机数状态。

model.state_dict() 主要包含参数和持久化缓冲区，不包含优化器状态，也不包含随机数状态。model.train() 和 model.eval() 只切换模块行为，不会保存或恢复参数。

如果验证阶段忘记调用 model.eval()：

  • Dropout 仍会随机丢弃激活；
  • BatchNorm 可能继续使用当前 batch 统计量，并更新运行统计量；
  • 验证结果可能波动或污染后续训练状态。

如果训练阶段忘记调用 model.train()，Dropout 和 BatchNorm 会一直处于推理行为，模型可能无法按预期训练。

六、检查点：恢复的不只是模型权重

1. state_dict 与完整检查点

只保存权重：

torch.save(model.state_dict(), "weights.pt")

适合发布模型或仅做推理，但不够恢复训练。完整训练检查点通常至少包含：

{
    "epoch": ...,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict(),
    "scaler": scaler.state_dict(),       # 使用混合精度时
    "best_metric": ...,
    "config": ...,
    "rng_states": ...,
}

恢复时顺序应与对象依赖一致：

checkpoint = torch.load(
    "checkpoint.pt",
    map_location=device,
    weights_only=False,
)

model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])

if checkpoint["scheduler"] is not None:
    scheduler.load_state_dict(checkpoint["scheduler"])

start_epoch = checkpoint["epoch"] + 1
best_metric = checkpoint["best_metric"]

先构造同样的模型结构、优化器和调度器，再加载状态。若改变了模型层名称、参数形状或优化器参数组，load_state_dict 可能报 missing/unexpected keys 或 shape mismatch；强行使用 strict=False 会隐藏未加载参数，必须检查返回结果。

2. 原子保存与故障路径

直接覆盖目标文件存在中断风险：进程在写入过程中崩溃、机器掉电或磁盘空间耗尽，目标文件可能损坏。常见的恢复策略是：

  1. 保存到临时文件；
  2. 写入完成后执行原子替换；
  3. 保留最近检查点和最佳指标检查点；
  4. 恢复时先校验文件能否读取和关键键是否存在。

前面的 save_checkpoint() 使用 os.replace() 完成同一文件系统内的原子替换。它不能解决所有问题：跨文件系统移动、对象存储上传和网络文件系统仍需额外校验，例如保存哈希、版本号或完成标记。

3. “恢复训练”不等于“加载权重后继续循环”

如果只加载 model.state_dict()，而不加载 Adam 状态、学习率调度器状态和混合精度 scaler，参数虽然相同，但下一步更新通常不同。若还重新初始化 DataLoader 的随机采样器，后续 batch 顺序也可能不同。

因此应区分：

  • 权重恢复：用于推理、微调初始化；
  • 训练恢复：希望尽可能接近中断前轨迹，需要模型、优化器、调度器、scaler、epoch/step、随机状态及数据采样状态。

即使保存了这些状态，也不保证跨硬件、跨 PyTorch 版本或跨非确定性 CUDA 算子完全复现。确定性设置可能降低性能，且部分算子仍可能不支持确定性实现。复现实验应同时保存代码版本、依赖版本、配置、数据版本和运行环境。

4. 安全加载边界

检查点本质上是序列化文件。不要从不可信来源直接反序列化任意对象。对于只包含权重的文件，优先使用当前 PyTorch 版本支持的安全权重加载选项；对于历史版本或包含自定义 Python 对象的检查点，必须把加载动作限制在可信环境中，并验证文件来源。

七、常见失败表现与诊断路径

损失不下降

依次检查：

  1. 输入和标签的 shape、dtype、device；
  2. 标签编码是否与损失函数匹配；
  3. 是否执行了 optimizer.zero_grad()、backward() 和 step()；
  4. 参数是否真的有梯度：
for name, parameter in model.named_parameters():
    if parameter.requires_grad:
        print(name, parameter.grad is None)
  1. 学习率是否过大或过小；
  2. 是否错误地在 forward 中使用了 detach()；
  3. 是否在每一步都重新初始化模型或优化器。

GPU 显存持续增长

常见原因是把带计算图的 Tensor 保存到了列表：

losses.append(loss)  # 可能保留整张计算图

若只需要日志值，应写：

losses.append(loss.item())

保存预测结果时通常使用：

predictions.append(prediction.detach().cpu())

另一个原因是验证阶段忘记使用 no_grad() 或 inference_mode()，导致不必要地构建计算图。

训练和验证指标异常接近或异常好

应检查数据泄漏：

  • 训练和验证是否包含同一实体；
  • 预处理统计量是否使用了验证或测试数据；
  • 标签或未来信息是否混入特征；
  • 是否错误地在测试集上反复调参。

指标计算本身也要与任务匹配。分类准确率在类别极不平衡时可能掩盖模型完全忽略少数类，应同时查看混淆矩阵、精确率、召回率或适合业务目标的代价函数。

恢复后结果与中断前不一致

检查：

  • 是否恢复了 optimizer 和 scheduler；
  • 是否从正确的 epoch 或 global_step 开始；
  • 是否恢复 AMP scaler；
  • 是否重新调用了与之前不同的 scheduler.step()；
  • 是否改变了 batch size、数据顺序或梯度累积步数；
  • 是否加载到了错误的“最佳模型”而不是“最近模型”。

保存“最佳检查点”和“最近检查点”是两种不同需求：最佳检查点用于部署或最终评测，最近检查点用于故障恢复。二者不应只保留一个文件。

八、从训练到生成式 AI 和模型服务

在 Transformer 或生成式 AI 中，Tensor 的形状语义会进一步扩展。例如隐藏状态常见形状为：

[B,T,H][B,T,H]

其中 BB 是批次大小，TT 是序列长度，HH 是隐藏维度。注意力计算还会产生与 T2T^2 相关的中间张量，因此批处理、序列长度和 GPU 容量直接相互制约。

训练时通常使用 teacher forcing：标签序列参与损失计算，模型一次处理整个序列；推理时则可能逐 token 生成。服务系统为提高吞吐，会把请求组成批次或连续批次，但必须处理不同请求的长度、padding、停止条件和 KV Cache 生命周期。训练检查点中的优化器状态通常不需要随推理服务加载；服务更关心模型权重、分词器、配置、量化格式和显存占用。

量化、混合精度和批处理都可能降低单请求成本或提高吞吐，但会改变数值行为和容量边界。部署前应分别测量：

  • 模型权重占用；
  • KV Cache 占用；
  • 激活和临时 Tensor 占用；
  • 不同 batch 与序列长度下的峰值显存；
  • 延迟、吞吐和错误率；
  • 不同精度或量化配置下的任务质量。

权限也属于同一生产系统的一部分。训练任务需要读取数据、写入检查点和日志；服务进程可能只需要读取模型和分词器。让在线服务拥有训练数据写权限，会扩大故障和数据泄漏影响面。检查点路径、数据路径和日志内容都应纳入访问控制与审计。

九、工程上的最小验收标准

一个可交付的 PyTorch 训练程序至少应能回答以下问题：

  • 每个 Tensor 的 shape、dtype 和 device 是否明确；
  • Dataset 返回的单样本结构能否被 DataLoader 正确批处理；
  • train() 和 eval() 是否在正确边界切换；
  • 每次参数更新前梯度是否正确清空；
  • loss 是否有限，梯度是否存在且没有异常爆炸；
  • 优化器和学习率调度器是否按预期更新；
  • 验证集是否独立于训练和调参过程；
  • 检查点能否在另一进程中加载；
  • 加载后模型是否能完成一次前向和评测；
  • 训练中断后是否能从正确的 epoch 或 step 继续；
  • 推理服务加载的是否是经过验证的权重和配置，而不是任意临时文件。

PyTorch 的 API 可以很短，但训练系统的正确性来自状态、形状、数据边界和生命周期都被明确管理。掌握 Tensor、Autograd、Dataset、训练循环和检查点，实质上是在掌握从数据进入模型、梯度改变参数、指标验证结果到故障恢复和部署交付的完整闭环。

系列导航与关联阅读

官方资料

本文依据研究论文、标准组织与主流框架官方文档重新梳理；正文、示例与工程清单由 WR BLOG 编写。