如何在批量输入输出场景下用torch.autograd.grad计算指定偏导数？

你的代码核心逻辑是正确的，能够获取到批量中每个样本的y1对x1、y2对x2的偏导数，但存在冗余代码和可优化的细节：

核心逻辑的正确性 #

你通过设置grad_outputs参数，分别指定只对y1、y2求导的方式是合理的：

• grad_output1[:, 1] = 0等价于给每个样本的y1分量赋予梯度权重1，y2分量权重0，因此dY1dX的每行对应[dy1/dx1, dy1/dx2]
• grad_output2[:, 0] = 0同理，dY2dX的每行对应[dy2/dx1, dy2/dx2]

只需提取对应列就能得到目标偏导数：

dy1dx1 = dY1dX[:, 0]  # 批量中每个样本y1对x1的偏导
dy2dx2 = dY2dX[:, 1]  # 批量中每个样本y2对x2的偏导

可优化的细节 #

1. 移除无用变量
   代码中的a10和a01未被使用，直接删除即可。

2. 简化输入梯度追踪
   PyTorch 0.4+版本已无需使用autograd.Variable，直接通过requires_grad_()方法设置梯度追踪更简洁：

   x_train = x_train.float().requires_grad_(True)
   

3. 按需调整梯度计算参数

   • 若不需要对偏导数本身再进行反向传播，可去掉create_graph=True，减少内存占用
   • 若两次梯度计算后无需保留计算图，可将第二次调用的retain_graph=True去掉

可选的高效实现方式 #

如果想减少autograd.grad的调用次数，可利用Jacobian矩阵直接提取目标元素（适合需要完整梯度矩阵的场景）：

def get_model_output(x):
    return model(x)

# 计算整个批量的Jacobian矩阵，形状为(100,2,100,2)
jacobian = torch.autograd.functional.jacobian(get_model_output, x_train)
# 提取每个样本的dy1/dx1和dy2/dx2
dy1dx1 = torch.diagonal(jacobian[:, 0, :, 0])
dy2dx2 = torch.diagonal(jacobian[:, 1, :, 1])

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