PyTorch中如何使用独立于反向传播的autograd计算单独函数的导数

问题原因 #

你遇到的报错是因为torch.autograd.functional.jacobian默认不会保留梯度计算过程的计算图，对x求导的结果J和参数theta的计算图被切断，后续计算loss时没有对应的梯度传播路径，就触发了does not have a grad_fn的错误。而你手写的解析导数是直接在theta的计算图上做运算，梯度链路完整，所以可以正常反向传播。

解决方案 #

PyTorch完全支持你需要的用法，只需要在调用jacobian时手动开启create_graph=True参数，让梯度计算过程本身也被记录到计算图中，就能保留对theta的梯度链路：

修改你的df函数即可：

def df(theta, x):
    # 新增create_graph=True参数，保留梯度计算的计算图
    J = jacobian(lambda x: f(theta, x), x, create_graph=True)
    return J

修改后运行完整代码，loss.backward()就能正常执行，打印theta.grad可以得到和解析导数完全一致的梯度结果tensor([8., 32.])。

你也可以用更底层的torch.autograd.grad实现相同逻辑，写法如下：

def df(theta, x):
    f_val = f(theta, x)
    # 对x求导，同时保留计算图
    J = torch.autograd.grad(outputs=f_val, inputs=x, create_graph=True)[0]
    return J

和JAX的差异说明 #

JAX的grad默认会保留高阶导数的计算链路，而PyTorch出于性能考虑，默认关闭了这个行为，需要手动传入create_graph=True显式开启，这不是框架限制，只是默认配置的差异。

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