聊聊pytorch中Optimizer与optimizer.step()的用法

当我们在使用PyTorch进行神经网络模型训练时，一般需要定义一个优化器，通过这个优化器去更新神经网络模型中的参数。那么，本篇攻略就来详细讲解一下PyTorch中Optimizer与optimizer.step()的用法。

什么是Optimizer

在进行模型训练时，我们在每一轮迭代中，都需要更新模型中的参数，使其更好地拟合数据。这个过程的实现就需要使用到优化器。优化器是一个算法，它根据我们指定的损失函数和一些超参数，来自动地调整模型参数，以达到最佳的效果。

在PyTorch中，我们可以通过torch.optim模块提供的优化器来实现。PyTorch提供了一些常用的优化器类，例如：torch.optim.SGD、torch.optim.Adam、torch.optim.Adagrad等。

Optimizer的用法

下面我们来具体讲解一下Optimizer的用法：

1. 创建Optimizer对象

我们首先需要创建一个优化器对象。要创建一个优化器对象，需要提供以下两个参数：
* 第一个参数：需要更新的模型参数，也就是PyTorch中的张量。我们一般通过model.parameters()来获取模型中的参数张量。
* 第二个参数：优化器的超参数。不同的优化器具有不同的超参数，例如：学习率lr、动量momentum、权重衰减weight_decay等。我们需要根据实际情况来选择合适的超参数。

以SGD优化器为例，创建方式如下：

import torch.optim as optim

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

2. 计算梯度

在优化器更新模型参数之前，我们需要先计算出模型中各个参数的梯度。梯度计算一般使用autograd模块的backward()函数来实现。例如：

loss = criterion(output, target)
loss.backward()   # 计算loss对模型参数的梯度

3. 更新模型参数

在计算了各个参数的梯度之后，我们可以调用优化器的step()函数来更新模型中的参数。此时，优化器会根据上述给定的超参数，在梯度方向上对模型参数进行调整，使得模型的损失函数值更小。

以SGD优化器为例，更新方式如下：

optimizer.step()

4. 清空梯度

在完成了一次参数更新之后，我们需要将之前计算的参数梯度清空，以便进行下一轮迭代。清空操作一般使用zero_grad()函数来实现。例如：

optimizer.zero_grad()   # 将参数的梯度清零

示例说明

下面我们举两个例子来说明Optimizer与optimizer.step()的用法。

示例一：使用SGD优化器训练线性回归模型

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

class LinearRegression(nn.Module):
    def __init__(self):
        super(LinearRegression, self).__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        out = self.linear(x)
        return out

# 随机生成数据
x = torch.randn(100, 1)
y = 3 * x + 0.4 * torch.randn(100, 1)

# 定义模型、损失函数和优化器
model = LinearRegression()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(1000):
    # 前向传播
    output = model(x)

    # 计算损失函数并反向传播
    loss = criterion(output, y)
    loss.backward()

    # 更新参数
    optimizer.step()

    # 清空梯度
    optimizer.zero_grad()

    # 打印损失值
    if (epoch + 1) % 100 == 0:
        print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, 1000, loss.item()))

示例二：使用Adam优化器训练卷积神经网络

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

# 定义模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=5)
        self.pool1 = nn.MaxPool2d(kernel_size=2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5)
        self.pool2 = nn.MaxPool2d(kernel_size=2)
        self.fc1 = nn.Linear(4 * 4 * 32, 1024)
        self.fc2 = nn.Linear(1024, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = nn.functional.relu(x)
        x = self.pool2(x)
        x = x.view(-1, 4 * 4 * 32)
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.fc2(x)
        return x

# 加载数据集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=torchvision.transforms.ToTensor(), download=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=100, shuffle=True)

# 定义模型、损失函数和优化器
model = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(10):
    for i, (images, labels) in enumerate(train_loader):
        # 前向传播
        output = model(images)

        # 计算损失函数并反向传播
        loss = criterion(output, labels)
        loss.backward()

        # 更新参数
        optimizer.step()

        # 清空梯度
        optimizer.zero_grad()

        # 打印训练进度
        if (i + 1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, 10, i+1, len(train_dataset)//100, loss.item()))

通过上述示例，我们可以看到，实现优化器的代码非常简单，而且能够大大简化模型训练的过程。