PyTorch中如何获取中间张量梯度值

PyTorch获取中间张量梯度值的核心原理

PyTorch的autograd模块会构建动态计算图，在反向传播时默认只会保留最终输出张量对叶子节点的梯度，中间张量的梯度会在反向传播完成后被释放，这是为了避免不必要的内存占用。如果我们需要获取中间张量的梯度，就需要主动干预这个默认流程。

方法一：使用register_hook钩子函数

钩子函数是PyTorch提供的用于干预计算图梯度流程的接口，我们可以给中间张量注册一个钩子，在反向传播过程中捕获该张量的梯度值。

具体实现步骤如下：

• 先正常定义前向计算过程，得到中间张量
• 调用中间张量的register_hook方法，传入一个回调函数，回调函数会接收该张量的梯度作为输入
• 在回调函数中保存梯度值，之后就可以在反向传播后使用

下面是完整的代码示例：

import torch

# 定义输入张量，设置requires_grad为True，作为叶子节点
x = torch.tensor([2.0], requires_grad=True)
# 定义中间计算层
y = x * 3  # 中间张量
z = y ** 2  # 最终输出张量

# 定义保存梯度的变量
y_grad = None

# 给中间张量y注册钩子
def save_grad(grad):
    global y_grad
    y_grad = grad
    return grad  # 返回梯度，不影响原有梯度传播

hook_handle = y.register_hook(save_grad)

# 反向传播
z.backward()

# 打印中间张量y的梯度
print("中间张量y的梯度值：", y_grad)
# 打印叶子节点x的梯度，验证梯度传播正常
print("叶子节点x的梯度值：", x.grad)

# 移除钩子，避免后续计算受影响
hook_handle.remove()

这种方法的优势是不会改变原有计算图的梯度传播逻辑，只是额外捕获梯度，适合大多数需要临时获取中间梯度的场景。

方法二：设置张量的retain_grad属性

PyTorch的张量对象有一个retain_grad方法，调用该方法后，该张量的梯度会在反向传播后被保留，不会自动释放。

实现方式非常简单，只需要在得到中间张量后调用它的retain_grad方法即可，示例代码如下：

import torch

x = torch.tensor([2.0], requires_grad=True)
y = x * 3  # 中间张量
# 调用retain_grad保留梯度
y.retain_grad()
z = y ** 2

# 反向传播
z.backward()

# 直接访问y的grad属性获取梯度
print("中间张量y的梯度值：", y.grad)
print("叶子节点x的梯度值：", x.grad)

这种方法比钩子函数更简洁，但是需要注意，如果中间张量很多，全部调用retain_grad会增加内存占用，因为梯度会被一直保留直到张量被释放。

方法三：拆分计算图获取梯度

如果我们需要的中间张量是某个子计算的输出，也可以把这部分计算拆分成独立的计算图，单独计算梯度。

示例代码如下：

import torch

x = torch.tensor([2.0], requires_grad=True)
# 第一部分计算，得到中间张量y
y = x * 3
# 第二部分计算，得到最终输出z
z = y ** 2

# 单独对y计算梯度，此时把y当作叶子节点，grad_outputs是上游传来的梯度
y_grad = torch.autograd.grad(z, y, grad_outputs=torch.ones_like(z))[0]

print("中间张量y的梯度值：", y_grad)
# 再正常反向传播获取x的梯度
z.backward()
print("叶子节点x的梯度值：", x.grad)

这种方法适合需要灵活控制梯度计算流程的场景，但是需要手动处理上游梯度，逻辑相对复杂一些。

不同方法的适用场景对比

我们可以通过下面的表格快速选择适合自己场景的方法：

注意事项

• 只有requires_grad属性为True的张量才会有梯度，如果中间张量是由不需要梯度的张量计算得到的，那么无法获取它的梯度
• 钩子函数如果返回了新的梯度，会替换原有梯度，影响后续的梯度传播，所以如果只是想捕获梯度，回调函数直接返回输入的梯度即可
• 训练过程中如果多次迭代使用钩子，需要注意每次迭代前清空之前保存的梯度值，避免数据混淆

PyTorch中间张量梯度值autograd反向传播修改时间：2026-06-12 07:39:14