PyTorch RNN训练：手动参数更新与optim.step()差异及正确用法

问题背景 #

在复现PyTorch字符级RNN姓名分类教程时，原教程采用手写参数更新逻辑可正常收敛，改用optim.Adam配合optimizer.step()实现参数更新时，模型完全没有学习效果，损失无下降趋势，异常训练结果参考：。
教程给出的基础RNN模型定义如下：

import torch.nn as nn

class RNN(nn.Module):
  def __init__(self, input_size, hidden_size, output_size):
    super(RNN, self).__init__()
    self.hidden_size = hidden_size
    
    self.i2h = nn.Linear(input_size + hidden_size, hidden_size)
    self.i2o = nn.Linear(input_size + hidden_size, output_size)
    self.softmax = nn.LogSoftmax(dim=1)
  
  def forward(self, input, hidden):
    combined = torch.cat((input, hidden), 1)
    hidden = self.i2h(combined)
    output = self.i2o(combined)
    output = self.softmax(output)
    return output, hidden

  def initHidden(self):
    return torch.zeros(1, self.hidden_size)

n_hidden = 128
rnn = RNN(n_letters, n_hidden, n_categories)
learning_rate = 0.005
criterion = nn.NLLLoss()

原教程可正常收敛的手写更新训练逻辑：

def train(category_tensor, name_tensor):
  hidden = rnn.initHidden()
  rnn.zero_grad()

  for i in range(name_tensor.size()[0]):
    output, hidden = rnn(name_tensor[i], hidden)
  
  loss = criterion(output, category_tensor)
  loss.backward()

  # 手写Vanilla SGD更新
  for p in rnn.parameters():
    p.data.add_(p.grad.data, alpha=-learning_rate)

  return output, loss.item()

改写后不收敛的optimizer实现：

optimizer = optim.Adam(rnn.parameters(), lr = learning_rate)  

def train(category_tensor, name_tensor):
  hidden = rnn.initHidden()
  rnn.zero_grad()

  for i in range(name_tensor.size()[0]):
    output, hidden = rnn(name_tensor[i], hidden)
  
  loss = criterion(output, category_tensor)
  
  optimizer.zero_grad()
  loss.backward()
  optimizer.step()

  return output, loss.item()

一、两种更新方式的核心差异 #

两种实现的差异可分为两点，其中代码逻辑错误是不收敛的核心原因，和优化器本身无关：

• 优化算法本质不同
  原教程手写实现的是最朴素的Vanilla SGD：无动量、无自适应学习率，直接拿当前步计算出的梯度乘以负学习率，直接覆盖更新参数值，没有任何内部状态，更新逻辑完全确定。
  改写版本使用的Adam是自适应优化器，内部会为每个参数维护梯度一阶动量、二阶动量两个滑动平均状态，更新时会对梯度做偏差修正，不同参数的实际更新步长会自适应调整，本身收敛特性和Vanilla SGD有区别，但这不是不收敛的诱因。
• 梯度清零逻辑混乱冗余
  改写代码先后两次执行梯度清零：第一次在前向传播前调用rnn.zero_grad()，第二次在loss计算完成、反向传播前调用optimizer.zero_grad()。虽然语法上第二次清零后才做反向传播，梯度值本身计算没有错误，但冗余的清零操作容易和优化器状态更新冲突，加上直接套用了Vanilla SGD的学习率，直接导致Adam更新失效。

二、optimizer.step()正确实现的调整点 #

只需要做3处简单修改即可正常收敛，收敛速度通常会快于原手写SGD：

• 统一梯度清零逻辑，删除冗余调用
  删掉前向传播前的rnn.zero_grad()，将optimizer.zero_grad()统一放在train函数最开头、所有前向操作之前，避免逻辑冲突。梯度清零的核心作用是防止上一轮迭代的梯度累加到当前步，只要在当前轮反向传播前执行即可，放在最前是最稳妥的工业界写法。
• 调整Adam适配的学习率
  原代码的learning_rate=0.005是针对Vanilla SGD调试的，Adam对学习率敏感度远高于普通SGD，将学习率调整到1e-3 ~ 3e-3区间即可，不要直接套用SGD的学习率参数。
• 隐状态逻辑不需要修改：当前代码每个训练样本重新初始化零隐状态，符合字符级分类任务的要求，不需要对hidden做额外detach操作，每轮迭代的计算图会在迭代结束后自动释放。

修正后的可运行训练代码如下：

import torch.optim as optim

# 调整Adam适配的学习率
learning_rate = 0.001
optimizer = optim.Adam(rnn.parameters(), lr=learning_rate)
criterion = nn.NLLLoss()

def train(category_tensor, name_tensor):
  # 统一在最开头清零梯度
  optimizer.zero_grad()

  hidden = rnn.initHidden()
  for i in range(name_tensor.size()[0]):
    output, hidden = rnn(name_tensor[i], hidden)
  
  loss = criterion(output, category_tensor)
  loss.backward()
  optimizer.step()

  return output, loss.item()

内容的提问来源于stack exchange，提问作者UiJin