PyTorch如何对两个输入参数计算联合偏导数？

你当前调用的jacobian接口仅用于计算一阶偏导数，因此只能得到$\frac{\partial f}{\partial x}$、$\frac{\partial f}{\partial y}$两个结果，无法直接得到二阶混合偏导$\frac{\partial^2 f}{\partial x \partial y}$，可以参考以下两种方案修改代码：

方案1：嵌套调用autograd.grad手动求二阶偏导 #

这个方法逻辑更直观，适合只需要计算特定混合偏导的场景：

import torch
def function(x,y):
    f = x**3 + y**3
    return f

a = torch.tensor([4., 5., 6.], requires_grad=True)
b = torch.tensor([1., 2., 6.], requires_grad=True)

# 第一次求导得到一阶偏导 ∂f/∂x，create_graph=True保留计算图用于二次求导
df_dx = torch.autograd.grad(function(a,b), a, create_graph=True)[0]
# 第二次对y求导得到混合偏导 ∂²f/(∂x∂y)
d2f_dxdy = torch.autograd.grad(df_dx.sum(), b)[0]

print(d2f_dxdy)

运行后会输出全0的张量，和数学推导结果一致：$f=x^3+y3$对x求一阶导为$3x^2$，再对y求导结果恒为0。

方案2：调用hessian接口提取混合偏导 #

如果需要一次性计算所有二阶偏导，可以直接调用Hessian计算接口，再提取对应混合偏导部分：

import torch
def function(x,y):
    f = x**3 + y**3
    return f

a = torch.tensor([4., 5., 6.], requires_grad=True)
b = torch.tensor([1., 2., 6.], requires_grad=True)

# 计算完整Hessian矩阵
full_hessian = torch.autograd.functional.hessian(lambda params: function(params[0], params[1]), (a,b))
# 提取∂²f/∂x∂y对应的子矩阵，取对角元素即为每个输入位置的混合偏导
d2f_dxdy = torch.diag(full_hessian[0][1])

print(d2f_dxdy)

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