不能直接在 torch.Tensor 上写反向逻辑,因为 autograd 禁止手动修改 tensor.grad 以避免破坏计算图一致性;必须继承 torch.autograd.Function 并实现 forward 与 backward 静态方法,用 ctx.save_for_backward 保存中间变量,backward 返回梯度元组需与 forward 输入参数一一对应,且须使用原生 torch 运算。

为什么不能直接在 torch.Tensor 上写反向逻辑?
PyTorch 的 autograd 不允许你手动修改梯度张量(比如 tensor.grad)来“伪造”反向传播路径——这会破坏计算图的拓扑一致性,导致 backward() 报错或梯度静默丢失。真正可控的方式是继承 torch.autograd.Function,把前向和反向逻辑封装成一个原子操作。
如何定义一个带自定义反向的 Function?
必须同时实现 forward 和 backward 两个静态方法,且 backward 的输入必须与 forward 的输出数量、顺序严格对应(不是输入参数!)。关键点:
-
forward中用ctx.save_for_backward()保存需要在反向中复用的中间变量(如输入张量、mask、临时计算结果),否则反向里拿不到 -
backward返回的梯度元组,长度必须等于forward输入参数个数,每个元素对应相应输入的梯度(None表示该输入不可导) - 所有张量运算必须用
torch原生函数(如torch.sum),不能混用 NumPy 或原地操作(如.mul_())
示例:实现一个只对正数求导的 ReLU 变体:
class PositiveOnlyReLU(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
ctx.save_for_backward(x)
return torch.where(x > 0, x, torch.zeros_like(x))
<pre class="brush:php;toolbar:false;">@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
grad_input = torch.where(x > 0, grad_output, torch.zeros_like(grad_output))
return grad_input
调用自定义 Function 时为什么报 RuntimeError: Trying to backward through the graph a second time?
这是最常见的误用:直接调用 PositiveOnlyReLU.apply(x) 后又对 x 多次调用 backward(),而没清空计算图。根本原因在于 apply 返回的仍是叶子张量(leaf tensor),但它的梯度累加行为受 retain_graph 控制。正确做法:
- 确保每次
backward()都作用在**标量 loss** 上(非向量),除非显式传retain_graph=True - 如果要多次反向(如二阶梯度),必须在第一次
backward()时加retain_graph=True,否则图被释放 - 避免把
apply()结果赋给全局变量或反复复用——它绑定的是某次前向的上下文,重复使用可能触发缓存冲突
什么时候该用 torch.compile 或 torch.jit.script 包裹自定义 Function?
基本不用。PyTorch 2.x 的 torch.compile 目前不支持自定义 Function 的图融合,会直接 fallback 到 eager 模式;torch.jit.script 对 Function 支持极弱,容易因 ctx 或动态控制流报错。真要部署,更稳妥的做法是把自定义逻辑重构成纯 nn.Module + 标准算子(比如用 torch.nn.functional.relu 加条件掩码),仅在必须绕过标准梯度规则时才用 Function。
最易被忽略的一点:自定义 backward 中的数值稳定性。比如你在反向里写了 1 / x,但没检查 x 是否接近零——这不会在前向报错,却会让训练中途梯度爆炸,且难以定位。
Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!











