Python 基础体系 · 第 98/112 篇。示例统一以 Python 3.14 为语言基线；第三方库使用与其兼容的现代稳定版本，版本敏感行为会单独说明。

Python PyTorch 工程：Tensor、Autograd、Dataset、训练和检查点

PyTorch 工程的核心不是“调用一个模型然后得到结果”，而是维护一条可验证的数据流：

原始数据
  ↓
Dataset / DataLoader
  ↓
Tensor：形状、dtype、device
  ↓
nn.Module 前向计算
  ↓
loss
  ↓
Autograd 反向传播
  ↓
optimizer 更新参数
  ↓
验证、日志、检查点

其中：

  • Tensor 表示参与数值计算的数据；
  • Autograd 记录计算过程并计算梯度；
  • Dataset 定义“一个样本是什么”；
  • DataLoader 定义如何批量、打乱、并发读取样本；
  • 训练 是重复执行前向、损失、反向和参数更新；
  • 检查点 保存足够多的状态，使训练可以恢复，而不只是重新加载模型参数。

PyTorch 官方当前的本地安装页面将 Python 3.10–3.14 列为支持范围，并要求根据操作系统、安装方式和 CPU/CUDA/ROCm 平台选择对应命令；因此本文代码以 Python 3.14 为目标，但实际安装命令应以目标平台生成器为准。(pytorch.org)

一、环境和第一个 Tensor

建议先创建隔离环境：

python3.14 -m venv .venv
source .venv/bin/activate        # Linux/macOS
# Windows PowerShell:
# .venv\Scripts\Activate.ps1

python -m pip install --upgrade pip

然后根据机器平台安装 PyTorch。CPU 环境可以使用本机安装页面生成的 CPU 命令；有 NVIDIA GPU 时，应选择与驱动和 CUDA 运行时匹配的构建，而不是随意复制其他机器的 CUDA 命令。

安装后验证：

# verify_torch.py
import torch

print("torch:", torch.__version__)
print("python-compatible tensor:", torch.rand(2, 3))
print("cuda available:", torch.cuda.is_available())

torch.rand(2, 3) 创建一个形状为 (2, 3) 的浮点 Tensor。形状表示两个维度的长度，不等于“二维数组一定只能有两个轴”；例如图像批次通常使用 (batch, channel, height, width) 四个轴。

1. Tensor 的五个基本属性

一个 Tensor 至少需要从以下几个方面理解：

import torch

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

print(x)
print("shape:", x.shape)
print("ndim:", x.ndim)
print("dtype:", x.dtype)
print("device:", x.device)
print("stride:", x.stride())

典型输出：

tensor([[ 0.,  1.,  2.,  3.],
        [ 4.,  5.,  6.,  7.],
        [ 8.,  9., 10., 11.]])
shape: torch.Size([3, 4])
ndim: 2
dtype: torch.float32
device: cpu
stride: (4, 1)

这里：

  • shape=(3, 4) 表示三行四列；
  • dtype=torch.float32 表示每个元素使用 32 位浮点数；
  • device=cpu 表示数据存放在 CPU；
  • stride=(4, 1) 表示沿第一个维度移动一个元素，需要跨过 4 个底层元素，沿第二个维度移动一个元素只需跨过 1 个底层元素。

stride 与 NumPy 的内存布局概念相似。转置通常只改变形状和步长，不立即复制底层数据：

y = x.T

print("y.shape:", y.shape)
print("y.stride:", y.stride())
print("y.is_contiguous:", y.is_contiguous())

y 的形状变为 (4, 3)，但通常不是 contiguous。需要连续内存的算子可以显式调用：

y_contiguous = y.contiguous()

contiguous() 可能分配并复制内存，因此不能把它当成无成本的“格式转换”。

2. 创建 Tensor 时的复制和共享

以下 API 的语义不同：

import numpy as np
import torch

array = np.array([1.0, 2.0, 3.0], dtype=np.float32)

a = torch.tensor(array)       # 通常复制数据
b = torch.from_numpy(array)   # 与 NumPy 数组共享存储
c = torch.as_tensor(array)    # 尽量避免复制

修改 array 可能影响 b 和 c，因为 from_numpy() 会共享 NumPy 数组的存储；torch.tensor() 则创建没有原始自动微分历史的新 Tensor，并复制输入数据。(docs.pytorch.org)

可以验证共享关系：

array[0] = 99.0

print(a)  # 通常仍为 tensor([1., 2., 3.])
print(b)  # tensor([99.,  2.,  3.])
print(c)  # 通常也会看到 99.

工程中最常见的错误之一是把一个已经是 Tensor 的对象写成：

x = torch.tensor(existing_tensor)

这会复制数据并切断原有 Autograd 历史。若目的是复制，应明确写出意图：

x_copy = existing_tensor.detach().clone()

若目的是保留计算历史，则不要重新构造；若只是转换 dtype 或 device，应使用：

x = existing_tensor.to(dtype=torch.float32, device="cuda")

二、形状、广播和设备：Tensor 运算的前置条件

1. 矩阵乘法不是逐元素乘法

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

c = a @ b
print(c.shape)  # torch.Size([2, 4])

