PyTorch中基于神经网络的ODE求解：输出对输入求导问题

你遇到的问题核心是autograd.grad的grad_output参数设置逻辑：当模型输出是多维度（比如[N,2]）时，默认用torch.ones_like(y)作为grad_output，会把每个样本的所有输出维度对输入的梯度做求和，最终得到[N,1]的结果，而不是每个输出维度单独的梯度。

要通过autograd.grad得到预期的[N,2]形状导数，需要针对每个输出维度单独计算梯度，再拼接起来。这里有两种实现方式：

方法1：循环遍历输出维度 #

适合输出维度较少的场景，代码直观：

import torch

# 假设model是1输入2输出的神经网络，t是形状[N,1]的输入张量
y = model(t)  # y形状[N,2]
dy_dt = []
for dim in range(y.shape[1]):
    # 构造仅当前输出维度为1，其余为0的grad_output
    grad_output = torch.zeros_like(y)
    grad_output[:, dim] = 1.0
    # 计算当前维度对输入的梯度
    grad = torch.autograd.grad(y, t, grad_output, create_graph=True)[0]
    dy_dt.append(grad)
# 拼接得到[N,2]的导数张量
dy_dt = torch.cat(dy_dt, dim=1)

方法2：批量构造grad_output（无循环） #

适合输出维度较多的场景，效率更高：

import torch

y = model(t)  # y形状[N,2]
# 构造形状为[N,2,2]的单位矩阵扩展张量，每个样本对应一个2x2单位矩阵
grad_outputs = torch.eye(y.shape[1], device=y.device).unsqueeze(0).repeat(y.shape[0], 1, 1)
# 计算所有输出维度的梯度，得到形状[N,2,1]的结果
dy_dt = torch.autograd.grad(y, t, grad_outputs, create_graph=True)[0]
# 挤压最后一维得到[N,2]
dy_dt = dy_dt.squeeze(-1)

原理说明 #

autograd.grad的grad_output参数本质是输出张量的梯度权重：

• 当传入torch.ones_like(y)时，相当于计算所有输出维度对输入梯度的加权和（权重全为1），因此每个样本只得到一个求和后的梯度值。
• 当传入单位矩阵形式的grad_output时，相当于分别计算每个输出维度对输入的梯度（对应维度权重为1，其余为0），最终得到每个输出维度的单独梯度。

这种方式比直接用torch.autograd.functional.jacobian更轻量化，不需要处理复杂的雅可比矩阵形状，更适合ODE求解器中对导数的需求。

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