PyTorch 模型训练技术文档：求解器、参数配置与训练循环

文档索引

阅读建议：先读总览建立「数据 → 前向 → 损失 → 反向 → 优化器/调度器」的全局图景，再按需跳转优化器、调度器、参数组或配置经验；实现时按「构建 optimizer → 可选 scheduler → 循环内 zero_grad/backward/step/scheduler.step」顺序对照各章。

一、总览：整体架构与文档脉络

1.1 训练闭环在做什么

目标：在给定数据上，通过迭代更新模型参数 $\theta$，使损失 $L(\theta)$ 下降，从而拟合数据或任务目标。

最小闭环：每个 step 内依次执行：

1. 前向：输入 → 模型 → 输出
2. 损失：输出与目标代入损失函数 → 标量 loss
3. 反向：loss.backward() → 计算梯度 $\nabla_\theta L$
4. 更新：优化器根据梯度更新参数，如 $\theta \leftarrow \theta - \eta \nabla_\theta L$（SGD 为例）
5. 清空梯度：optimizer.zero_grad()，为下一轮 backward 做准备

其中「更新」由 优化器（Optimizer / 求解器） 完成；学习率 $\eta$ 可固定，也可由 学习率调度器（LRScheduler） 按 step 或 epoch 调整。因此：优化器决定「如何用梯度更新参数」；调度器决定「学习率随时间如何变化」。

1.2 整体架构与数据流

┌─────────────────────────────────────────────────────────────────────────────────┐
│                           训练准备（一次）                                         │
│  model = MyModel()                                                               │
│  optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)   ← 求解器            │
│  scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)     │
│  loss_fn = nn.CrossEntropyLoss()                                                 │
└─────────────────────────────────────────────────────────────────────────────────┘
                                        │
                                        ▼
┌─────────────────────────────────────────────────────────────────────────────────┐
│                           每个 step（或每 N 个 step 做一次 scheduler.step）       │
│  batch = next(dataloader)                                                        │
│  out = model(batch)          →  loss = loss_fn(out, target)  →  loss.backward()  │
│  optimizer.step()            →  optimizer.zero_grad()       →  [scheduler.step()] │
└─────────────────────────────────────────────────────────────────────────────────┘

1.3 知识结构（文档脉络）

• 基础层：优化器是什么、param_groups、state、step/zero_grad（第二节）；损失函数与训练循环的角色（第七、八节）。
• 算法层：SGD 族（第三节）、Adam/AdamW 等自适应算法（第四节）；参数组实现差异化 lr/weight_decay（第五节）。
• 调度层：LRScheduler 基类与 step 时机（第六节）；与 optimizer 的绑定关系。
• 工程层：梯度累积、梯度裁剪（第七节）；参数配置经验（第九节）；完整示例（第十节）。

各节关系：优化器 负责单步更新；调度器 在每 step 或每 epoch 后改 optimizer 的 lr；参数组 让同一 optimizer 内不同参数用不同 lr/weight_decay；训练循环 把 DataLoader、loss、optimizer、scheduler 串起来；配置经验 指导选 lr、warmup、weight_decay 等。

1.4 各部分职责与关联

1.5 优缺点与适用场景对比

1.5.1 SGD vs Adam / AdamW

AdamW 相对 Adam：weight decay 不进入动量与方差，等价于解耦的 L2 正则，通常比 Adam 更稳、泛化更好，推荐作默认。

1.5.2 学习率调度器选型

二、优化器概念与基类

2.1 优化器是什么、为什么需要

是什么：torch.optim.Optimizer 是持有「待优化参数」和「更新规则」的对象；每次调用 step() 时，根据当前梯度按该规则更新参数。

为什么需要：梯度只给出下降方向，步长（学习率）、动量、自适应缩放等若手写易错且难以复用；抽象成 Optimizer 后，算法统一、接口一致，且便于实现参数组、状态保存与恢复。

解决什么问题：把「梯度 → 参数更新」标准化，支持 SGD、Adam、AdamW 等多种算法，并支持 per-parameter 或 per-group 的 lr、weight_decay 等。

