如何在PyTorch中实现全批梯度下降和Nesterov动量？

Question

我正在做一个机器学习项目，使用的是PyTorch。我需要用一种叫做全批量梯度下降的方法来优化我的模型。这个方法的关键要求是，每次更新时都要使用数据集中所有的数据点。现在我遇到的问题是，现有的torch.optim.SGD优化器并不支持一次性使用整个数据集进行更新。这对我的项目来说非常重要，因为我需要优化过程考虑到所有的数据点，以确保模型参数的更新尽可能准确。

另外，我还想在优化过程中保留Nesterov动量的使用。我知道可以通过把批量大小设置为整个数据集来模拟全批量更新，但我更想知道有没有更优雅或直接的方法，在PyTorch中实现一个真正的梯度下降优化器，并且支持Nesterov动量。

理想情况下，我希望找到一个解决方案或指导，来实现或配置一个符合以下标准的优化器：

• 每次参数更新都使用整个数据集（真正的梯度下降行为）。
• 结合Nesterov动量，以提高收敛效率。
• 与PyTorch的其他部分兼容，能够继承torch.optim.Optimizer。

Answer 1

PyTorch中的SGD（随机梯度下降）实现其实和批量处理没有直接关系！它只使用在反向传播过程中计算并存储在参数的.grad属性中的梯度。因此，用于计算的批量大小和用于优化的批量大小是可以分开的。

你现在可以选择：

a) 将所有样本作为一个大批量输入模型，也就是把批量大小设置为数据集的大小，或者

b) 在进行一次优化步骤之前，先对多个小批量的梯度进行累积（伪代码）：

model = YourModel()
data = YourDataSetOrLoader()
optim = torch.optim.SGD(model.parameters())
for full_batch_step in range(100)
   #this sets the accumulated gradient to zero
   optim.zero_grad()
   for batch in data:
      f=model(data)
      # this adds the gradient wrt to the parameters for the current datapoint to the model paramters
      f.backward()
  # now after we summed the gradient for all samples, we do a GD step.
  optim.step()