PyTorch 基础教程（三）：掌握梯度下降算法

2024-01-06 06:00:45

引言：理解梯度下降

在机器学习中，优化算法对于寻找模型参数的最优值至关重要，以最小化损失函数。梯度下降是一种迭代算法，它通过反复向损失函数的负梯度方向移动来优化模型参数。这种移动有助于逐渐减少损失，从而提高模型的性能。

案例分析：梯度下降的直观理解

让我们从一个简单的案例开始，考虑一个一元函数 y = x^2。对于该函数，梯度为 y' = 2x。梯度下降算法从某个初始点（例如 x = 1）开始，然后沿负梯度方向（-2x）移动，更新 x 的值。

x_new = x_old - learning_rate * y'

其中，learning_rate 是控制步长的超参数。

通过多次迭代，x 的值将逐步接近函数的最小值（x = 0）。这种直观的理解有助于理解梯度下降算法的基本原理。

在 PyTorch 中实现梯度下降

在 PyTorch 中，可以使用 torch.optim.SGD 类来实现梯度下降。以下代码片段展示了如何使用 SGD 优化器优化一个简单的线性回归模型：

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

# 定义模型
model = nn.Linear(1, 1)

# 定义损失函数
criterion = nn.MSELoss()

# 定义优化器
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(1000):
    # 正向传播
    y_pred = model(x)
    loss = criterion(y_pred, y)

    # 反向传播
    loss.backward()

    # 梯度下降更新参数
    optimizer.step()

    # 清除梯度
    optimizer.zero_grad()