PyTorch梯度累积模拟大batch训练是否需设置retain_graph=True？

问题1：是否需要给loss.backward()添加retain_graph=True参数？ #

不需要，你提供的原生写法是正确的，无需额外添加该参数。

• 原因：retain_graph=True的作用是反向传播后保留当前的计算图，仅当你需要对同一份计算图/同一份loss多次调用反向传播时才需要使用。而梯度累积的场景下，每个batch都会独立执行前向传播生成新的计算图，单次反向传播结束后上一个batch的计算图就可以释放，加了retain_graph=True反而会导致计算图无法释放，显存占用持续升高，甚至触发显存溢出。

问题2：每次调用backward()时各参数的梯度是否是直接累加的？ #

是的，PyTorch的默认逻辑就是梯度累加，这也是梯度累积方案能生效的核心基础。

• 原因：PyTorch中模型参数的.grad属性默认在调用backward()时，会将当前batch计算得到的梯度值叠加到原有值上，而非直接覆盖。只有你显式调用optimizer.zero_grad()时，梯度才会被清空。你每N个batch才执行一次梯度清空和参数更新，正好就实现了N个batch的梯度累积效果，等效于使用N倍大小的batch size训练。
• 可选优化：如果要和直接使用N倍batch size的训练效果完全对齐，建议在反向传播前把loss除以N，抵消多batch梯度累加带来的梯度量级放大：

def train(model, optimizer, dataloader, num_epochs, N):
     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 = loss / N
             loss.backward()
             if batch_num % N == 0:
                 optimizer.step()
                 optimizer.zero_grad(set_to_none=True)

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