Python PyTorch jacobian用法及代码示例

本文简要介绍python语言中 torch.autograd.functional.jacobian 的用法。

用法:

torch.autograd.functional.jacobian(func, inputs, create_graph=False, strict=False, vectorize=False)

参数：

• func(函数) -一个 Python 函数，它接受张量输入并返回张量元组或张量。

• inputs(张量元组或者Tensor) -函数 func 的输入。

• create_graph(bool,可选的) -如果 True ，雅可比将以可微分的方式计算。请注意，当 strict 为 False 时，结果不能需要梯度或与输入断开连接。默认为 False 。

• strict(bool,可选的) -如果 True ，当我们检测到存在一个输入使得所有输出都独立于它时，将引发错误。如果 False ，我们返回一个零张量作为所述输入的雅可比，这是预期的数学值。默认为 False 。

• vectorize(bool,可选的) -此函数是实验性的，请自行承担使用风险。在计算雅可比时，通常我们在雅可比的每一行调用一次autograd.grad。如果此标志是 True ，我们使用 vmap 原型函数作为后端来矢量化对 autograd.grad 的调用，因此我们只调用它一次而不是每行一次。这应该会在许多用例中带来性能改进，但是，由于此函数不完整，可能会出现性能悬崖。请使用torch._C._debug_only_display_vmap_fallback_warnings(True) 显示任何性能警告，如果您的用例存在警告，请向我们提交问题。默认为 False 。

返回：

如果只有一个输入和输出，这将是一个包含线性化输入和输出的雅可比行列式的张量。如果两者中的一个是元组，那么雅可比将是张量的元组。如果它们都是元组，那么雅可比将是张量元组的元组，其中 Jacobian[i][j] 将包含第 i 输出和 j 输入的雅可比，并且将具有大小的串联相应的输出和相应的输入将具有与相应输入相同的 dtype 和设备。

返回类型：

雅可比行列式(Tensor 或张量的嵌套元组)

计算给定函数的雅可比行列式的函数。

示例

>>> def exp_reducer(x):
...   return x.exp().sum(dim=1)
>>> inputs = torch.rand(2, 2)
>>> jacobian(exp_reducer, inputs)
tensor([[[1.4917, 2.4352],
         [0.0000, 0.0000]],
        [[0.0000, 0.0000],
         [2.4369, 2.3799]]])

>>> jacobian(exp_reducer, inputs, create_graph=True)
tensor([[[1.4917, 2.4352],
         [0.0000, 0.0000]],
        [[0.0000, 0.0000],
         [2.4369, 2.3799]]], grad_fn=<ViewBackward>)

>>> def exp_adder(x, y):
...   return 2 * x.exp() + 3 * y
>>> inputs = (torch.rand(2), torch.rand(2))
>>> jacobian(exp_adder, inputs)
(tensor([[2.8052, 0.0000],
        [0.0000, 3.3963]]),
 tensor([[3., 0.],
         [0., 3.]]))