2.2 参数组（param_groups）

是什么：优化器内部把参数分成若干 参数组；每个组是一个 dict，至少包含 'params'（该组参数列表），以及可选的 'lr'、'weight_decay' 等，与构造时传入的 defaults 合并。

为什么需要：不同层或不同模块往往需要不同学习率（如 backbone 小 lr、head 大 lr），或 bias 不加 weight_decay；参数组允许在同一 optimizer 里为不同子集设置不同超参。

结构：optimizer.param_groups 是 list of dict，例如：

optimizer = torch.optim.SGD([
    {'params': model.base.parameters(), 'lr': 1e-2},
    {'params': model.head.parameters(), 'lr': 1e-1}
], momentum=0.9)
# param_groups[0]['params'] 为 base 参数，lr=1e-2
# param_groups[1]['params'] 为 head 参数，lr=1e-1

2.3 状态（state）

是什么：每个被优化参数在 optimizer 内有一个 id，对应 state[id] 的 dict，用于存该算法的内部状态（如 SGD 的 momentum buffer、Adam 的 exp_avg/exp_avg_sq）。

为什么需要：动量类、自适应类算法需要跨 step 记忆历史信息；state 在 step() 时被读入并写回。

用户通常不直接改 state；保存/恢复训练时用 optimizer.state_dict() 与 optimizer.load_state_dict() 即可。

2.4 基类方法完整说明

注意：zero_grad() 与 step() 的先后关系。常见写法是「先 zero_grad，再 forward/backward，再 step」；若做梯度累积，则多个 backward 后再一次 step，且每步开始 zero_grad。

2.5 示例：基类用法

import torch
import torch.nn as nn

model = nn.Linear(10, 2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 模拟一步
x = torch.randn(4, 10)
y = model(x).sum()
optimizer.zero_grad()
y.backward()
optimizer.step()

# 查看 param_groups
assert len(optimizer.param_groups) == 1
assert optimizer.param_groups[0]['lr'] == 0.01

# 保存/恢复（仅示例，实际会连同 model）
state = optimizer.state_dict()
optimizer.load_state_dict(state)

三、SGD 族

3.1 结构说明

SGD 族包括：SGD（可选 momentum）、Nesterov。本节只讨论 torch.optim.SGD 的完整参数与含义；Momentum 与 Nesterov 均通过同一类实现。

3.2 关键概念与公式

• SGD： $$\theta_{t+1} = \theta_t - \eta \nabla L(\theta_t)$$ 即用当前梯度乘以学习率更新。
• SGD with Momentum：先累积动量 $v_{t+1} = \mu v_t + \nabla L(\theta_t)$，再更新 $\theta_{t+1} = \theta_t - \eta v_{t+1}$，减轻震荡、加速收敛。
• Nesterov：在计算梯度前先做「临时更新」再算梯度，使动量更前瞻一步，公式上等价于对动量项做修正（PyTorch 中 nesterov=True 即启用）。

为什么需要 Momentum：纯 SGD 在病态曲面上震荡大、收敛慢；动量起到平滑梯度、利用历史方向的作用。

3.3 SGD 构造函数与参数

torch.optim.SGD(params, lr=<required>, momentum=0, dampening=0,
                weight_decay=0, nesterov=False, *, maximize=False, foreach=None, differentiable=False)

3.4 使用方式与示例

# 纯 SGD
opt = torch.optim.SGD(model.parameters(), lr=0.01)

# SGD + Momentum
opt = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# SGD + Nesterov
opt = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, nesterov=True)

# 带 weight_decay（L2 正则）
opt = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)

3.5 适用场景

• 大 batch、长训练、追求泛化时常用 SGD + Momentum（或 Nesterov）。
• 小 batch、快速实验更常用 Adam/AdamW；若从 Adam 换到 SGD，需要重新调 lr（通常 SGD 的 lr 比 Adam 大一个数量级左右）。

四、自适应学习率优化器

4.1 结构说明

