PyTorch网络中如何计算多输出对各输入的独立梯度（雅可比矩阵）？

• 第一种用torch.autograd.grad的实现，本质是对二维输出所有元素做全1权重的加权求和后再求梯度，返回的是聚合后的梯度值，无法拆分出每个独立输出分量对输入的偏导。
• 调用torch.autograd.functional.jacobian时的报错根因：
  • 该接口第一个入参要求传入可调用的前向计算函数，传入已经计算完成的输出张量y/net(x)时，就会触发'Tensor' object is not callable类型错误
  • 直接实例化网络调用net()时，没有给前向过程绑定输入x的传参映射，会触发前向函数缺少x位置参数的报错。

方法1：适配jacobian接口的入参要求 #

只要把网络的前向过程封装成一个仅接收输入张量、返回输出张量的可调用函数，就能正常调用接口计算雅可比：

# 封装符合接口入参要求的前向函数
def model_forward(x_input):
    return net(x_input)

# 计算完整雅可比
full_jacobian = torch.autograd.functional.jacobian(
    func=model_forward,
    inputs=x,
    create_graph=True
)

返回张量的形状完全匹配需求：[批量样本数, 输出维度, 输入维度]，本例中形状为(3, 2, 3)，索引对应关系为：

• full_jacobian[sample_idx, out_dim_idx, in_dim_idx] = 对应样本、对应输出分量对对应输入分量的独立偏导，完全覆盖你需要的所有梯度值。

方法2：手动逐分量计算雅可比 #

如果需要更灵活的控制计算图留存逻辑，可以逐输出通道构造梯度掩码，单独计算每个通道的梯度：

jacobian = []
# 遍历所有2个输出通道
for channel in range(y.shape[-1]):
    # 构造梯度掩码：仅当前遍历的通道位置为1，其余为0，避免多输出梯度求和
    channel_mask = torch.zeros_like(y)
    channel_mask[:, channel] = 1.0
    # 计算当前通道对输入的梯度
    channel_grad = torch.autograd.grad(
        outputs=y,
        inputs=x,
        grad_outputs=channel_mask,
        retain_graph=True
    )[0]
    jacobian.append(channel_grad)
# 堆叠调整维度到[batch, out_dim, in_dim]的目标形状
manual_jacobian = torch.stack(jacobian, dim=-1).permute(0, 2, 1)

按照你设置的全1权重、零偏置的网络参数，任意输出分量对输入分量的偏导恒为20：

• 第一层全连接3→10，权重全1，无偏置，输入对隐藏层的梯度为1
• ReLU层对所有正输入不改变梯度
• 第二层全连接10→2，权重全1，无偏置，隐藏层对输出的梯度为10
  最终偏导为1*10=20，和计算结果完全一致。

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