PyTorch中optim的用法

PyTorch中optim的用法

torch.optim是PyTorch中用于实现各种优化算法的模块，在训练神经网络时非常有用。它提供了各种类型的优化器，如随机梯度下降、Adam、Adagrad等，并允许我们根据需要自定义优化器。

定义优化器

在torch.optim中，我们可以使用各种类型的优化器来更新神经网络模型的参数。这些优化器通常被定义为类，并通过指定要优化的参数和学习率等超参数来进行初始化。以下是一些常用的优化器及其用法：

• 随机梯度下降（SGD）优化器：

import torch.optim as optim

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

# 在训练循环中使用优化器
optimizer.zero_grad()
loss.backward()
optimizer.step()

• Adam优化器：

import torch.optim as optim

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

# 在训练循环中使用优化器
optimizer.zero_grad()
loss.backward()
optimizer.step()

• Adagrad优化器：

import torch.optim as optim

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

# 在训练循环中使用优化器
optimizer.zero_grad()
loss.backward()
optimizer.step()

调整学习率

在训练神经网络时，我们通常需要根据训练的进程调整学习率。在torch.optim中，我们可以通过以下方式调整学习率：

import torch.optim as optim

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

# 调整学习率
for epoch in range(num_epochs):
    if epoch == 30:
        optimizer.param_groups[0]['lr'] = 0.001
    elif epoch == 50:
        optimizer.param_groups[0]['lr'] = 0.0001

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

自定义优化器

在torch.optim中，我们还可以根据需要自定义优化器。我们可以通过继承torch.optim.Optimizer类并实现__init__和step方法来创建自己的优化器。

import torch.optim as optim

class MyOptimizer(optim.Optimizer):
    def __init__(self, params, lr=0.01):
        self.lr = lr
        super(MyOptimizer, self).__init__(params)
    
    def step(self, closure=None):
        for group in self.param_groups:
            for param in group['params']:
                if param.grad is None:
                    continue
                grad = param.grad.data
                param.data -= self.lr * grad

# 使用自定义优化器
optimizer = MyOptimizer(model.parameters(), lr=0.01)

以上是对PyTorch中optim模块的用法介绍，它们可以帮助我们更方便地实现各种优化算法，并调整学习率以提高训练效果。