Python 现代开发学习笔记

(十)梯度累积与每轮清理

# 项目文件: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。接下来的反向再写入这一批梯度。

有意把多个小批次累积为一个大批次时,可以延后清理和更新,并按总样本数组织损失;这与忘记清理导致上一轮梯度混进来分别对应不同的代码顺序。先完成本章每批一次更新的流程,再尝试累积练习。

同一权重两次重新前向,各自产生负十二梯度并累加为负二十四,清理成 None 后再反向重新得到负十二

笔记正文与全部题解 · Ctrl / ⌘ K