《PyTorch 是怎么跟踪计算依赖的，顺便分享下我平时防显存泄漏的写法》

在 PyTorch 里跑模型训练时，前向传播只管照着公式往前算，等到调一句 loss.backward()，所有参数的梯度就自动算好了。这套自动求导机制之所以知道该沿着什么路径往回传梯度，全靠张量身上的两个核心属性在背后记账，一个是 requires_grad，另一个是 grad_fn。
直接用 torch.tensor() 新建普通张量时，requires_grad 默认是 False，代表不需要对它求导，而模型里的权重参数 nn.Parameter 默认都是 True。只要参与运算的输入张量里有一个带了 requires_grad=True，算出来的新张量也会自动染上 requires_grad=True。

import torch

w = torch.tensor([2.0], requires_grad=True)
x = torch.tensor([3.0])
y = w * x + 1.0

print(w.is_leaf, w.grad_fn)  # True None
print(y.is_leaf, y.grad_fn)  # False <AddBackward0 object ...>

跑上面这段代码就能看出源头张量和中间张量的区别。像 w 这种直接创建出来、不依赖其他张量计算的源头变量，在计算图里叫叶子张量，is_leaf 属性是 True，身上的 grad_fn 是 None。而 y 是由乘法和加法一步步算出来的中间结果，is_leaf 变成了 False，同时它的 grad_fn 记录了生成它的最后一步反向传播函数 AddBackward0，如果顺着 y.grad_fn.next_functions 往回摸，还能摸到上一步乘法的 MulBackward0 以及指向叶子张量 w 的梯度累加节点。
前向传播每算一步，PyTorch 就在后台实时拼出一个有向无环图，同时为了待会儿反向传播算梯度时有数可用，图里的节点还会把前向传播的中间激活值拽在显存里。比如算乘法 w * x 对 w 的导数需要用到 x 的值。哪怕前向代码已经走到了最后一层，前面各层的中间张量只要还挂在这张图上，显存就绝对不会释放。
正常情况下，调完 loss.backward() 算出梯度后，PyTorch 会自动把这张图里缓存的中间激活值全部清空，把显存腾出来给下一个 batch 用。但写训练和验证代码时只要没注意切断这套依赖跟踪，显存很快就会爆掉。
我在写训练循环统计整个 epoch 的平均损失时，以前踩过一个很隐蔽的坑，就是在循环里直接写 total_loss += loss。因为每个 batch 算出来的 loss 本身是个带 grad_fn 的张量，直接用加号把它累加到 total_loss 上，等于在计算图末尾又接了一个新的 AddBackward0 节点，导致就算调完了 loss.backward()，由于 Python 变量 total_loss 还拽着整条计算图的尾巴，每一个 batch 的计算图引用都释放不掉，跑几十个步长显存就会一路涨到 OOM。现在我凡是做日志打印或者指标累加，都会固定写成 total_loss += loss.item()，先把单元素张量转成纯粹的 Python 浮点数，彻底剥离掉它背后的计算图引用。
跑验证集评估或者线上推理时，我还会固定套上 with torch.no_grad() 上下文。光不调反向传播没用。有时候觉得验证阶段反正不跑 loss.backward()，加不加这行好像都能跑通，但实际上只要模型的权重参数还带着 requires_grad=True，前向传播照样会在后台老老实实记录 grad_fn 并把每一层的中间激活值全塞进显存里备着，直到当前 batch 的输出变量生命周期结束才回收。单次前向的峰值显存会白白翻倍。套上 torch.no_grad() 后，引擎从源头关掉了依赖记录，算出来的张量 grad_fn 全是 None，前向算完一层就能立刻复用上一层的显存。
如果是在生成对抗网络或者多阶段训练里，需要把前一个模块算出来的张量继续留在 GPU 上传给下一个模块用，又不想让梯度顺着这个张量传回前一个模块，转成 Python 标量的 .item() 就不适用了。这种场景我会调 tensor.detach()，它不会复制底层显存数据，而是直接返回一个共享同一块存储空间的新张量，只是把新张量的 requires_grad 设成了 False，grad_fn 也置为 None，相当于把计算图从这个节点直接剪断，前面分支的中间缓存就能及时释放掉。