PyTorch中高效计算图像分类器各类别单独梯度的方法求助

问题背景 #

需要为PyTorch图像分类器单独计算每个类别的梯度，初始实现采用循环调用torch.autograd.grad并设置retain_graph=True：

outputs = net(inputs)[0] # 仅考虑batch中的第一个样本
grads = [torch.autograd.grad(outputs[i], inputs, retain_graph=True) 
         for i in range(len(outputs))]

但torch.autograd.grad文档指出：

Note that in nearly all cases setting this option (retain_graph) to True is not needed and often can be worked around in a much more efficient way

尝试Bing AI建议的代码时出现形状不匹配错误：

identity = torch.eye(len(outputs))
outputs.backward(gradient=identity)

报错信息：

RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([9, 9]) and outputs has a shape of torch.Size([9]).

（注：所用分类器包含9个类别）

高效实现方法 #

方法一：使用torch.autograd.grad批量计算 #

利用torch.autograd.grad的grad_outputs参数批量传入梯度，一次性完成所有类别的梯度计算，无需循环和retain_graph=True：

# 确保输入张量开启梯度追踪
inputs.requires_grad_(True)

# 获取输出（单样本，形状为[num_classes]）
outputs = net(inputs.unsqueeze(0))[0]

# 生成对应每个类别的one-hot梯度矩阵，形状为[num_classes, num_classes]
grad_outputs = torch.eye(len(outputs), device=outputs.device)

# 一次性计算所有类别的梯度，返回结果形状为[num_classes, *inputs.shape]
grads = torch.autograd.grad(outputs, inputs, grad_outputs=grad_outputs)[0]

# 若需按类别单独提取梯度，可直接索引
class_0_grad = grads[0]
class_1_grad = grads[1]

此方法通过一次反向传播完成所有计算，效率远高于循环调用。

方法二：使用torch.autograd.functional.jacobian #

直接计算输出向量关于输入的雅可比矩阵，矩阵的每一行对应一个类别输出的梯度：

# 确保输入张量开启梯度追踪
inputs.requires_grad_(True)

# 定义适配单样本的输出计算函数
def get_single_sample_output(x):
    # 为输入添加batch维度，适配模型输入要求
    return net(x.unsqueeze(0))[0]

# 计算雅可比矩阵，形状为[num_classes, *inputs.shape]
jacobian = torch.autograd.functional.jacobian(get_single_sample_output, inputs)

# 按行提取每个类别的梯度
grads = [jacobian[i] for i in range(jacobian.shape[0])]

该方法无需手动处理梯度参数，PyTorch会自动高效计算整个雅可比矩阵。

错误原因说明 #

Bing AI建议的代码报错是因为outputs形状为[9]，而传入的torch.eye(9)形状为[9,9]，两者形状不匹配。backward()的gradient参数要求与输出张量形状一致（或可广播），而批量计算梯度更适合用torch.autograd.grad的grad_outputs参数实现。

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