【PyTorch】优化器进阶：从基础使用到自定义与调优实战

1. PyTorch优化器基础操作指南第一次接触PyTorch优化器时我盯着optim.SGD的参数列表发呆了半小时——这些参数到底该怎么设置后来才发现掌握优化器的基本使用就像学骑自行车刚开始会摔倒几次但一旦掌握就能自由驰骋。1.1 优化器的创建与初始化所有PyTorch优化器的使用都遵循一个固定套路先准备待优化的参数再设置优化选项。下面这段代码展示了最基础的优化器创建方式import torch import torch.optim as optim # 假设我们有一个简单的全连接网络 class SimpleNet(torch.nn.Module): def __init__(self): super().__init__() self.fc torch.nn.Linear(10, 2) model SimpleNet() # 创建优化器的标准姿势 optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9)这里有几个关键点需要注意model.parameters()会返回模型中所有需要训练的参数lr学习率是优化器最重要的超参数之一momentum是SGD特有的参数可以加速收敛我曾经犯过一个典型错误忘记把模型参数传给优化器。结果训练时模型参数完全没更新loss居高不下调试了半天才发现问题。1.2 优化三步曲zero_grad、backward、step每个训练迭代中优化器的使用遵循固定的三步流程for inputs, labels in dataloader: # 第一步清空梯度 optimizer.zero_grad() # 前向传播计算损失 outputs model(inputs) loss criterion(outputs, labels) # 第二步反向传播 loss.backward() # 第三步更新参数 optimizer.step()zero_grad()这一步特别容易忘记。我早期训练模型时经常纳闷为什么我的模型完全不收敛后来发现是因为梯度在每一步都累积了导致参数更新方向混乱。记住每次参数更新前都要清空梯度1.3 优化器状态管理优化器内部会维护一些状态信息比如动量缓冲器等。我们可以通过state_dict()方法获取这些状态# 获取优化器当前状态 optimizer_state optimizer.state_dict() # 保存到文件通常和模型状态一起保存 torch.save({ model_state: model.state_dict(), optimizer_state: optimizer_state, }, checkpoint.pth) # 恢复训练时加载状态 checkpoint torch.load(checkpoint.pth) model.load_state_dict(checkpoint[model_state]) optimizer.load_state_dict(checkpoint[optimizer_state])在实际项目中我习惯在每个epoch结束时保存一次状态。这样当训练意外中断时可以从最近的检查点恢复避免从头开始训练。2. 高级优化技巧分层参数配置当网络结构变得复杂时对所有层使用相同的优化参数往往不是最佳选择。PyTorch允许我们为不同层设置不同的优化参数。2.1 不同层不同学习率假设我们有一个包含特征提取器和分类头的模型通常希望特征提取器部分学习得更慢一些optimizer optim.SGD([ {params: model.feature_extractor.parameters(), lr: 1e-4}, {params: model.classifier.parameters(), lr: 1e-3} ], momentum0.9)这种技术在迁移学习中特别有用。当使用预训练模型时我们通常会让底层特征提取器的学习率比顶层分类器小10倍左右这样既能微调特征提取器又不会破坏预训练学到的有用特征。2.2 差异化权重衰减权重衰减L2正则化也可以分层设置。例如我们可能希望对偏置项不使用权重衰减# 分离权重和偏置参数 params_with_wd [] params_without_wd [] for name, param in model.named_parameters(): if bias in name: params_without_wd.append(param) else: params_with_wd.append(param) optimizer optim.AdamW([ {params: params_with_wd, weight_decay: 0.01}, {params: params_without_wd, weight_decay: 0} ], lr1e-3)这种设置方式在实践中很常见因为偏置项通常不需要正则化。我第一次尝试这种技巧是在训练一个大型Transformer模型时发现它能显著提升模型性能。2.3 参数冻结与解冻有时候我们希望在训练的不同阶段冻结或解冻某些层的参数# 初始阶段冻结特征提取器 for param in model.feature_extractor.parameters(): param.requires_grad False # 只优化分类器 optimizer optim.Adam(model.classifier.parameters(), lr1e-3) # 训练几个epoch后解冻特征提取器 for param in model.feature_extractor.parameters(): param.requires_grad True # 现在优化所有参数 optimizer.add_param_group({params: model.feature_extractor.parameters(), lr: 1e-5})这种渐进式解冻策略在微调大型模型时非常有效。我在处理医学图像分类任务时使用这种方法使模型准确率提升了约3个百分点。3. 学习率动态调整策略学习率是影响模型训练最关键的超参数之一。PyTorch提供了多种学习率调度器可以帮助我们在训练过程中动态调整学习率。3.1 内置学习率调度器PyTorch的torch.optim.lr_scheduler模块提供了多种调度器from torch.optim.lr_scheduler import StepLR, ExponentialLR, ReduceLROnPlateau # 创建优化器 optimizer optim.SGD(model.parameters(), lr0.1) # 每30个epoch将学习率乘以0.1 scheduler1 StepLR(optimizer, step_size30, gamma0.1) # 指数衰减学习率 scheduler2 ExponentialLR(optimizer, gamma0.95) # 当验证损失不再下降时降低学习率 scheduler3 ReduceLROnPlateau(optimizer, modemin, factor0.1, patience5) # 使用示例 for epoch in range(100): train(...) val_loss validate(...) # 不同调度器的调用方式不同 scheduler1.step() # 按epoch数调整 scheduler3.step(val_loss) # 按验证损失调整我个人的经验法则是对于简单任务使用StepLR对于复杂任务使用ReduceLROnPlateau。后者虽然计算开销稍大但通常能获得更好的最终性能。3.2 自定义学习率调度有时候内置调度器不能满足需求我们可以自定义学习率变化def custom_lr_scheduler(optimizer, epoch): 自定义学习率调度函数 if epoch 10: lr 1e-3 elif 10 epoch 20: lr 5e-4 else: lr 1e-4 for param_group in optimizer.param_groups: param_group[lr] lr # 在训练循环中使用 for epoch in range(100): custom_lr_scheduler(optimizer, epoch) train(...)我曾经在一个图像分割任务中使用类似的热启动策略先使用较大学习率快速收敛再逐步降低学习率精细调整最终比固定学习率提升了约2%的IoU。3.3 学习率预热技巧对于某些优化器特别是Adam在训练初期使用学习率预热warmup可以带来更好的效果def warmup_scheduler(optimizer, current_step, warmup_steps): 线性学习率预热 lr_scale min(1.0, float(current_step) / warmup_steps) for param_group in optimizer.param_groups: param_group[lr] param_group[initial_lr] * lr_scale # 使用示例 warmup_steps 1000 for batch_idx, (inputs, labels) in enumerate(train_loader): warmup_scheduler(optimizer, batch_idx, warmup_steps) ...在训练Transformer模型时学习率预热几乎是标配。我实验发现适当的预热可以显著提高模型训练的稳定性。4. 自定义优化器开发指南虽然PyTorch提供了丰富的内置优化器但有时候我们需要实现自己的优化算法。PyTorch的优化器基类设计得非常灵活便于扩展。4.1 理解Optimizer基类所有PyTorch优化器都继承自torch.optim.Optimizer基类。要实现自定义优化器我们需要了解它的核心结构class MyOptimizer(Optimizer): def __init__(self, params, lr1e-3, beta10.9, beta20.999): defaults dict(lrlr, beta1beta1, beta2beta2) super().__init__(params, defaults) def step(self, closureNone): 执行单次参数更新 loss None if closure is not None: loss closure() for group in self.param_groups: for p in group[params]: if p.grad is None: continue # 获取梯度 grad p.grad.data # 在这里实现参数更新逻辑 # ... # p.data.add_(-group[lr], grad) return loss基类已经处理了参数组管理、状态维护等通用功能我们只需要关注具体的参数更新逻辑。4.2 实现简单的自定义优化器让我们实现一个带有动量的简单SGD变种class MySGD(Optimizer): def __init__(self, params, lr1e-3, momentum0.9): defaults dict(lrlr, momentummomentum) super().__init__(params, defaults) def step(self, closureNone): loss None if closure is not None: loss closure() for group in self.param_groups: momentum group[momentum] for p in group[params]: if p.grad is None: continue grad p.grad.data # 初始化状态 state self.state[p] if momentum_buffer not in state: state[momentum_buffer] torch.zeros_like(p.data) # 更新动量缓冲 buf state[momentum_buffer] buf.mul_(momentum).add_(grad) # 更新参数 p.data.add_(-group[lr], buf) return loss这个优化器虽然简单但包含了自定义优化器的所有关键要素参数初始化、状态管理和更新逻辑。4.3 测试自定义优化器实现优化器后我们需要验证它的正确性# 测试函数 def test_optimizer(): # 创建一个简单的优化问题最小化 y (x - 5)^2 x torch.tensor([10.0], requires_gradTrue) opt MySGD([x], lr0.1, momentum0.9) for _ in range(100): opt.zero_grad() y (x - 5).pow(2) y.backward() opt.step() print(fx {x.item():.4f}, y {y.item():.4f}) assert abs(x.item() - 5) 1e-3 test_optimizer()在实际项目中我通常会先用这种简单问题验证优化器的基本功能然后再应用到复杂模型上。这样可以快速发现算法实现中的问题。4.4 性能优化技巧自定义优化器的性能也很重要。以下是一些优化建议尽量减少在step()中创建临时张量使用原地操作如mul_而不是mul对于复杂优化器可以考虑用C扩展实现关键部分我曾经实现过一个复杂的二阶优化器最初的Python版本比标准SGD慢10倍。通过将核心部分用C实现最终只比SGD慢2倍左右。