PyTorch模型训练

有了张量和自动求导，训练一个模型就是反复做四件事：前向计算得到预测、用损失函数衡量误差、反向传播求梯度、用优化器更新参数。PyTorch 为每一步都提供了标准组件。

定义模型：nn.Module

所有模型都继承 nn.Module：在 __init__ 里声明层，在 forward 里写前向计算。

import torch
from torch import nn
 
class MLP(nn.Module):
    def __init__(self, in_dim=2, hidden=64, n_classes=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, n_classes),
        )
 
    def forward(self, x):
        return self.net(x)   # 输出 logits，不做 softmax
 
model = MLP()
model(torch.randn(5, 2)).shape      # torch.Size([5, 3])：直接调用 model(x)，不要调用 model.forward(x)

在 __init__ 里赋值给 self 的层会被自动注册，它们的参数可以统一取出：

for name, p in model.named_parameters():
    print(name, tuple(p.shape))
# net.0.weight (64, 2)
# net.0.bias (64,)
# net.2.weight (64, 64)
# ...
sum(p.numel() for p in model.parameters())   # 4547：参数总量

「一个 7B 模型」说的就是这个数字约为 70 亿。每一层的含义见神经网络与反向传播和激活函数。

准备数据：Dataset 与 DataLoader

• Dataset 负责「第 i 个样本是什么」，实现 __len__ 和 __getitem__ 两个方法。
• DataLoader 负责把样本组成批次、打乱顺序、多进程加载。

from torch.utils.data import Dataset, DataLoader
 
class MyDataset(Dataset):
    def __init__(self, x, y):
        self.x, self.y = x, y
 
    def __len__(self):
        return len(self.x)
 
    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]
 
x = torch.randn(128, 2)
y = torch.randint(0, 3, (128,))
loader = DataLoader(MyDataset(x, y), batch_size=64, shuffle=True)
for xb, yb in loader:     # 每次取出一个批次，第 0 维是 batch
    ...

训练集要 shuffle=True，验证集不需要打乱。

损失函数与优化器

分类任务最常用的是交叉熵损失 nn.CrossEntropyLoss。它的输入是未经 softmax 的 logits 和整数类别标签，内部会自己做 log-softmax，所以模型最后一层不要再加 softmax。

logits = torch.tensor([[2.0, 0.5, -1.0]])
target = torch.tensor([0])
nn.CrossEntropyLoss()(logits, target)        # tensor(0.2413)
-torch.log_softmax(logits, dim=-1)[0, 0]     # 同样是 0.2413

语言模型预测下一个 token，本质上也是在整个词表上做分类，用的也是交叉熵，见损失函数和语言模型。

优化器负责根据梯度更新参数。当前训练 Transformer 的默认选择是 AdamW：

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)

优化器的原理见梯度下降。

标准训练循环

每个批次都是固定的五步：

model.train()                    # 切换到训练模式
for x, y in train_loader:
    x, y = x.to(device), y.to(device)
    logits = model(x)            # 1. 前向计算
    loss = loss_fn(logits, y)    # 2. 计算损失
 
    optimizer.zero_grad()        # 3. 清空上一步的梯度（梯度默认会累加）
    loss.backward()              # 4. 反向传播，求出每个参数的梯度
    optimizer.step()             # 5. 按梯度更新参数

zero_grad() 忘写是最常见的错误之一：梯度会一直累加，训练很快发散。

评估模式

model.eval()                     # 切换到评估模式
with torch.no_grad():            # 不记录计算图，省内存、提速
    logits = model(x_val)

model.train() 和 model.eval() 改变的是 Dropout、BatchNorm 这类层的行为：Dropout 在训练时随机丢弃、评估时不丢弃；BatchNorm 在训练时用当前批次的统计量、评估时用训练中累积的统计量。它们不会关闭梯度计算，所以评估时还要配合 torch.no_grad()。相关概念见过拟合与正则化和归一化。

学习率调度

学习率通常不是固定的，而是随训练进度变化。PyTorch 提供了 torch.optim.lr_scheduler，每个 epoch 或每一步调用一次 scheduler.step()：

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)
for epoch in range(num_epochs):
    train_one_epoch()
    scheduler.step()

大模型训练常用「先线性预热（warmup），再余弦衰减」，可以用 LambdaLR 自己写：

import math
from torch.optim.lr_scheduler import LambdaLR
 
warmup, total = 10, 100          # 预热步数、总步数
 
def lr_lambda(step):             # 返回值是基础学习率的倍数
    if step < warmup:
        return (step + 1) / warmup
    progress = (step - warmup) / (total - warmup)
    return 0.5 * (1 + math.cos(math.pi * progress))
 
scheduler = LambdaLR(optimizer, lr_lambda)
# 每个训练步：optimizer.step(); scheduler.step()

为什么要这样调度，见优化器、学习率调度与混合精度。

保存与加载

推荐只保存参数字典 state_dict，而不是整个模型对象：

