PyTorch中如何为中间节点注册钩子以获取梯度？

问题背景 #

我有如下简单代码：

import torch

x = torch.tensor([4.0], requires_grad=True)
y = torch.tensor([2.0], requires_grad=True)
output = x * y + x / y

grad_x = torch.ones_like(output)
torch.autograd.grad(output, x, grad_outputs=grad_x)

对应的计算图如下：

现在我希望访问MulBackward0、DivBackward0和AddBackward0这些中间节点的梯度。我知道PyTorch默认不会存储这些梯度，也了解可以通过显式设置retain_grad等方式实现，但能否无需遍历计算图，直接为这些节点附加钩子？

解决方案 #

可以直接为这些反向节点附加钩子，不需要手动遍历计算图，核心是利用PyTorch的张量反向钩子或全局模块反向钩子机制，直接关联前向运算结果与对应反向节点。

方法1：为中间前向张量注册反向钩子 #

PyTorch中，前向运算的结果张量（比如乘法、除法、加法的输出）和对应的反向传播节点（MulBackward0、DivBackward0、AddBackward0）是一一绑定的。为这些中间张量注册register_hook，就能直接捕获对应反向节点的梯度：

import torch

x = torch.tensor([4.0], requires_grad=True)
y = torch.tensor([2.0], requires_grad=True)

# 拆分前向运算，保存中间结果张量
mul_result = x * y
div_result = x / y
add_result = mul_result + div_result

# 定义钩子函数，用于捕获并打印梯度
def capture_grad(grad):
    print(f"捕获到的梯度值: {grad}")
    return grad  # 可修改梯度，这里直接返回原梯度

# 为每个中间张量注册反向钩子
mul_hook = mul_result.register_hook(capture_grad)
div_hook = div_result.register_hook(capture_grad)
add_hook = add_result.register_hook(capture_grad)

# 执行反向传播计算
grad_x = torch.ones_like(add_result)
torch.autograd.grad(add_result, x, grad_outputs=grad_x)

# 用完钩子后移除，避免内存泄漏
mul_hook.remove()
div_hook.remove()
add_hook.remove()

运行后，钩子会依次打印出MulBackward0、DivBackward0、AddBackward0对应的梯度，无需手动遍历计算图。

方法2：全局反向钩子批量捕获 #

如果需要批量捕获所有反向节点的梯度，可以为所有PyTorch模块注册全局反向钩子：

import torch

# 定义全局反向钩子，捕获反向节点信息与梯度
def global_grad_hook(module, grad_input, grad_output):
    print(f"反向节点类型: {type(module).__name__}")
    print(f"输入梯度: {grad_input}, 输出梯度: {grad_output}\n")

# 注册全局模块反向钩子
torch.nn.modules.module.register_module_full_backward_hook(global_grad_hook)

x = torch.tensor([4.0], requires_grad=True)
y = torch.tensor([2.0], requires_grad=True)
output = x * y + x / y

grad_x = torch.ones_like(output)
torch.autograd.grad(output, x, grad_outputs=grad_x)

这种方式会自动捕获所有反向传播节点的梯度信息，包括目标的MulBackward0、DivBackward0、AddBackward0，完全不需要手动遍历计算图结构。

核心原理 #

前向运算的结果张量内部会记录对应的反向传播函数节点，通过为张量注册反向钩子，PyTorch会自动在反向传播时将对应节点的梯度传入钩子函数，从而实现无需遍历计算图就能捕获中间反向节点的梯度。

内容的提问来源于stack exchange，提问作者xotix