PyTorch mini batch训练时optimizer.zero_grad()的正确调用位置是哪里？

第二种写法（在每个batch迭代内部调用optimizer.zero_grad()）完全符合PyTorch标准mini-batch训练的规范，第一种写法属于特殊场景下的自定义策略，不适用常规训练流程。

PyTorch的张量梯度默认采用累加机制：调用tensor.backward()计算得到的梯度，会直接累加到对应参数的.grad属性上，不会自动覆盖上一次计算得到的梯度值。优化器执行optimizer.step()更新参数时，读取的就是参数.grad属性里存储的累加值。

第一种写法的问题（epoch开头调用zero_grad） #

• 每个epoch启动时仅执行1次梯度清零
• 整个epoch内所有batch的反向传播操作，都会把当前batch计算出的梯度，累加上之前所有batch残留在.grad里的梯度值
• 最终优化器每次执行step()时，拿到的是跨多个batch混合累加的梯度，完全不是单个mini-batch对应的真实梯度，会导致训练梯度失真，偏离正常的mini-batch收敛逻辑
• 举个具体执行流程：
  1. epoch开始，清零参数梯度为0
  2. 第1个batch：反向传播得到梯度值2，step()用梯度2更新参数，此时参数.grad残留值为2
  3. 第2个batch：因为没有重新清零，反向传播得到当前batch的梯度值3，累加后总梯度为2+3=5，step()直接用混合后的5更新参数，完全违背单batch训练的预期

第二种写法的正确性（batch循环内部调用zero_grad） #

• 每一次进入单个batch的迭代流程，只要在调用cost.backward()之前完成梯度清零，就能保证当前反向传播计算得到的梯度完全对应这一个batch的前向计算结果，不会被之前batch的残留梯度污染
• 代码里zero_grad()放在前向传播之后、计算loss的前后位置，只要处于backward()调用之前，都不影响最终效果，完全满足梯度隔离的要求

特殊场景补充：第一种写法并不是完全错误，它是梯度累加训练策略的基础实现逻辑——当显存不足需要模拟更大batch的训练效果时，会跳过每个batch后的梯度清零，连续累加N个batch的梯度后再执行一次optimizer.step()和zero_grad()。但这种是自定义训练策略，并非常规mini-batch训练的默认实现。

内容的提问来源于stack exchange，提问作者이준재