矩阵乘法要求 a 的最后一个维度与 b 的倒数第二个维度相等，即这里的 3 必须匹配 3。

逐元素乘法则要求形状相同，或满足广播规则：

x = torch.ones(2, 3)
scale = torch.tensor([1.0, 2.0, 3.0])

print(x * scale)

scale 的形状为 (3,)，会沿第一个维度广播为 (2, 3)。

广播不是复制数据本身，而是让运算在逻辑上扩展维度。因此下面的代码成立：

x = torch.randn(8, 3)
bias = torch.randn(3)

y = x + bias

但下面通常不成立：

x = torch.randn(8, 3)
bias = torch.randn(8)

y = x + bias

(8, 3) 与 (8,) 从末尾对齐后比较的是 3 和 8，二者既不相等，也没有一个是 1。

如果 bias 表示每个样本一个标量，应显式写成 (8, 1)：

bias = torch.randn(8, 1)
y = x + bias

2. dtype 和 device 必须满足算子要求

模型参数通常是浮点 Tensor，分类标签通常是整数索引：

features = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
labels = torch.tensor([0, 1], dtype=torch.int64)

使用 nn.CrossEntropyLoss 时，标签通常应为 torch.long，并且取值范围是 [0, 类别数)。把标签误写成浮点数，或把 one-hot Tensor 直接传给需要类别索引的损失函数，都会导致错误或语义不符。

设备也必须统一：

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

x = x.to(device)
model = model.to(device)

不能让模型参数在 GPU、输入在 CPU：

Expected all tensors to be on the same device

model.to(device) 会移动参数和持久化 buffer，但不会自动移动任意保存在模块普通属性中的 Tensor。若一个 Tensor 是模型状态的一部分，应注册为 buffer：

class Normalizer(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.register_buffer("mean", torch.zeros(3))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x - self.mean

这样 model.to(device) 时，mean 会随模型移动，并进入 state_dict()。

三、Autograd：从计算图到梯度

Autograd 是 PyTorch 的自动微分系统。它不是“自动修改参数”，而是记录计算并计算导数；真正修改参数的是优化器。

1. 一个完整的标量推导

考虑：

y=2x+1y = 2x + 1

再定义损失：

L=y2L = y^2

根据链式法则：

∂L∂y=2y\frac{\partial L}{\partial y}=2y

∂y∂x=2\frac{\partial y}{\partial x}=2

因此：

∂L∂x=∂L∂y∂y∂x=2y⋅2=4y\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial y}{\partial x} = 2y \cdot 2 = 4y

当 x=3 时：

y=7,L=49,∂L∂x=28y=7,\qquad L=49,\qquad \frac{\partial L}{\partial x}=28

PyTorch 代码：

import torch

x = torch.tensor(3.0, requires_grad=True)
y = 2 * x + 1
loss = y.square()

print("x:", x.item())
print("y:", y.item())
print("loss:", loss.item())
print("y.grad_fn:", y.grad_fn)
print("loss.grad_fn:", loss.grad_fn)

loss.backward()

print("x.grad:", x.grad)

输出的关键结果是：

x: 3.0
y: 7.0
loss: 49.0
x.grad: tensor(28.)

当一个 Tensor 的 requires_grad=True 时，涉及它的运算会被 Autograd 记录。前向计算期间，输出通常带有 grad_fn，它指向用于反向传播的节点。PyTorch 使用反向自动微分：前向构建实际执行过的计算图，反向从输出沿图向输入应用链式法则。这个图默认在每次迭代后重新创建，因此 Python 条件分支可以改变每次迭代的计算路径。(docs.pytorch.org)

2. leaf Tensor、非 leaf Tensor 和 .grad

x 是 leaf Tensor，因为它由用户直接创建并要求梯度：

x = torch.tensor(3.0, requires_grad=True)
y = x * 2

y 是由运算产生的非 leaf Tensor。默认情况下，反向传播后主要保留 leaf Tensor 的 .grad：

loss = (x * 2).square()
loss.backward()

print(x.grad)  # 有值
# print(y.grad)  # 默认通常为 None

如果确实需要查看中间 Tensor 的梯度，需要在反向传播前调用：

x = torch.tensor(3.0, requires_grad=True)
y = x * 2
y.retain_grad()

loss = y.square()
loss.backward()

print(y.grad)

3. 梯度会累积

这是训练循环中必须理解的状态变化：

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

(x * 3).backward()
print(x.grad)  # tensor(3.)

(x * 4).backward()
print(x.grad)  # tensor(7.)

第二次 backward() 不会覆盖第一次梯度，而是累加：

3+4=73+4=7

因此每个训练批次通常需要：

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

现代代码也常见：

optimizer.zero_grad(set_to_none=True)

这会把梯度设为 None，而不是写入全零 Tensor，通常可以减少不必要的内存写入。若某个参数本轮没有参与计算，None 也能明确表示“没有梯度”，而不是“梯度恰好为零”。

4. backward() 的标量条件

对标量损失直接调用：

loss.backward()

等价于对该标量求导。

如果输出是向量，例如：

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x.square()

y 不是标量，直接执行 y.backward() 通常会报错，因为系统不知道应该对哪个标量目标求导。可以显式指定外部梯度：