本节覆盖：Adam、AdamW、RAdam、NAdam、RMSprop、Adagrad、Adadelta。它们共同点是按「每个参数」维护标量或向量状态，用梯度的一阶/二阶信息自适应地缩放步长。

4.2 Adam

是什么：对每个参数维护一阶矩估计 $m_t$ 和二阶矩估计 $v_t$，做偏差修正后用 $\hat{m}_t/(\sqrt{\hat{v}_t}+\epsilon)$ 方向更新，步长为 lr。

为什么需要：不同参数梯度量纲和曲率差异大；固定 lr 要么对某些参数过大要么过小；自适应步长能减少对 lr 的敏感度，收敛快、调参简单。

公式（简要）：

偏差修正：$\hat{m}_t = m_t/(1-\beta_1^t)$，$\hat{v}_t = v_t/(1-\beta_2^t)$。
更新：$\theta_{t+1} = \theta_t - \eta \cdot \hat{m}_t/(\sqrt{\hat{v}_t}+\epsilon)$。

构造函数：

torch.optim.Adam(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, *, amsgrad=False, foreach=None, maximize=False, capturable=False, differentiable=False, fused=None)

4.3 AdamW

是什么：Adam 的变体，把 weight decay 从「加在梯度里」改为「直接对参数做衰减」，即每步先做 $\theta \leftarrow \theta - \eta\lambda\theta$ 再按 Adam 更新。这样 weight decay 不参与一阶/二阶矩计算，等价于解耦的 L2 正则。

为什么需要：原 Adam 里 weight decay 与梯度耦合，正则效果和收敛行为不如解耦形式；AdamW 在多数任务上更稳、泛化更好，推荐作为默认选择。

构造函数：与 Adam 相同，仅实现不同；weight_decay 默认 0，常用 0.01。

torch.optim.AdamW(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01, ...)

4.3.1 Adam 与 AdamW 对比

• Adam：weight_decay 在实现上等价于 L2 正则（对梯度加 weight_decay·θ，再参与动量与二阶矩），与自适应步长耦合后，对泛化有时不利。
• AdamW：将 weight decay 从梯度更新中解耦，单独做 $\theta \leftarrow \theta - \eta\lambda\theta$，再按 Adam 公式更新；weight decay 不进入 $m_t/v_t$，更符合「参数衰减」语义，通常泛化更好，推荐默认使用 AdamW。

4.4 RAdam / NAdam

• RAdam：在训练早期二阶矩方差大时，用修正项避免步长过大，减少 warmup 需求。
• NAdam：把 Nesterov 思想融进 Adam，用「当前步的动量预测」更新，有时收敛略快。

两者接口与 Adam 类似，可作为替代尝试；默认仍推荐 AdamW。

4.5 RMSprop / Adagrad / Adadelta

• RMSprop：只维护二阶矩的指数平均，用 $\eta/\sqrt{v_t+\epsilon}$ 缩放梯度；适合 RNN、非平稳目标。
• Adagrad：累积平方梯度，步长单调降；适合稀疏梯度，但易过早变小。
• Adadelta：在 Adagrad 基础上用滑动平均替代累积，避免步长趋 0。

日常训练中 Adam/AdamW 使用更多；RMSprop 在部分 RL 或旧代码中仍见。

4.6 PyTorch 内置优化器一览（无遗漏）

以下为 torch.optim 中常见优化器及典型用途，便于按场景选择或查漏。

4.7 使用示例

# 默认推荐
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)

# 仅 Adam
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999))

五、参数组与差异化配置

5.1 结构说明

参数组即构造 optimizer 时传入的「list of dict」，每个 dict 一组参数 + 可选覆盖的 lr、weight_decay 等；优化器内部对所有组统一执行同一种算法，但每组用各自的超参。

5.2 常见用法

1）Backbone 小 lr、Head 大 lr（微调/迁移学习）

optimizer = torch.optim.AdamW([
    {'params': model.backbone.parameters(), 'lr': 1e-5},
    {'params': model.head.parameters(), 'lr': 1e-3}
], weight_decay=0.01)

2）Bias 不加 weight_decay