torch.save(model.state_dict(), "mlp.pt")
 
model = MLP().to(device)                              # 模型和输入要放在相同设备
model.load_state_dict(torch.load("mlp.pt", map_location=device, weights_only=True))

如果要从中断处继续训练，还要一起保存优化器和调度器的 state_dict，以及当前的步数。

完整例子

在合成的三类螺旋数据上训练一个 MLP。三类点在平面上互相缠绕，线性模型分不开，需要多层非线性网络。这段代码可以直接运行，CPU 上几秒钟就能跑完。

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader, random_split
 
torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
class SpiralDataset(Dataset):
    """二维平面上的三类螺旋点。"""
 
    def __init__(self, n_per_class=300, n_classes=3):
        xs, ys = [], []
        for c in range(n_classes):
            r = torch.linspace(0.0, 1.0, n_per_class)
            theta = torch.linspace(c * 4.0, (c + 1) * 4.0, n_per_class) + 0.2 * torch.randn(n_per_class)
            xs.append(torch.stack([r * torch.sin(theta), r * torch.cos(theta)], dim=1))
            ys.append(torch.full((n_per_class,), c))
        self.x = torch.cat(xs)
        self.y = torch.cat(ys)
 
    def __len__(self):
        return len(self.x)
 
    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]
 
 
class MLP(nn.Module):
    def __init__(self, in_dim=2, hidden=64, n_classes=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, n_classes),
        )
 
    def forward(self, x):
        return self.net(x)
 
 
dataset = SpiralDataset()
train_set, val_set = random_split(dataset, [720, 180])
train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
val_loader = DataLoader(val_set, batch_size=256)
 
model = MLP().to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2, weight_decay=1e-4)
num_epochs = 100
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)
 
 
def evaluate(model, loader):
    model.eval()
    correct, total, loss_sum = 0, 0, 0.0
    with torch.no_grad():
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            logits = model(x)
            loss_sum += loss_fn(logits, y).item() * len(y)
            correct += (logits.argmax(dim=-1) == y).sum().item()
            total += len(y)
    return loss_sum / total, correct / total
 
 
for epoch in range(num_epochs):
    model.train()
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        loss = loss_fn(model(x), y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    scheduler.step()
 
    if (epoch + 1) % 20 == 0:
        val_loss, val_acc = evaluate(model, val_loader)
        print(f"epoch {epoch + 1:3d}  lr {scheduler.get_last_lr()[0]:.5f}  "
              f"val_loss {val_loss:.4f}  val_acc {val_acc:.3f}")
 
torch.save(model.state_dict(), "mlp.pt")

在 CPU 上的一次运行输出如下，验证集准确率约 99%：

epoch  20  lr 0.00905  val_loss 0.0429  val_acc 0.983
epoch  40  lr 0.00655  val_loss 0.0287  val_acc 0.994
epoch  60  lr 0.00345  val_loss 0.0206  val_acc 0.994
epoch  80  lr 0.00095  val_loss 0.0189  val_acc 0.994
epoch 100  lr 0.00000  val_loss 0.0175  val_acc 0.994

两个常用技巧

梯度累积：把有效大批次拆成几个微批次，累计梯度后更新一次。微批次等大时，可以把各自的平均损失除以累积步数；大小不等或最后一组不足时，要按实际样本数加权。下面针对前面的单标签分类例子，先累计损失之和，再按整组样本数归一化，同时处理尾批次。语言模型则通常应按参与损失的有效 token 数归一化。带 BatchNorm 等跨样本运算时，累积也未必等价于一次完整大批次前向。

梯度裁剪：把梯度的整体范数限制在一个上限内，防止偶发的梯度爆炸破坏训练，大模型训练几乎都会用。

accum_steps = 4
model.train()
optimizer.zero_grad(set_to_none=True)
group_samples = 0
for step, (x, y) in enumerate(train_loader):
    x, y = x.to(device), y.to(device)
    loss_sum = nn.functional.cross_entropy(model(x), y, reduction="sum")
    loss_sum.backward()
    group_samples += y.numel()
    if (step + 1) % accum_steps == 0 or step + 1 == len(train_loader):
        for p in model.parameters():
            if p.grad is not None:
                p.grad.div_(group_samples)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
        group_samples = 0

在 AI 中的作用

大模型的训练代码虽然复杂得多（混合精度、分布式并行、检查点等），但骨架仍然是这里的五步循环。Hugging Face 的 Trainer、TRL 的 SFTTrainer 等高层工具，本质上都是把这个循环封装了起来。

参考资料

• PyTorch：梯度累积与梯度裁剪

相关

• PyTorch张量与自动求导：训练循环的底层机制。
• 梯度下降：优化器背后的原理。
• 过拟合与正则化：为什么要划分验证集、使用权重衰减和 Dropout。
• Hugging Face生态：直接加载预训练模型，用封装好的训练器训练。
• 分布式训练：单卡训练循环如何扩展到多机多卡。