y.backward(torch.ones_like(y))
print(x.grad)  # tensor([2., 4.])

这实际上计算的是：

∂(y1+y2)∂x\frac{\partial (y_1+y_2)}{\partial x}

也可以先聚合：

x = torch.tensor([1.0, 2.0], requires_grad=True)
loss = x.square().sum()
loss.backward()

5. detach()、no_grad() 和 inference_mode()

推理阶段通常不需要构建反向图：

model.eval()

with torch.no_grad():
    prediction = model(inputs)

model.eval() 与 torch.no_grad() 是两件不同的事：

  • model.eval() 改变模块行为，例如 Dropout 关闭随机丢弃，BatchNorm 使用运行统计量；
  • torch.no_grad() 禁止记录梯度，减少计算图和内存开销。

只调用其中一个都可能不完整：

model.eval()       # 仍然可能记录梯度
with torch.no_grad():
    model(inputs)  # Dropout 仍可能处于训练行为

detach() 则是从某个 Tensor 的计算历史中切断：

features = encoder(images)
features_for_numpy = features.detach().cpu().numpy()

直接调用 features.numpy() 在需要梯度或位于 GPU 时通常不合法；正确路径通常是 detach().cpu().numpy()。

torch.inference_mode() 是更强的推理上下文，适合明确不会参与 Autograd 的推理路径：

model.eval()

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

但如果后续代码需要对这些 Tensor 建立梯度关系，就不应使用它。

四、从参数到模型：nn.Module 的状态

PyTorch 模型通常继承 torch.nn.Module：

import torch
from torch import nn

class LinearModel(nn.Module):
    def __init__(self, input_dim: int) -> None:
        super().__init__()
        self.linear = nn.Linear(input_dim, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear(x)

nn.Linear(input_dim, 1) 内部包含两个可学习参数：

  • 权重 weight，形状为 (1, input_dim)；
  • 偏置 bias，形状为 (1,)。

它实现：

y^=xWT+b\hat{y}=xW^T+b

调用：

model = LinearModel(2)

for name, parameter in model.named_parameters():
    print(name, parameter.shape, parameter.requires_grad)

可以看到：

linear.weight torch.Size([1, 2]) True
linear.bias torch.Size([1]) True

Module 会递归注册子模块、参数和 buffer。只有被注册的参数才会被 model.parameters() 找到，只有被注册的子模块才会随着 .to(device)、.train()、.eval() 一起改变。

下面的写法有问题：

class BadModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layers = [nn.Linear(2, 4), nn.ReLU(), nn.Linear(4, 1)]

普通 Python 列表不会自动注册其中的子模块。应该使用：

class GoodModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(2, 4),
            nn.ReLU(),
            nn.Linear(4, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(x)

或者使用 nn.ModuleList：

self.layers = nn.ModuleList([
    nn.Linear(2, 4),
    nn.ReLU(),
    nn.Linear(4, 1),
])

五、Dataset 和 DataLoader：把数据变成批次

1. Dataset 的契约

Dataset 表达一个可索引的数据集。最基本的 map-style Dataset 需要实现：

__len__()
__getitem__(index)

含义是：

  • len(dataset) 返回样本总数；
  • dataset[i] 返回第 i 个样本。

例如构造一个线性回归数据集：

from torch.utils.data import Dataset
import torch

class ToyRegressionDataset(Dataset[tuple[torch.Tensor, torch.Tensor]]):
    def __init__(self, n_samples: int = 256) -> None:
        generator = torch.Generator().manual_seed(42)

        self.x = torch.randn(n_samples, 2, generator=generator)
        noise = 0.1 * torch.randn(n_samples, 1, generator=generator)

        # 真实关系：y = 3*x0 - 2*x1 + 1 + noise
        self.y = 3 * self.x[:, 0:1] - 2 * self.x[:, 1:2] + 1 + noise

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

    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
        return self.x[index], self.y[index]

这里每个样本是：

x: shape=(2,)
y: shape=(1,)

当 DataLoader 组装成批次后，形状变为：

x_batch: shape=(batch_size, 2)
y_batch: shape=(batch_size, 1)

这是因为默认的 collation 会沿新维度堆叠同类型样本。

2. DataLoader 的职责

from torch.utils.data import DataLoader

dataset = ToyRegressionDataset(10)
loader = DataLoader(
    dataset,
    batch_size=4,
    shuffle=True,
    num_workers=0,
)

for x_batch, y_batch in loader:
    print(x_batch.shape, y_batch.shape)

可能输出：

torch.Size([4, 2]) torch.Size([4, 1])
torch.Size([4, 2]) torch.Size([4, 1])
torch.Size([2, 2]) torch.Size([2, 1])

batch_size=4 并不保证最后一个批次也有 4 个样本。若模型或损失函数严格要求固定批次，可以使用 drop_last=True，但这会丢弃最后不足一个完整批次的数据。

shuffle=True 只影响每轮迭代的索引顺序，不会改变样本本身。训练集通常需要打乱；验证集和测试集一般不需要。

3. 多进程读取的状态边界

当 num_workers>0 时，DataLoader 会使用 worker 进程读取数据。此时 Dataset 中的状态可能被复制到多个进程：

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    persistent_workers=True,
)

常见故障包括：

  • Dataset 在主进程可用，但 worker 启动方式下无法序列化；
  • worker 内部打开的文件句柄无法安全复用；
  • 随机增强在不同 worker 中使用相同随机状态；
  • Windows 和某些启动方式下缺少 if __name__ == "__main__": 保护；
  • worker 读取异常被包装后，堆栈信息不如 num_workers=0 直观。

诊断数据问题时，应先把 num_workers 降为 0，确认单进程读取正确，再逐步增加并发。

如果 GPU 训练中 CPU 到 GPU 的拷贝成为瓶颈，可以尝试：

loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=2,
    pin_memory=True,
)

