如何在PyTorch中高效计算多样本梯度（无for循环）？

问题分析 #

你遇到的RuntimeError: element 0 of tensors does not require grad并非因为输入x未开启梯度，而是torch.autograd.grad与vmap的配合问题：vmap批量执行时，无法正确追踪全局模型参数的梯度上下文，导致autograd.grad找不到需要计算梯度的有效张量。

解决方案：使用torch.func（原Functorch）的vmap+grad组合 #

PyTorch官方推荐用torch.func的函数式API配合vmap实现批量梯度计算，避免for循环的开销，同时解决梯度追踪问题。

完整修改代码 #

import torch
from torch import nn
from torch.func import vmap, grad, functional_call

device = "cpu"
input_size = 2
output_size = 2

# 定义模型
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Sequential(
            nn.Linear(input_size, 16, bias=False),
            nn.Linear(16, output_size, bias=False),
        )

    def forward(self, x):
        return self.linear(x)

# 初始化模型、损失函数与数据
model = NeuralNetwork().to(device)
loss_fn = nn.MSELoss(reduction='sum')  # 用sum保证单个样本损失独立计算

# 简化数据维度：去掉冗余的中间维度(1,)，vmap会自动识别第一个维度为批量维度
x = torch.randn(10, input_size).to(device)
y = torch.randn(10, output_size).to(device)

# 定义函数式损失计算：输入模型参数、单个样本数据与标签，返回对应损失
def compute_loss(params, x_single, y_single):
    pred = functional_call(model, params, (x_single,))
    return loss_fn(pred, y_single)

# 提取模型参数为字典格式，供functional_call使用
params = dict(model.named_parameters())

# 生成单个样本的梯度计算函数，再用vmap批量应用到所有样本
sample_grad_fn = grad(compute_loss)
batch_grads = vmap(sample_grad_fn)(params, x, y)

# 输出结果：每个参数的批量梯度形状为[样本数, 参数原有形状]
print(batch_grads['linear.0.weight'].shape)  # torch.Size([10, 16, 2])
print(batch_grads['linear.1.weight'].shape)  # torch.Size([10, 2, 16])

关键说明 #

1. 函数式模型调用：用functional_call在函数式上下文中调用模型，避免全局参数的梯度追踪冲突，让vmap能正确处理每个样本的独立计算。
2. 损失函数设置：reduction='sum'确保每个样本的损失是独立计算的，不会被平均，这样梯度对应单个样本的贡献。
3. vmap+grad组合：grad(compute_loss)生成计算损失对模型参数梯度的函数，vmap将该函数批量应用到所有样本上，实现无for循环的高效计算。

其他可选方案（不推荐） #

如果不愿使用torch.func，可以手动用torch.autograd.grad配合批量损失的reduction='none'，将所有样本的损失堆叠后计算梯度，但这种方式的效率远不如torch.func方案，且代码复杂度更高：

def compute_batch_grads(x, y):
    model.zero_grad()
    preds = model(x)
    losses = loss_fn(preds, y, reduction='none').sum(dim=-1)  # 每个样本的损失
    # 计算每个损失对参数的梯度
    grads = [torch.autograd.grad(loss, model.parameters(), retain_graph=True) for loss in losses]
    # 整理成批量梯度张量
    batch_grads = tuple(torch.stack([g[i] for g in grads]) for i in range(len(grads[0])))
    return batch_grads

batch_grads = compute_batch_grads(x, y)
print(batch_grads[0].shape)  # torch.Size([10, 16, 2])

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