PyTorch 基础学习（9）- 训练优化器

系列文章：
PyTorch 基础学习（1） - 快速入门
PyTorch 基础学习（2）- 张量 Tensors
PyTorch 基础学习（3） - 张量的数学操作
PyTorch 基础学习（4）- 张量的类型
PyTorch 基础学习（5）- 神经网络
PyTorch 基础学习（6）- 函数API
PyTorch 基础学习（7）- 自动微分
PyTorch 基础学习（8）- 多进程并发

介绍

PyTorch 提供了一个强大的优化器模块 torch.optim，它实现了多种常用的优化算法，用于训练神经网络模型。在本教程中，我们将探讨如何使用这些优化器，以及它们的应用场景和示例代码。

基本概念

优化器的主要任务是根据损失函数的梯度来更新模型参数，从而最小化损失函数。PyTorch 提供了各种优化算法，如随机梯度下降（SGD）、Adam、RMSprop 等，这些算法在不同的场景下有着各自的优势。

构建优化器

使用 torch.optim 中的优化器，你需要首先构建一个优化器对象，并将模型的参数传递给它。优化器对象会保存这些参数的当前状态，并在每次调用 step() 时根据梯度更新参数。

基本步骤：

1. 定义模型参数：通过 model.parameters() 获取模型的所有参数。
2. 选择优化算法：如 SGD、Adam、RMSprop 等。
3. 设置优化器参数：如学习率（lr）、权重衰减（weight_decay）等。

代码示例：

import torch.optim as optim

# 使用随机梯度下降（SGD）优化器
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 使用 Adam 优化器
optimizer = optim.Adam(model.parameters(), lr=0.0001)

为每个参数组单独设置选项

有时你可能希望为模型的不同部分设置不同的优化选项。比如，你可能希望在神经网络的不同层上使用不同的学习率。PyTorch 支持为每个参数组单独设置选项。

代码示例：

optimizer = optim.SGD([
    {'params': model.base.parameters(), 'lr': 1e-2},
    {'params': model.classifier.parameters(), 'lr': 1e-3}
], momentum=0.9)

在这个例子中，model.base 的参数使用 1e-2 的学习率，而 model.classifier 的参数使用 1e-3 的学习率。同时，0.9 的动量因子会应用于所有参数。

执行参数更新

执行参数更新的关键步骤是调用 optimizer.step()，这会根据梯度更新模型的所有参数。通常，这个步骤是在每次前向传播和反向传播之后执行的。

基本步骤：

1. 清零梯度：在反向传播之前，使用 optimizer.zero_grad() 清除之前的梯度。
2. 前向传播：计算输出和损失。
3. 反向传播：计算梯度。
4. 更新参数：调用 optimizer.step() 更新参数。

代码示例：

for input, target in dataset:
    optimizer.zero_grad()  # 清除梯度
    output = model(input)  # 前向传播
    loss = loss_fn(output, target)  # 计算损失
    loss.backward()  # 反向传播
    optimizer.step()  # 更新参数

使用闭包进行多次优化

对于一些高级的优化算法，如 L-BFGS，可能需要多次计算损失和梯度。在这种情况下，你需要传入一个闭包函数，该函数会在每次迭代时重新计算损失和梯度。

代码示例：

for input, target in dataset:
    def closure():
        optimizer.zero_grad()  # 清除梯度
        output = model(input)  # 前向传播
        loss = loss_fn(output, target)  # 计算损失
        loss.backward()  # 反向传播
        return loss

    optimizer.step(closure)  # 使用闭包进行参数更新

常见优化算法及其应用场景

1. SGD（随机梯度下降）：

   • 应用场景：适用于大规模数据集和深度神经网络训练，特别是在需要简单且计算资源有限的场景。
   • 代码示例：

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

2. Adam：

   • 应用场景：适用于需要快速收敛的场景，常用于深度学习中的图像处理和自然语言处理任务。
   • 代码示例：

   optimizer = optim.Adam(model.parameters(), lr=0.001)
   

3. RMSprop：

   • 应用场景：适用于处理非平稳目标的场景，如处理递归神经网络中的长序列数据。
   • 代码示例：

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

综合应用实例：图像分类任务中的优化器应用

在一个图像分类任务中，假设我们有一个卷积神经网络模型 model，我们将使用 Adam 优化器来训练该模型。

代码示例：

import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 定义一个简单的卷积神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载数据集
train_dataset = datasets.MNIST('../data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('../data', train=False, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

# 实例化模型和优化器
model = Net()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 损失函数
loss_fn = nn.CrossEntropyLoss()

# 训练循环
for epoch in range(10):
    model.train()  # 设置模型为训练模式
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()

    print(f'Epoch {epoch + 1}, Loss: {loss.item()}')

# 验证模型在测试集上的性能
model.eval()  # 设置模型为评估模式
test_loss = 0
correct = 0
with torch.no_grad():
    for data, target in test_loader:
        output = model(data)
        test_loss += loss_fn(output, target).item()  # 累加测试损失
        pred = output.argmax(dim=1, keepdim=True)  # 获取预测结果
        correct += pred.eq(target.view_as(pred)).sum().item()

test_loss /= len(test_loader.dataset)  # 计算平均损失
accuracy = 100. * correct / len(test_loader.dataset)  # 计算准确率

print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')

这个例子展示了如何使用 Adam 优化器来训练一个简单的卷积神经网络进行图像分类任务。在每个 epoch 中，模型的参数都会根据损失函数的梯度进行更新。

通过本教程，你应该已经掌握了 PyTorch 中 torch.optim 模块的基本使用方法，并了解了不同优化算法的应用场景及其实现方式。在实际项目中，你可以根据任务的具体需求选择合适的优化器来提高模型的性能。