随后：

x_batch = x_batch.to(device, non_blocking=True)
y_batch = y_batch.to(device, non_blocking=True)

pin_memory=True 只对 CPU Tensor 有意义；它并不把数据直接放到 GPU，也不保证一定加速。数据量小、worker 多或预处理很轻时，额外进程和内存开销可能反而更大。

4. 可迭代 Dataset 与 map-style Dataset

如果数据来自无限流、日志管道或无法随机索引的生成器，可以实现 IterableDataset。这类 Dataset 的关键不是 __getitem__，而是 __iter__：

from torch.utils.data import IterableDataset

class StreamDataset(IterableDataset[torch.Tensor]):
    def __iter__(self):
        for value in range(10):
            yield torch.tensor([float(value)])

使用多 worker 时，必须根据 worker ID 切分数据，否则每个 worker 可能重复读取完整数据流。对于有限、可索引的数据，map-style Dataset 更容易计算长度、划分训练集和验证集。

六、训练循环：四个状态变化

一次标准训练批次可以写成：

1. 清理旧梯度
2. 前向计算得到预测值
3. 计算损失
4. 反向计算梯度
5. 优化器更新参数

代码：

optimizer.zero_grad()
prediction = model(x_batch)
loss = loss_fn(prediction, y_batch)
loss.backward()
optimizer.step()

顺序不能随意交换。

1. 为什么要先清理梯度

参数的梯度默认累积。如果不执行 zero_grad()，第 k 批次的梯度大致变成：

gstored=g1+g2+⋯+gkg_{\text{stored}} = g_1+g_2+\cdots+g_k

这只有在你明确想做梯度累积时才是正确的。普通 mini-batch 训练需要每批次重新开始。

2. 损失函数和优化器

对于回归问题，均方误差为：

L=1N∑i=1N(y^i−yi)2L = \frac{1}{N} \sum_{i=1}^{N} (\hat{y}_i-y_i)^2

其中：

  • NN 是批次中的样本数；
  • y^i\hat{y}_i 是模型预测；
  • yiy_i 是目标值。

对应代码：

loss_fn = nn.MSELoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)

优化器根据参数梯度改变参数。以最简单的梯度下降为例：

θt+1=θt−η∇θL\theta_{t+1} = \theta_t-\eta\nabla_\theta L

其中：

  • θt\theta_t 是当前参数；
  • η\eta 是学习率；
  • ∇θL\nabla_\theta L 是损失对参数的梯度。

optimizer.step() 才是参数发生变化的地方。loss.backward() 只计算梯度，不更新模型。

3. train() 和 eval() 是模块状态切换

model.train()

通常用于训练阶段，会让 Dropout、BatchNorm 等模块处于训练行为。

model.eval()

通常用于验证和推理，会让这些模块处于评估行为。

验证代码：

model.eval()

total_loss = 0.0
total_count = 0

with torch.no_grad():
    for x_batch, y_batch in valid_loader:
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)

        prediction = model(x_batch)
        loss = loss_fn(prediction, y_batch)

        batch_size = x_batch.shape[0]
        total_loss += loss.item() * batch_size
        total_count += batch_size

valid_loss = total_loss / total_count

这里必须按样本数加权。错误写法是直接平均每个批次的 loss：

valid_loss = sum(batch_losses) / len(batch_losses)

如果最后一个批次比其他批次小，这种计算会让小批次与完整批次拥有相同权重。

4. 梯度累积的正确实现

如果显存不足，可以把多个 micro-batch 的梯度累积后再更新一次：

accumulation_steps = 4
optimizer.zero_grad(set_to_none=True)

for step, (x_batch, y_batch) in enumerate(loader):
    x_batch = x_batch.to(device)
    y_batch = y_batch.to(device)

    prediction = model(x_batch)
    loss = loss_fn(prediction, y_batch)

    # 缩放损失，使多批次累积后的梯度近似大批次平均梯度
    (loss / accumulation_steps).backward()

    should_step = (
        (step + 1) % accumulation_steps == 0
        or (step + 1) == len(loader)
    )

    if should_step:
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

如果不除以 accumulation_steps，梯度大小通常会变为原来的数倍，等价于改变了有效学习率。最后一个不足完整累积周期的批次还需要单独处理，否则其梯度可能永远不会执行 step()。

七、一个端到端可运行的训练示例

下面的程序使用前面定义的合成数据集训练一个线性模型。它包含：

  • Dataset；
  • DataLoader；
  • CPU/GPU 设备选择；
  • 训练与验证；
  • 最优模型保存；
  • 最终检查点保存；
  • 重新加载检查点并验证。
