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

Question

我有两个变量， x和theta 。 我试图仅在theta方面最小化我的损失，但作为我的损失函数的一部分，我需要不同函数（ f ）关于x 。 这个导数本身与最小化无关，只与它的输出有关。 但是，在 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)

张量（[2., 4.]）

张量（[2., 4.]）

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

给出以下错误

运行时错误：张量的元素 0 不需要 grad 并且没有 grad_fn

如果我提供一个解析导数（ 2*theta*x ），它工作正常：

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

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

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

聚苯乙烯

我想象的解决方案类似于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) – 如果为 True，雅可比行列式将以可微分的方式计算