如何将 autograd 用于独立于 PyTorch 中的反向传播的单独函数？

Question

我有两个变量，x 和 theta。我试图仅将关于 theta 的损失降至最低，但作为损失函数的一部分，我需要关于 x 的不同函数 (f) 的导数。这个导数本身与最小化无关，只与它的输出有关。但是，在 PyTorch 中实现此功能时出现运行时错误。

一个最小的例子如下：

# minimal example of two different autograds
import torch

from torch.autograd.functional import jacobian
def f(theta, x):
    return torch.sum(theta * x ** 2)

def df(theta, x):
    J = jacobian(lambda x: f(theta, x), x)
    return J

# example evaluations of the autograd gradient
x = torch.tensor([1., 2.])
theta = torch.tensor([1., 1.], requires_grad = True)

# derivative should be 2*theta*x (same as an analytical)
with torch.no_grad():
    print(df(theta, x))
    print(2*theta*x)

tensor([2., 4.])

# define some arbitrary loss as a fn of theta
loss = torch.sum(df(theta, x)**2)
loss.backward()

给出如下错误

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

如果我提供解析导数 (2*theta*x)，它工作正常：

loss = torch.sum((2*theta*x)**2)
loss.backward()

有没有办法在 PyTorch 中做到这一点？或者我在某些方面受到限制？

如果有人需要更多详细信息，请告诉我。

PS

我想这个解决方案类似于 JAX 进行 autograd 的方式，因为这是我更熟悉的。我的意思是，在 JAX 中，我相信您会这样做：

from jax import grad
df = grad(lambda x: f(theta, x))

然后 df 只是一个可以在任何时候调用的函数。但是 PyTorch 是一样的吗？或者 .backward() 中是否存在导致此错误的冲突？

Answer 1

除非您明确要求，否则 PyTorch 的 jacobian 不会创建计算图

J = jacobian(lambda x: f(theta, x), x, create_graph=True)

.. 带有 create_graph 参数。

文档说的很清楚

create_graph (bool, optional) – If True, the Jacobian will be computed in a differentiable manner

如何将 autograd 用于独立于 PyTorch 中的反向传播的单独函数？

How do I use autograd for a separate function independent of backpropagate in PyTorch?

python

gradient

pytorch

autograd