# train_regression.py
from __future__ import annotations

import random
from pathlib import Path

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


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


class ToyRegressionDataset(Dataset[tuple[torch.Tensor, torch.Tensor]]):
    def __init__(self, n_samples: int = 512) -> None:
        generator = torch.Generator().manual_seed(1234)

        self.x = torch.randn(n_samples, 2, generator=generator)
        noise = 0.1 * torch.randn(n_samples, 1, generator=generator)

        self.y = (
            3.0 * self.x[:, 0:1]
            - 2.0 * self.x[:, 1:2]
            + 1.0
            + noise
        )

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

    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
        return self.x[index], self.y[index]


class LinearModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = nn.Linear(2, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear(x)


def evaluate(
    model: nn.Module,
    loader: DataLoader,
    loss_fn: nn.Module,
    device: torch.device,
) -> float:
    model.eval()

    total_loss = 0.0
    total_count = 0

    with torch.no_grad():
        for x_batch, y_batch in loader:
            x_batch = x_batch.to(device)
            y_batch = y_batch.to(device)

            prediction = model(x_batch)
            loss = loss_fn(prediction, y_batch)

            batch_size = x_batch.shape[0]
            total_loss += loss.item() * batch_size
            total_count += batch_size

    return total_loss / total_count


def main() -> None:
    seed_everything(2026)

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

    dataset = ToyRegressionDataset(n_samples=512)

    train_size = 400
    valid_size = len(dataset) - train_size

    split_generator = torch.Generator().manual_seed(2026)
    train_dataset, valid_dataset = random_split(
        dataset,
        [train_size, valid_size],
        generator=split_generator,
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=32,
        shuffle=True,
        num_workers=0,
    )
    valid_loader = DataLoader(
        valid_dataset,
        batch_size=64,
        shuffle=False,
        num_workers=0,
    )

    model = LinearModel().to(device)
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=1e-2,
        weight_decay=1e-4,
    )

    output_dir = Path("artifacts")
    output_dir.mkdir(exist_ok=True)

    best_valid_loss = float("inf")
    best_path = output_dir / "best_model.pt"
    last_path = output_dir / "last_checkpoint.pt"

    epochs = 50

    for epoch in range(1, epochs + 1):
        model.train()
        train_loss_sum = 0.0
        train_count = 0

        for x_batch, y_batch in train_loader:
            x_batch = x_batch.to(device)
            y_batch = y_batch.to(device)

            optimizer.zero_grad(set_to_none=True)

            prediction = model(x_batch)
            loss = loss_fn(prediction, y_batch)

            loss.backward()
            optimizer.step()

            batch_size = x_batch.shape[0]
            train_loss_sum += loss.item() * batch_size
            train_count += batch_size

        train_loss = train_loss_sum / train_count
        valid_loss = evaluate(
            model,
            valid_loader,
            loss_fn,
            device,
        )

        if valid_loss < best_valid_loss:
            best_valid_loss = valid_loss
            torch.save(
                {
                    "model_state": model.state_dict(),
                    "valid_loss": valid_loss,
                    "epoch": epoch,
                },
                best_path,
            )

        torch.save(
            {
                "epoch": epoch,
                "model_state": model.state_dict(),
                "optimizer_state": optimizer.state_dict(),
                "best_valid_loss": best_valid_loss,
            },
            last_path,
        )

        if epoch == 1 or epoch % 10 == 0:
            print(
                f"epoch={epoch:02d} "
                f"train_loss={train_loss:.6f} "
                f"valid_loss={valid_loss:.6f}"
            )

    checkpoint = torch.load(
        best_path,
        map_location=device,
    )

    restored_model = LinearModel().to(device)
    restored_model.load_state_dict(checkpoint["model_state"])

    restored_valid_loss = evaluate(
        restored_model,
        valid_loader,
        loss_fn,
        device,
    )

    print("best checkpoint:", best_path)
    print("restored valid loss:", restored_valid_loss)

    print("learned weight:", restored_model.linear.weight.detach().cpu())
    print("learned bias:", restored_model.linear.bias.detach().cpu())


if __name__ == "__main__":
    main()

运行：

python train_regression.py

预期现象是：

device: cpu
epoch=01 train_loss=...
epoch=10 train_loss=...
epoch=10 ... valid_loss=...
...
restored valid loss: ...
learned weight: tensor([[ 3..., -1...]])
learned bias: tensor([1...])

由于数据中加入了噪声，最终权重不会严格等于 [[3.0, -2.0]]，但通常会接近它。若训练损失下降而验证损失异常升高，可能是过拟合、数据划分差异、验证流程错误或数据泄漏。

八、检查点：保存什么，恢复什么

1. state_dict 是什么

模块的 state_dict() 是参数和持久化 buffer 的名称到 Tensor 的映射：

state = model.state_dict()

for name, value in state.items():
    print(name, value.shape)

例如线性模型可能得到：

linear.weight torch.Size([1, 2])
linear.bias torch.Size([1])

模块的 state_dict 不只包含可学习参数，也包含影响计算但不需要梯度更新的持久化 buffer，例如 BatchNorm 的运行均值和方差。官方文档将参数和 persistent buffers 都视为模块状态。(docs.pytorch.org)

