如何在PyTorch中获取向量输入函数的Jacobian及正确使用torch.autograd.grad

问题背景 #

自定义的derivative函数仅在输出为标量或多个独立标量时能返回梯度之和，但当输出是张量时，无法得到完整的雅可比矩阵（输出每个分量对输入每个分量的梯度构成的矩阵）。

核心原因 #

torch.autograd.grad默认会将所有输出元素的梯度求和后返回，而非保留每个输出分量对应的梯度信息。要获取完整雅可比，需要针对每个输出分量单独计算梯度，或利用grad_outputs参数批量处理。

方法1：循环计算每个输出分量的梯度 #

通过遍历输出张量的每个元素，单独计算其对输入的梯度，最后堆叠成雅可比矩阵。注意需要设置retain_graph=True以保留计算图，避免后续求导失败。

示例代码：

import torch

def f(x):
    # 定义一个输出为3维张量的函数
    return torch.stack([x[0]**2, x[0]+x[1], x[1]*x[2]])

def compute_jacobian(y, x):
    jacobian_list = []
    # 遍历每个输出分量
    for i in range(y.shape[0]):
        # 单独计算第i个输出分量对输入x的梯度
        grad = torch.autograd.grad(y[i], x, retain_graph=True)[0]
        jacobian_list.append(grad)
    # 堆叠为(output_dim, input_dim)形状的雅可比矩阵
    return torch.stack(jacobian_list)

# 测试
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = f(x)
jacobian = compute_jacobian(y, x)
print(jacobian)

输出结果：

tensor([[2., 0., 0.],
        [1., 1., 0.],
        [0., 3., 2.]])

方法2：利用grad_outputs参数批量计算 #

通过设置grad_outputs为单位矩阵，让torch.autograd.grad一次性计算所有输出分量的梯度，避免循环。单位矩阵的每个行对应一个输出分量的权重，确保每个分量的梯度被单独保留。

示例代码：

def compute_jacobian_batch(y, x):
    # 创建与输出y同形状的单位矩阵
    grad_outputs = torch.eye(y.shape[0], device=y.device, dtype=y.dtype)
    # 计算梯度，返回结果为(input_dim, output_dim)，转置后得到标准雅可比矩阵
    jacobian = torch.autograd.grad(y, x, grad_outputs=grad_outputs, retain_graph=True)[0].T
    return jacobian

# 测试
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = f(x)
jacobian = compute_jacobian_batch(y, x)
print(jacobian)

输出结果与方法1完全一致。

方法3：使用torch.func.jacfwd（PyTorch 2.0+） #

PyTorch 2.0及以上版本提供了torch.func.jacfwd（前向模式自动微分），可以直接计算函数的雅可比矩阵，无需手动处理计算图或循环，代码更简洁高效。

示例代码：

from torch.func import jacfwd

def f(x):
    return torch.stack([x[0]**2, x[0]+x[1], x[1]*x[2]])

# 测试
x = torch.tensor([1.0, 2.0, 3.0])
jacobian = jacfwd(f)(x)
print(jacobian)

输出结果同样与前两种方法一致。

补充：原函数的行为说明 #

当输出y是张量时，直接调用torch.autograd.grad(y, x)会返回所有输出分量梯度的和。例如上述测试中，原derivative函数会返回tensor([3., 4., 2.])，这是2x0 +1 +0、0+1 +x2、0+0 +x1的结果，并非完整雅可比矩阵。

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