PyTorch中如何更简便计算输出各元素对参数矩阵的梯度？

我有一个由N×M大小矩阵参数化的PyTorch网络，输出尺寸为N×1。想要计算网络输出对其参数的导数，期望得到维度为N×N×M的结果（每个参数对应N个导数，对应每个输出元素），但调用output.backward(torch.ones_like(output))后，.grad属性中的梯度仅为N×M维度。

目前的解决方案是遍历output的每个元素，逐个调用.backward()并堆叠梯度，但每次调用后需要手动清零参数的.grad属性，操作繁琐。有没有更简洁高效的方法？

示例代码与期望结果 #

import torch
import torch.nn as nn

class Network(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Network, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = Network(4, 5, 2)
inputs = torch.randn(1, 4)
output = model(inputs)
output.backward(torch.ones_like(output))
gradients = [param.grad for param in model.parameters()]

• 现有结果：gradients是包含[5×4张量、5×1张量、2×5张量、2×1张量]的列表
• 期望结果：包含[2×5×4张量、2×5×1张量、2×2×5张量、2×2×1张量]的列表

直接用torch.autograd.grad函数可以一次性计算所有输出元素对参数的梯度，无需循环调用.backward()，也不用手动清零梯度。该函数支持传入多个输出目标，并返回每个目标对参数的梯度组成的列表。

实现代码 #

import torch
import torch.nn as nn

class Network(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Network, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = Network(4, 5, 2)
inputs = torch.randn(1, 4)
output = model(inputs)

# 获取所有可训练参数
params = list(model.parameters())
# 计算每个输出元素对所有参数的梯度：将输出拆分为单个元素作为目标
gradients = torch.autograd.grad(outputs=output.unbind(), inputs=params)

# 调整bias梯度维度以匹配期望格式（可选）
adjusted_gradients = []
for grad in gradients:
    if len(grad.shape) == 2:
        adjusted_gradients.append(grad.unsqueeze(-1))
    else:
        adjusted_gradients.append(grad)

# 打印验证维度
for grad in adjusted_gradients:
    print(grad.shape)

结果说明 #

运行上述代码后，adjusted_gradients列表的元素维度完全符合期望：

• adjusted_gradients[0]：torch.Size([2, 5, 4])（对应fc1的weight）
• adjusted_gradients[1]：torch.Size([2, 5, 1])（对应fc1的bias）
• adjusted_gradients[2]：torch.Size([2, 2, 5])（对应fc2的weight）
• adjusted_gradients[3]：torch.Size([2, 2, 1])（对应fc2的bias）

原理说明 #

• output.unbind()将输出张量按第一个维度拆分为单个元素的元组（示例中为2个元素）
• torch.autograd.grad会为每个输出元素计算其对所有输入参数的梯度，返回的列表中，每个元素的第一个维度对应输出元素的索引，后续维度匹配参数的原始维度
• 该方法无需手动管理.grad属性的清零，计算效率远高于循环调用.backward()

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