保存模型参数：

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

加载：

model = LinearModel()
state = torch.load("model.pt", map_location="cpu")
model.load_state_dict(state)

load_state_dict() 会根据键名和形状匹配参数。常见错误包括：

Missing key(s) in state_dict
Unexpected key(s) in state_dict
size mismatch

这些错误通常说明模型结构、命名、包装方式或版本不一致。

2. 只保存模型参数不等于恢复训练

如果只保存：

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

可以恢复推理模型，但不能完整恢复训练过程。原因是优化器通常还有自己的状态，例如 AdamW 会保存参数的一阶、二阶统计量。如果丢失这些状态，继续训练时的更新轨迹会发生变化。

一个可恢复训练的检查点通常包括：

checkpoint = {
    "epoch": epoch,
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
    "scheduler_state": scheduler.state_dict(),
    "best_valid_loss": best_valid_loss,
    "config": {
        "batch_size": 32,
        "learning_rate": 1e-2,
    },
}

如果使用混合精度，还应保存 GradScaler 状态；如果使用分布式训练，还要考虑每个 rank 的状态和分片检查点。

3. 从检查点继续训练

checkpoint = torch.load(
    "artifacts/last_checkpoint.pt",
    map_location=device,
)

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

start_epoch = int(checkpoint["epoch"]) + 1
best_valid_loss = float(checkpoint["best_valid_loss"])

for epoch in range(start_epoch, epochs + 1):
    ...

加载优化器状态后，应确认它的参数确实对应当前模型。最安全的顺序是：

model = LinearModel().to(device)

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-2,
)

checkpoint = torch.load(
    "artifacts/last_checkpoint.pt",
    map_location=device,
)

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

如果先创建了不同参数集合的优化器，再加载状态，可能出现参数组数量或形状不匹配。

4. 保存和加载的安全边界

PyTorch 的常规序列化格式依赖 Python pickle 机制。不要加载来源不可信的 .pt 或 .pth 文件，因为反序列化可能执行任意代码。检查点应来自可信来源，并在加载前校验文件来源、哈希和权限。

对于只包含模型权重的场景，可以根据所使用 PyTorch 版本支持情况考虑安全加载选项，例如：

state = torch.load(
    "model.pt",
    map_location="cpu",
    weights_only=True,
)

但 weights_only 的行为和默认值属于版本敏感内容，应以当前安装版本的 torch.load 文档和实际测试为准，不能把一个版本的默认行为假定为所有版本都相同。

九、检查点的原子写入和恢复验证

训练过程中直接覆盖正式文件存在风险：

torch.save(checkpoint, "last_checkpoint.pt")

如果进程在写文件过程中被杀死，可能留下损坏文件。更稳妥的方式是先写临时文件，再原子替换：

from pathlib import Path
import os
import torch


def atomic_torch_save(
    state: dict,
    path: Path,
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)

    temporary_path = path.with_suffix(path.suffix + ".tmp")
    torch.save(state, temporary_path)

    # 同一文件系统内的 replace 通常是原子替换
    os.replace(temporary_path, path)

保存后还应验证：

checkpoint = {
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
    "epoch": epoch,
}

atomic_torch_save(checkpoint, Path("artifacts/last_checkpoint.pt"))

test = torch.load(
    "artifacts/last_checkpoint.pt",
    map_location="cpu",
)

required_keys = {"model_state", "optimizer_state", "epoch"}
missing = required_keys - test.keys()

if missing:
    raise RuntimeError(f"checkpoint is incomplete: {missing}")

这只能验证文件可读取和键存在，不能保证训练语义正确。更进一步，可以：

  1. 新建同结构模型；
  2. 加载模型状态；
  3. 在固定样本上计算预测；
  4. 与保存前的预测比较；
  5. 检查参数形状和 dtype；
  6. 若要续训，再检查优化器状态和 epoch 是否一致。

十、随机性、可复现性和检查点边界

训练结果受到多个随机源影响：

  • Python random；
  • NumPy 随机数；
  • PyTorch CPU 随机数；
  • CUDA 随机数；
  • DataLoader worker 随机数；
  • 数据增强；
  • 某些 GPU 算子的非确定性实现；
  • 并发执行顺序。

最小的种子设置：

import random
import numpy as np
import torch

seed = 2026

random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)

if torch.cuda.is_available():
    torch.cuda.manual_seed_all(seed)

但“设置随机种子”不等于完全复现。若需要严格复现，还要处理确定性算法、worker 初始化、数据顺序、软件版本、硬件和编译配置。确定性设置可能降低性能，某些算子也可能没有确定性实现。

检查点如果只保存模型和优化器状态，通常无法保证从中断位置逐样本复现。更完整的恢复状态还包括：

rng_state = {
    "python": random.getstate(),
    "numpy": np.random.get_state(),
    "torch_cpu": torch.get_rng_state(),
}

if torch.cuda.is_available():
    rng_state["torch_cuda"] = torch.cuda.get_rng_state_all()

恢复：

random.setstate(checkpoint["rng_state"]["python"])
np.random.set_state(checkpoint["rng_state"]["numpy"])
torch.set_rng_state(checkpoint["rng_state"]["torch_cpu"])