bias_params = [p for n, p in model.named_parameters() if 'bias' in n]
other_params = [p for n, p in model.named_parameters() if 'bias' not in n]
optimizer = torch.optim.AdamW([
    {'params': other_params, 'weight_decay': 0.01},
    {'params': bias_params, 'weight_decay': 0}
], lr=1e-3)

3）BatchNorm 等不 weight_decay（按需）

def get_param_groups(model, base_lr, wd):
    decay, no_decay = [], []
    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if 'bias' in n or 'norm' in n or 'bn' in n.lower():
            no_decay.append(p)
        else:
            decay.append(p)
    return [
        {'params': decay, 'lr': base_lr, 'weight_decay': wd},
        {'params': no_decay, 'lr': base_lr, 'weight_decay': 0}
    ]
optimizer = torch.optim.AdamW(get_param_groups(model, 1e-3, 0.01))

5.3 add_param_group

微调中动态加入新模块时，可把新参数加进已有 optimizer：

optimizer.add_param_group({'params': model.new_module.parameters(), 'lr': 1e-4})

注意：新组的其他键（如 weight_decay）若未写，会使用 optimizer 的 defaults。

六、学习率调度器

6.1 概念与基类

是什么：torch.optim.lr_scheduler.LRScheduler 绑定一个 optimizer，在每次 scheduler.step() 时根据策略更新 optimizer.param_groups[*]['lr']，从而在训练过程中改变学习率。

为什么需要：固定 lr 前期可能震荡、后期可能过大会影响收敛；按 step 或 epoch 衰减或先升后降，往往收敛更稳、最终效果更好。

调用顺序：必须先 optimizer.step()，再 scheduler.step()（若按 step 调）；若按 epoch 调，则每个 epoch 结束调一次 scheduler.step()。

6.2 基类方法

6.3 常用调度器一览

6.4 按 step 与按 epoch

• 按 step：每个 training step 后调用 scheduler.step()，如 OneCycleLR、部分 Cosine 实现。
• 按 epoch：每个 epoch 结束调用一次 scheduler.step()，如 StepLR、MultiStepLR、ReduceLROnPlateau（plateau 的 step 传 validation 指标）。

示例（按 epoch）：

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

示例（按 step，OneCycleLR）：

total = len(dataloader) * num_epochs
scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=1e-2, total_steps=total, pct_start=0.1
)
for batch in dataloader:
    ...
    optimizer.step()
    scheduler.step()

6.5 学习率调度器完整列表（无遗漏）

以下为 torch.optim.lr_scheduler 中常见调度器及要点，便于速查与补全。

基类方法：step()（部分调度器需传入 metrics）、get_last_lr()、state_dict()、load_state_dict()；子类实现 get_lr() 返回当前各 param_group 的学习率列表。

七、训练循环与梯度

7.1 最小闭环

model.train()
for batch in dataloader:
    inputs, targets = batch
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()
    # 若按 step 调 lr：scheduler.step()

7.2 梯度累积

当有效 batch = N * 单次 batch 时，每 N 步才 optimizer.step() 一次，每步只 backward() 不 step()，梯度累加；最后一步前 zero_grad()，然后 backward() 再 step()。

accum_steps = 4
for i, batch in enumerate(dataloader):
    inputs, targets = batch
    if i % accum_steps == 0:
        optimizer.zero_grad()
    outputs = model(inputs)
    loss = loss_fn(outputs, targets) / accum_steps
    loss.backward()
    if (i + 1) % accum_steps == 0:
        optimizer.step()

7.3 梯度裁剪

防止梯度爆炸，在 backward() 之后、step() 之前对梯度做范数裁剪：

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
# 或按值裁剪
# torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)
optimizer.step()

常用于 RNN、Transformer、大 lr 或长序列训练。

7.4 与 DDP 的衔接

在 分布式训练 中，每 rank 独立维护自己的 optimizer；DDP 在 backward 时已做梯度 AllReduce，因此各 rank 上梯度一致，各自 optimizer.step() 后参数保持同步。无需对 optimizer 做额外集体通信。

