(十)梯度累积与每轮清理
# 项目文件:demos/gradient_accumulation.py
import torch
w = torch.tensor(1.0, requires_grad=True)
for step in range(2):
loss = (w * 2 - 5).square() # 每一轮重新前向,建立这一轮的图。
loss.backward()
print("第几次 backward", step + 1, "grad", w.grad.item())
w.grad = None # 下一轮从尚未写入梯度开始,而不是继续加到 -24。
(w * 2 - 5).square().backward()
print("清理后", w.grad.item())第一次反向得到 -12,第二次重新前向、再反向后 .grad 为 -24。两次没有更新参数,当前梯度仍是 -12,而 .grad 会把新梯度加到已有值上。
这也解释了训练循环里的清理步骤:开始计算新一批数据前,调用优化器的 zero_grad(set_to_none=True),将它管理的参数梯度恢复为 None。接下来的反向再写入这一批梯度。
有意把多个小批次累积为一个大批次时,可以延后清理和更新,并按总样本数组织损失;这与忘记清理导致上一轮梯度混进来分别对应不同的代码顺序。先完成本章每批一次更新的流程,再尝试累积练习。