如何用PyTorch计算神经网络对输入与参数的梯度？

问题核心 #

你的网络输出是2维向量，而PyTorch的backward()默认只支持标量输入（因为它计算的是标量损失对变量的梯度），直接调用会报错。要计算向量输出对输入x和参数θ的梯度，需要针对每个输出分量单独求导，或者用torch.autograd.grad直接生成梯度矩阵。

完整实现代码 #

1. 初始化网络与输入 #

首先确保输入张量开启梯度追踪，否则无法计算对它的梯度：

import torch
import torch.nn as nn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(2, 20, dtype=float),
            nn.ReLU(),
            nn.Linear(20, 20, dtype=float),
            nn.ReLU(),
            nn.Linear(20, 2, dtype=float)
        )

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

M = NeuralNetwork().to(device)
# 开启输入的梯度追踪
input_x = torch.tensor([1.0, 0.0], device=device, requires_grad=True)

2. 计算对输入x的梯度（dM/dx） #

输出是2维向量，我们需要得到每个输出分量对输入x的梯度，最终形成2×2的矩阵（每行对应一个输出分量的梯度）：

方法一：分两次调用backward

# 先清空输入的梯度
input_x.grad = None

y = M(input_x)
# 计算第一个输出分量对x的梯度，retain_graph=True保留计算图
y[0].backward(retain_graph=True)
grad_x_0 = input_x.grad.clone()

# 清空梯度后计算第二个分量的梯度
input_x.grad = None
y[1].backward()
grad_x_1 = input_x.grad.clone()

# 组合成完整的梯度矩阵
jacobian_x = torch.stack([grad_x_0, grad_x_1])
print("dM/dx:\n", jacobian_x)

方法二：用torch.autograd.grad一键生成

这种方法更简洁，不用手动管理梯度和计算图：

y = M(input_x)
# grad_outputs传入单位矩阵，相当于分别对两个输出分量求导
jacobian_x = torch.autograd.grad(outputs=y, inputs=input_x, grad_outputs=torch.eye(2, device=device))[0]
print("dM/dx:\n", jacobian_x)

3. 计算对参数θ的梯度（dM/dθ） #

每个参数会对应两个梯度值（分别对应两个输出分量的梯度），可以用以下两种方式获取：

方法一：分两次调用backward

# 清空所有参数的梯度
M.zero_grad()

y = M(input_x)
# 计算第一个输出分量对参数的梯度
y[0].backward(retain_graph=True)
grad_params_0 = [p.grad.clone() for p in M.parameters()]

# 清空梯度后计算第二个分量的梯度
M.zero_grad()
y[1].backward()
grad_params_1 = [p.grad.clone() for p in M.parameters()]

# 把每个参数的两个梯度组合起来
combined_grad_params = []
for g0, g1 in zip(grad_params_0, grad_params_1):
    combined_grad_params.append(torch.stack([g0, g1]))

# 示例：查看第一个线性层权重的梯度形状（2, 20, 2）
print("第一个线性层权重的dM/dθ形状:", combined_grad_params[0].shape)

方法二：用torch.autograd.grad直接计算

y = M(input_x)
# 直接获取所有参数对应两个输出分量的梯度
grad_params = torch.autograd.grad(outputs=y, inputs=M.parameters(), grad_outputs=torch.eye(2, device=device))

# 示例：查看第一个线性层权重的梯度形状
print("第一个线性层权重的dM/dθ形状:", grad_params[0].shape)

关键注意点 #

• retain_graph=True：多次调用backward()时必须加，否则第一次求导后计算图会被销毁，无法继续求导。
• requires_grad=True：输入张量必须开启这个属性，否则PyTorch不会追踪对它的梯度。
• torch.autograd.grad：返回的是梯度元组，直接取第一个元素就是我们要的梯度矩阵/张量。

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