if torch.cuda.is_available() and "torch_cuda" in checkpoint["rng_state"]:
    torch.cuda.set_rng_state_all(checkpoint["rng_state"]["torch_cuda"])

即使这样，也要注意 DataLoader 的索引生成器、worker 调度和已经消费过的数据位置。对于普通工程，检查点的目标通常是“从相近训练状态继续优化”，而不是保证每一个后续 batch 与原运行完全一致。

十一、混合精度和梯度裁剪

1. 混合精度的因果关系

混合精度使用较低精度执行部分计算，以减少显存和提高吞吐；但低精度可能使小梯度下溢。梯度缩放器通过先放大损失，使反向梯度暂时变大，再在参数更新前还原。

典型结构：

scaler = torch.amp.GradScaler("cuda")

for x_batch, y_batch in train_loader:
    x_batch = x_batch.to(device)
    y_batch = y_batch.to(device)

    optimizer.zero_grad(set_to_none=True)

    with torch.autocast(
        device_type=device.type,
        dtype=torch.float16,
        enabled=device.type == "cuda",
    ):
        prediction = model(x_batch)
        loss = loss_fn(prediction, y_batch)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

具体 dtype 和设备支持取决于 PyTorch、硬件和算子。不能假定所有设备都适合 float16；某些 CPU 场景更适合 bfloat16，应根据当前版本文档和硬件测试。

若需要梯度裁剪，必须先把缩放梯度还原：

scaler.scale(loss).backward()

scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(
    model.parameters(),
    max_norm=1.0,
)

scaler.step(optimizer)
scaler.update()

直接对尚未 unscale 的梯度裁剪，会把缩放因子误认为真实梯度大小。

2. 为什么需要梯度裁剪

当梯度范数过大时，单次更新可能把参数推到数值不稳定区域。梯度裁剪限制更新前的梯度范数：

g′=g⋅min⁡(1,c∥g∥)g' = g \cdot \min\left(1,\frac{c}{\|g\|}\right)

其中 cc 是 max_norm。它不能修复错误的损失、学习率或数据归一化，但可以缓解部分梯度爆炸问题。

诊断时可以记录：

total_norm = torch.nn.utils.clip_grad_norm_(
    model.parameters(),
    max_norm=1.0,
)

print("gradient norm before clipping:", float(total_norm))

如果梯度几乎每一步都被严重裁剪，根因可能是学习率过高、损失尺度异常、输入未归一化或模型结构不稳定，而不是简单地把 max_norm 继续调小。

十二、训练失败时先看什么

1. Loss 是 nan 或 inf

常见原因：

  • 输入中已经包含 nan 或 inf；
  • 学习率过大；
  • log(0)、除以零或非法平方根；
  • 混合精度溢出；
  • 标签范围错误；
  • 损失函数输入形状不符合预期。

可以在训练前加入：

def check_finite(name: str, value: torch.Tensor) -> None:
    if not torch.isfinite(value).all():
        raise ValueError(f"{name} contains NaN or Inf")

check_finite("x_batch", x_batch)
check_finite("y_batch", y_batch)

也可以检查参数和梯度：

for name, parameter in model.named_parameters():
    if parameter.grad is not None:
        check_finite(f"gradient:{name}", parameter.grad)

2. Loss 不下降

不要先盲目增加 epoch。按数据流顺序检查：

  1. Dataset[0] 是否返回正确；
  2. 一个 batch 的形状和 dtype 是否正确；
  3. 标签是否与输入对齐；
  4. 模型输出形状是否与损失函数要求一致；
  5. loss.backward() 后参数是否存在梯度；
  6. optimizer.step() 前后参数是否真的变化；
  7. 学习率是否合理；
  8. 训练集上能否过拟合极小样本。

极小样本过拟合测试：

small_x = torch.randn(8, 2)
small_y = 3 * small_x[:, 0:1] - 2 * small_x[:, 1:2] + 1

for _ in range(1000):
    optimizer.zero_grad(set_to_none=True)
    prediction = model(small_x)
    loss = loss_fn(prediction, small_y)
    loss.backward()
    optimizer.step()

如果一个容量足够的模型连 8 个固定样本都无法拟合，问题通常在数据、损失、梯度或优化器连接，而不是泛化能力。

3. 训练准确率很高，验证结果很差

这可能是过拟合，但也可能是验证流程错误：

  • 忘记 model.eval()；
  • 验证时仍在执行随机数据增强；
  • 验证集与训练集预处理不一致；
  • 训练集和验证集发生重复；
  • 指标聚合方式错误；
  • 标签编码不一致。

对于含 BatchNorm 或 Dropout 的网络，验证阶段尤其要同时使用：

model.eval()

with torch.no_grad():
    ...

4. 恢复后结果不一致

首先区分两个目标：

  • 恢复推理：模型参数和必要 buffer 足够；
  • 恢复训练：还需要优化器、学习率调度器、混合精度缩放器、epoch、随机状态和数据进度等。

只加载模型权重后继续训练，仍然可以运行，但它不是从原训练状态严格续训。

十三、训练、验证和检查点的状态图

