使用 PyTorch 为更大的批量累积梯度

Question

为了模仿更大的批次大小，我希望能够在 PyTorch 中为模型每 N 个批次累积梯度，例如：

def train(model, optimizer, dataloader, num_epochs):
     model.train()
     model.cuda()
     for epoch_num in range(1, num_epochs+1):
         for batch_num, data in enumerate(dataloader):
             ims = data.to('cuda:0') 
             loss = model(ims)
             loss.backward()
             if batch_num % N == 0 and batch_num != 0:
                 optimizer.step()
                 optimizer.zero_grad(set_to_none=True)

对于这种方法，我需要添加标志 retain_graph=True，即

loss.backward(retain_graph=True)

以这种方式，每个反向调用的梯度是否简单地按每个参数求和？

Answer 1

如果你想对同一个计算图进行多次反向传播，你需要设置retain_graph=True，利用一次正向传播的中间结果。例如，如果您在计算 loss 一次后多次调用 loss.backward()，或者如果您从图表的不同部分有多个损失要从中反向传播（一个很好的解释可能是找到 ).

在你的例子中，对于每一次前向传播，你只反向传播一次。因此，一旦计算了梯度，您就不需要存储来自计算图的中间结果。

简而言之：

图中的中间输出在反向传递后被清除，除非使用retain_graph=True.
Gradients 默认累积，除非使用 zero_grad.

使用 PyTorch 为更大的批量累积梯度

Accumulating gradients for a larger batch size with PyTorch

pytorch