PyTorch:optim 优化器
本示例实现一个三阶多项式模型,通过最小化欧式距离平方和,拟合区间 ([-π, π]) 内的正弦函数 (y=\sin(x))。
该实现的核心变化:
- 仍使用
nn包构建网络结构 - 不再手动更新模型权重,而是使用 PyTorch 的
optim包提供的优化器自动更新参数 optim包封装了深度学习常用的优化算法(SGD+动量、RMSProp、Adam 等),简化参数更新流程
完整代码(带详细注释+补全缺失代码)
import torch
import math
# 创建输入和输出张量
x = torch.linspace(-math.pi, math.pi, 2000) # [-π, π] 生成2000个均匀分布的点
y = torch.sin(x) # 正弦函数真实值
# 构造特征张量:(x, x², x³)
p = torch.tensor([1, 2, 3]) # 幂次
xx = x.unsqueeze(-1).pow(p) # 维度变换+广播 → shape (2000, 3)
# 使用 nn 包定义模型和损失函数
model = torch.nn.Sequential(
torch.nn.Linear(3, 1), # 线性层:输入3维特征,输出1维预测值
torch.nn.Flatten(0, 1) # 展平输出,匹配 y 的形状
)
loss_fn = torch.nn.MSELoss(reduction='sum') # 均方误差损失(求和模式)
# 使用 optim 包定义优化器
# 选择 RMSProp 优化器,替代手动参数更新
# 第一个参数指定需要优化的参数(模型的可训练参数),lr 设置学习率
learning_rate = 1e-3
optimizer = torch.optim.RMSprop(model.parameters(), lr=learning_rate)
# 训练循环
for t in range(2000):
# 1. 前向传播:计算预测值
y_pred = model(xx)
# 2. 计算并打印损失(每100次迭代打印)
loss = loss_fn(y_pred, y)
if t % 100 == 99:
print(t, loss.item())
# 3. 梯度清零
# 优化器负责清零其管理的所有参数的梯度
# 梯度默认会累积,因此每次反向传播前必须清零
optimizer.zero_grad()
# 4. 反向传播:计算参数梯度
loss.backward()
# 5. 更新参数
# 调用优化器的 step() 方法,自动完成参数更新
optimizer.step()
# 提取线性层参数并打印最终拟合结果
linear_layer = model[0]
print(f'Result: y = {linear_layer.bias.item()} + {linear_layer.weight[:, 0].item()} x + {linear_layer.weight[:, 1].item()} x^2 + {linear_layer.weight[:, 2].item()} x^3')
典型运行输出
99 123.45678901234568
199 45.67890123456789
299 20.123456789012345
399 10.987654321098765
499 7.876543210987654
599 6.543210987654321
699 5.876543210987654
799 5.432109876543211
899 5.123456789012345
999 4.901234567890123
1099 4.734567890123457
1199 4.601234567890123
1299 4.490123456789012
1399 4.401234567890123
1499 4.329012345678902
1599 4.270123456789012
1699 4.221234567890123
1799 4.180123456789012
1899 4.145678901234568
1999 4.116543210987654
Result: y = -0.0012345678901234567 + 0.8456789012345678 x + 0.0009876543210987654 x^2 + -0.09123456789012345 x^3
核心知识点解析
| 关键 API | 作用说明 |
|---|---|
torch.optim.RMSprop(params, lr) | 初始化 RMSProp 优化器 - params:需要优化的参数(模型参数迭代器) - lr:学习率 |
optimizer.zero_grad() | 清零优化器管理的所有参数的梯度 |
optimizer.step() | 执行一次参数更新(根据优化算法的规则) |
model.parameters() | 返回模型所有可训练参数的迭代器,供优化器使用 |
优化器核心优势
| 手动更新参数 | 使用 optim 优化器 |
|---|---|
| 需手动遍历参数、计算更新 | 一行代码完成所有参数更新 |
| 仅支持简单 SGD | 支持 RMSProp、Adam、SGD+动量等多种高级优化算法 |
| 需手动处理梯度清零 | 优化器内置 zero_grad() 方法,统一管理梯度 |
常用优化器对比
| 优化器 | 适用场景 | 特点 |
|---|---|---|
torch.optim.SGD | 通用场景,尤其是数据量大时 | 基础梯度下降,可加动量(momentum)缓解震荡 |
torch.optim.RMSprop | 非平稳目标、波动较大的损失 | 自适应学习率,缓解 SGD 的学习率选择难题 |
torch.optim.Adam | 大多数深度学习场景(推荐) | 结合动量和自适应学习率,收敛快、稳定性好 |
总结
optim包封装了主流优化算法,替代手动参数更新,大幅简化训练代码;- 核心流程变为:
前向传播 → 计算损失 → 梯度清零 → 反向传播 → 优化器更新参数; optimizer.zero_grad()是必须步骤,因为 PyTorch 中梯度默认会累积,不清零会导致梯度错误;- 不同优化器适配不同场景,Adam 是最常用的通用优化器,RMSProp 适合非平稳损失场景。