PyTorch中输出对输入的梯度计算异常：形状不符问题排查

为什么梯度形状和input一致？ #

你使用torch.autograd.grad(outputs=output, inputs=input, grad_outputs=None)得到的结果是符合PyTorch默认行为的：

• 当grad_outputs设为None时，PyTorch会自动将output的所有元素求和为一个标量，再计算这个标量对input的梯度。
• 标量对张量的梯度形状必然和原张量（input）一致，所以你得到的结果是正确的，只是和你的预期需求不匹配。

如何得到与output逻辑对应的梯度？ #

你想要的本质是**output每个元素对input的导数**（即Jacobian矩阵），因为output是(number_of_rays, number_of_points_per_ray, 300)的张量，每个元素对input的导数是(3,)的向量，最终你需要的实际形状应该是(number_of_rays, number_of_points_per_ray, 300, 3)（每个output元素对应input对应位置的3维导数）。

结合你提到的output[i]仅依赖input[i]的结构，可以通过以下方式实现：

方法1：使用torch.autograd.functional.jacobian

直接计算整个张量的Jacobian，再利用结构简化形状：

import torch

# 示例inference函数（替换为你的模型）
def inference(input_tensor):
    # 模拟每个ray的output仅依赖对应ray的input
    linear = torch.nn.Linear(3, 300)
    return linear(input_tensor)

# 初始化输入张量
num_rays = 10
num_points_per_ray = 20
input = torch.randn(num_rays, num_points_per_ray, 3, requires_grad=True)

# 计算Jacobian，形状为(num_rays, num_points_per_ray, 300, num_rays, num_points_per_ray, 3)
jacobian = torch.autograd.functional.jacobian(inference, input)

# 利用output[i]仅依赖input[i]的特性，提取有效部分，最终形状为(num_rays, num_points_per_ray, 300, 3)
simplified_jac = jacobian.diagonal(dim1=0, dim2=3).diagonal(dim1=0, dim2=2)

方法2：用vmap批量优化计算

因为每个ray的计算独立，使用torch.func.vmap可以高效地对每个ray单独计算Jacobian：

from torch.func import vmap

# 定义单ray的推理函数
def single_ray_inference(ray_input):
    # ray_input形状为(num_points_per_ray, 3)
    linear = torch.nn.Linear(3, 300)
    return linear(ray_input)

# 批量计算每个ray的Jacobian，形状为(num_rays, num_points_per_ray, 300, num_points_per_ray, 3)
jac_per_ray = vmap(lambda x: torch.autograd.functional.jacobian(single_ray_inference, x))(input)

# 若每个点的output仅依赖对应点的input，进一步简化为(num_rays, num_points_per_ray, 300, 3)
jac_per_point = jac_per_ray.diagonal(dim1=1, dim2=3)

关键总结 #

• torch.autograd.grad默认计算的是标量损失对输入的梯度，所以形状与输入一致；
• 若要获取每个输出元素对输入的导数，需要计算Jacobian矩阵，可通过torch.autograd.functional.jacobian结合你的模型结构优化实现。

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