stateDiagram-v2
    [*] --> Initialized

    Initialized --> TrainMode: model.train()
    TrainMode --> BatchReady: DataLoader yields batch
    BatchReady --> Forward: model(x)
    Forward --> LossComputed: loss_fn(pred, y)
    LossComputed --> Backward: loss.backward()
    Backward --> Updated: optimizer.step()
    Updated --> BatchReady: next batch
    Updated --> Validate: epoch finished

    Validate --> EvalMode: model.eval()
    EvalMode --> Validated: no_grad + validation loop
    Validated --> SaveBest: valid_loss improves
    Validated --> SaveLast: every epoch
    SaveBest --> TrainMode: next epoch
    SaveLast --> TrainMode: next epoch

    SaveLast --> Interrupted: process failure
    Interrupted --> Restore: load checkpoint
    Restore --> TrainMode: restore model/optimizer/RNG state

关键路径是：

  • model.train() 和 model.eval() 控制模块行为；
  • zero_grad() 控制梯度累积边界；
  • backward() 产生梯度；
  • step() 修改参数；
  • 验证阶段不能执行参数更新；
  • 检查点至少要区分“最好模型”和“最近训练状态”。

“最好模型”适合部署或最终评估；“最近检查点”适合故障恢复。二者通常不是同一个文件。

十四、与 NumPy 和 OpenAI SDK 的边界

1. NumPy 是数据处理工具，Tensor 是训练状态载体

NumPy ndarray 和 PyTorch Tensor 都支持多维数组、切片、广播和向量化运算，但训练代码还需要：

  • Autograd 计算图；
  • GPU/其他设备；
  • 模型参数注册；
  • 优化器状态；
  • 训练和推理模式；
  • 检查点状态。

因此常见边界是：

# NumPy → Tensor
x = torch.from_numpy(array).float()

# Tensor → NumPy
array = tensor.detach().cpu().numpy()

转换时要同时考虑：

  • 是否复制；
  • 是否共享存储；
  • dtype 是否匹配；
  • Tensor 是否在 CPU；
  • 是否仍然需要梯度。

不要在每个训练 batch 中无意义地来回转换，否则会破坏设备上的连续数据流。

2. OpenAI SDK 与 PyTorch 的职责不同

OpenAI Python SDK 是访问 OpenAI API 的客户端库，提供同步和异步客户端、类型定义，以及 Responses API 等接口；它不是 PyTorch 的数据集、自动微分或训练框架。官方仓库当前说明其面向 Python 3.10+，并展示了 client.responses.create(...) 的基本用法。(github.com)

如果一个 PyTorch 项目需要把模型输出发送到外部 API，应把这部分放在推理或应用层，而不是放进需要高吞吐、可复现的核心训练循环：

import os
from openai import OpenAI

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

response = client.responses.create(
    model="your-configured-model",
    input="请解释这个模型输出的含义。",
)

print(response.output_text)

训练过程中调用网络 API 会引入新的故障路径：

  • 网络超时；
  • 限流；
  • 请求重试导致重复副作用；
  • API 响应延迟造成 GPU 空转；
  • 结果随服务端模型版本或配置变化；
  • API 密钥泄露；
  • 无法通过本地检查点完全恢复外部请求状态。

因此，若 API 参与数据标注或增强，通常应先离线生成并缓存结果，再由 Dataset 读取缓存。若必须在线调用，则应设计请求 ID、幂等缓存、超时、重试和失败样本队列，并明确记录外部服务配置。

OpenAI SDK 还支持同步和异步 Responses 流式接口，但异步网络并发不能替代 DataLoader 的 worker 并发；前者解决 API I/O，后者解决训练数据准备，二者的生命周期和故障恢复模型不同。(github.com)

十五、工程上应固定下来的最小约束

一个可维护的 PyTorch 工程，至少应让以下信息可观察：

print({
    "torch_version": torch.__version__,
    "device": str(device),
    "dtype": str(next(model.parameters()).dtype),
    "num_parameters": sum(
        p.numel() for p in model.parameters()
    ),
})

每个 batch 或 epoch 都应能确认：

  • 输入和标签形状；
  • 输入和标签 dtype；
  • 当前 device；
  • loss 是否有限；
  • 当前学习率；
  • 训练/验证模式；
  • 当前 epoch 和 global step；
  • 检查点路径及其验证结果。

模型、优化器和数据处理配置应一起版本化。仅保存一个 model.pt，却不记录输入归一化方式、类别映射、模型构造参数和 PyTorch 版本，往往无法可靠部署。

最终可以把 PyTorch 训练理解为一组明确的状态转换：

Tensor 状态正确
→ Dataset 能返回单样本
→ DataLoader 能组成批次
→ Module 能在同一 device 上前向
→ Loss 是有限标量
→ Autograd 产生正确梯度
→ Optimizer 更新参数
→ Eval 流程正确计算指标
→ Checkpoint 能恢复需要的状态

任何一个环节的形状、dtype、device 或生命周期不一致，后面的训练结果都可能失去意义。把这些边界逐一验证，比单纯增加模型规模或训练轮数更能提高 PyTorch 工程的可靠性。

系列导航与关联阅读

官方资料

本文依据 Python 官方文档、相关 PEP 与生态项目官方文档重新梳理；正文、示例与工程清单由 WR BLOG 编写。