PyTorch torch.optim 优化器模块

torch.optim 是 PyTorch 中用于优化神经网络参数的模块，提供了多种优化算法（如 SGD、Adam 等），用于最小化损失函数。以下是 torch.optim 模块的参考手册，涵盖核心功能、常用优化器、使用方法和示例，力求简洁且全面。

---

1. torch.optim 模块概述

torch.optim 提供了一系列优化算法，用于更新神经网络的参数以最小化损失函数。它与 torch.nn 和 torch.autograd 紧密配合，支持梯度下降和更复杂的优化策略。主要功能包括：

• 优化算法：如 SGD、Adam、RMSprop 等。
• 学习率调度：通过 torch.optim.lr_scheduler 动态调整学习率。
• 参数管理：支持分组参数和自定义优化设置。

---

2. 核心组件与使用方法

(1) 基本用法

优化器的核心是通过梯度更新模型参数。典型流程：

1. 定义模型和损失函数。
2. 初始化优化器，传入模型参数。
3. 计算损失、反向传播、更新参数。

import torch
import torch.nn as nn
import torch.optim as optim

# 简单模型
model = nn.Linear(10, 2)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练步骤
data, target = torch.randn(32, 10), torch.randn(32, 2)
optimizer.zero_grad()  # 清零梯度
output = model(data)
loss = criterion(output, target)
loss.backward()  # 计算梯度
optimizer.step()  # 更新参数

(2) 常用优化器

torch.optim 提供了多种优化算法，以下是常用的几种：

• optim.SGD：随机梯度下降，支持动量和 Nesterov 动量。

  optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, nesterov=True)

• 参数：
  • lr：学习率。
  • momentum：动量因子（默认 0）。
  • weight_decay：L2 正则化（默认 0）。
  • nesterov：是否使用 Nesterov 动量。
• optim.Adam：自适应矩估计，结合一阶和二阶动量，适合大多数任务。

  optimizer = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999), weight_decay=0)

• 参数：
  • lr：学习率（常用 0.001 或 0.0001）。
  • betas：一阶/二阶动量衰减率。
  • eps：数值稳定性常数（默认 1e-8）。
• optim.RMSprop：基于梯度平方均值的优化，适合非平稳目标。

  optimizer = optim.RMSprop(model.parameters(), lr=0.01, alpha=0.99)

• 其他优化器：
• optim.Adagrad：自适应学习率，适合稀疏数据。
• optim.AdamW：Adam 的正则化版本，改进权重衰减。
• optim.LBFGS：拟牛顿法，适合小数据集（需提供闭包函数）。

(3) 参数分组

优化器支持为不同参数组设置不同的超参数（如学习率）。

optimizer = optim.Adam([
    {'params': model.fc1.parameters(), 'lr': 0.001},
    {'params': model.fc2.parameters(), 'lr': 0.0001}
], lr=0.01)  # 默认学习率

(4) 学习率调度（torch.optim.lr_scheduler）

动态调整学习率以改善收敛。常用调度器：

• lr_scheduler.StepLR：每隔固定步数降低学习率。

  scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

• lr_scheduler.ReduceLROnPlateau：当指标（如验证损失）停止改善时降低学习率。

  scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10)

• lr_scheduler.CosineAnnealingLR：余弦退火学习率。

  scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)

使用调度器：

for epoch in range(epochs):
    optimizer.step()
    scheduler.step()  # 更新学习率

(5) 关键方法

• optimizer.zero_grad()：清零所有参数的梯度。
• optimizer.step()：根据梯度更新参数。
• optimizer.state_dict()：保存优化器状态（如动量）。
• optimizer.load_state_dict(state_dict)：加载优化器状态。

---

3. 完整训练示例

以下是一个结合优化器和学习率调度的完整示例：

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

# 数据
data = torch.randn(100, 10)
labels = torch.randint(0, 2, (100,))
dataset = TensorDataset(data, labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 2)

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

model = Net().to('cuda' if torch.cuda.is_available() else 'cpu')
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

# 训练
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.train()
for epoch in range(20):
    total_loss = 0
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    scheduler.step()  # 更新学习率
    print(f'Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}, LR: {scheduler.get_last_lr()[0]:.6f}')

---

4. 进阶用法

• 梯度裁剪：防止梯度爆炸。

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

• 梯度累积：处理大批量数据。

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

• 混合精度训练：结合 torch.cuda.amp 提高效率。

  from torch.cuda.amp import autocast, GradScaler
  scaler = GradScaler()
  with autocast():
      outputs = model(inputs)
      loss = criterion(outputs, targets)
  scaler.scale(loss).backward()
  scaler.step(optimizer)
  scaler.update()

• 自定义优化器：继承 torch.optim.Optimizer 定义新算法。

---

5. 常见问题与注意事项

• 梯度清零：每次迭代前调用 zero_grad()，否则梯度会累积。
• 学习率选择：Adam 通常用 0.001 或 0.0001，SGD 可能需要 0.01 或更高。
• 调度器调用：scheduler.step() 位置需注意，ReduceLROnPlateau 需要传入指标（如损失）。
• 设备一致性：确保模型、数据和优化器操作在同一设备（CPU/GPU）。
• 权重衰减：weight_decay 参数用于 L2 正则化，典型值为 1e-4 或 1e-5。

---

6. 参考资源

• 官方文档：
• torch.optim：https://pytorch.org/docs/stable/optim.html
• lr_scheduler：https://pytorch.org/docs/stable/optim.html#torch.optim.lr_scheduler
• 教程：https://pytorch.org/tutorials/beginner/basics/optimization_tutorial.html
• 社区论坛：https://discuss.pytorch.org/

---

7. 进一步帮助

如果你需要特定优化器的详细配置（如 AdamW 的权重衰减）、学习率调度策略、或针对特定任务（如 Transformer 训练）的优化建议，请提供更多细节，我可以为你提供更精准的代码或指导！