PyTorch多标签分类：如何单反向传播计算多输出对批量输入的梯度？

完全可以通过单次反向传播实现你的需求，核心是利用torch.autograd.grad的is_grads_batched参数，结合对每个样本正标签对应输出的求和操作完成。

实现步骤 #

1. 筛选正标签对应输出并求和：对每个样本，将模型输出中与真实标签为1的位置对应的值求和，得到形状为(batch_size,)的张量——每个元素代表对应样本所有正标签输出的总和。
2. 单次反向传播计算梯度：使用torch.autograd.grad并设置is_grads_batched=True，让autograd自动计算每个样本的求和输出对对应输入样本的梯度。

完整代码示例 #

import torch

# 假设model已定义并处于可用状态
model = ...
inputs = torch.randn((4,3,224,224), requires_grad=True)  # 需开启requires_grad
targets = torch.tensor([[1,0,1],[1,0,0],[0,0,1],[1,1,0]])

outputs = model(inputs)

# 1. 计算每个样本正标签对应输出的总和
mask = targets == 1
sum_pos_outputs = (outputs * mask).sum(dim=1)  # 形状: (4,)

# 2. 单次反向传播计算梯度
grads = torch.autograd.grad(sum_pos_outputs, inputs, is_grads_batched=True)[0]
# grads形状: (4, 3, 224, 224)，grads[i]即为input[i]对应的所有正标签输出对它的梯度之和

参数说明 #

is_grads_batched=True的作用是告诉PyTorch：我们需要计算每个目标元素（sum_pos_outputs[i]）对输入中对应样本（inputs[i]）的梯度，而非所有目标对整个输入张量的梯度。这正好匹配批量计算需求，避免了循环每个样本做反向传播的开销。

与单标签场景的对比 #

单标签场景中你通过gather提取单个标签的输出，再unbind后计算梯度；多标签场景下通过求和将每个样本的多个正标签输出合并为一个标量目标，再利用is_grads_batched参数直接完成批量梯度计算，逻辑更简洁且效率更高。

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