八、损失函数选型

8.1 常见 Loss 与适用任务

8.2 使用方式

# 分类
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits, targets)

# 回归
loss_fn = nn.MSELoss()
loss = loss_fn(pred, target)

# 多任务：加权和
loss = ce_loss(logits_cls, y_cls) + 0.1 * mse_loss(pred_reg, y_reg)
loss.backward()

九、参数配置经验

9.1 学习率与 batch size

• 线性缩放：batch 扩大 k 倍时，lr 常线性扩大 k 倍（如 32→256，lr 1e-3→8e-3）；再大 batch 可用 sqrt 缩放或 warmup 更长。
• SGD：常用 0.01～0.1；大 batch 可试 0.1 * batch/32。
• Adam/AdamW：1e-3 通用；2e-3～3e-3 也常见；再大易不稳定。

9.2 Warmup

前若干 step 或 epoch 从 0 或小 lr 线性升到目标 lr，减轻初期大梯度导致的震荡。Transformer、大 batch 训练常用。

# 可用 LambdaLR 或 OneCycleLR 的 pct_start 实现
def warmup_lambda(step):
    if step < warmup_steps:
        return step / warmup_steps
    return 1.0
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, warmup_lambda)

9.3 Weight decay

• AdamW：0.01 常用；0.1 也有用。
• SGD：1e-4～1e-2；bias / BN 通常设为 0。

9.4 常见坑

• 忘记 zero_grad()：梯度会累积，等价于放大 lr。
• scheduler.step() 在 optimizer.step() 之前其实并不会出错，但大多数调度器（如 StepLR、CosineAnnealingLR）推荐在 optimizer.step() 之后调用，这样每 step 用的是当前 step 的 lr，和论文/官方实现一致。部分调度器（如 ReduceLROnPlateau）需根据文档操作。
• 恢复 checkpoint 时先恢复 scheduler 再恢复 optimizer，否则 scheduler 会覆盖加载进来的 lr。
• ReduceLROnPlateau 的 step(metric) 要传验证指标，且注意 mode（‘min’/‘max’）。

十、完整示例

下面是一段可独立运行的训练脚本（CPU 即可），串联 DataLoader、模型、损失、优化器、调度器与梯度裁剪。

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

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    batch_size = 32
    num_epochs = 10
    lr = 1e-3
    weight_decay = 0.01
    max_grad_norm = 1.0

    # 数据
    X = torch.randn(1000, 20)
    y = (X.sum(dim=1, keepdim=True) > 0).long().squeeze(1)
    dataset = TensorDataset(X, y)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)

    # 模型、损失、优化器、调度器
    model = nn.Sequential(
        nn.Linear(20, 64),
        nn.ReLU(),
        nn.Linear(64, 2),
    ).to(device)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)

    model.train()
    for epoch in range(num_epochs):
        total_loss = 0.0
        for step, (inputs, targets) in enumerate(dataloader):
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()
            logits = model(inputs)
            loss = loss_fn(logits, targets)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            optimizer.step()
            total_loss += loss.item()
        avg_loss = total_loss / len(dataloader)
        scheduler.step()
        current_lr = scheduler.get_last_lr()[0]
        print(f"Epoch {epoch+1}/{num_epochs}  loss={avg_loss:.4f}  lr={current_lr:.2e}")

if __name__ == "__main__":
    main()

将上述代码保存为 train_example.py，在项目根目录执行 python train_example.py 即可运行（无 GPU 时自动用 CPU）。

十一、速查与小结

11.1 优化器速查

11.2 调度器速查

11.3 与其它文档的衔接

• 数据流：数据来自 Dataset → DataLoader → 本训练的 for batch in dataloader。
• 多卡：本训练的 optimizer/scheduler 每 rank 一份，与 DDP 无冲突；checkpoint 保存/加载时需同时保存 optimizer、scheduler 的 state_dict。

以上为模型训练中求解器、参数配置与训练循环的完整技术解读；按文档索引可快速定位到优化器、调度器、参数组或配置经验等小节。