torch.autograd.functional.jacobian¶
- torch.autograd.functional.jacobian(func, inputs, create_graph=False, strict=False, vectorize=False, strategy='reverse-mode')[源代码]¶
计算给定函数的雅可比矩阵。
- 参数
func (函数) – 一个 Python 函数,它接受张量输入并返回一个张量元组或一个张量。
create_graph (布尔值,可选) – 如果为
True
,则将以可微方式计算雅可比矩阵。注意,当strict
为False
时,结果不能要求梯度或与输入断开连接。默认为False
。strict (布尔值,可选) – 如果为
True
,当检测到存在一个输入使得所有输出都独立于它时,将引发错误。如果为False
,我们将返回一个零张量作为该输入的雅可比矩阵,这是预期的数学值。默认为False
。vectorize (布尔值,可选) – 此功能尚处于实验阶段。如果您正在寻找更不实验且更高效的方法,请考虑使用
torch.func.jacrev()
或torch.func.jacfwd()
。计算雅可比矩阵时,通常我们对雅可比矩阵的每一行调用一次autograd.grad
。如果此标志为True
,我们将只执行一次autograd.grad
调用,其中batched_grad=True
使用 vmap 原型功能。虽然这在很多情况下都应该会带来性能改进,但由于此功能仍在实验阶段,因此可能会出现性能下降。有关更多信息,请参阅torch.autograd.grad()
的batched_grad
参数。strategy (字符串,可选) – 设置为
"forward-mode"
或"reverse-mode"
以确定将使用前向或反向模式 AD 计算雅可比矩阵。当前,"forward-mode"
要求vectorized=True
。默认为"reverse-mode"
。如果func
的输出多于输入,则"forward-mode"
往往性能更高。否则,建议使用"reverse-mode"
。
- 返回值
如果只有一个输入和输出,这将是一个包含线性化输入和输出的雅可比矩阵的单个张量。如果其中一个是元组,那么雅可比矩阵将是张量元组。如果两者都是元组,那么雅可比矩阵将是张量元组的元组,其中
Jacobian[i][j]
将包含第i
个输出和第j
个输入的雅可比矩阵,并且其大小将是相应输出和相应输入大小的串联,并且具有与相应输入相同的 dtype 和设备。如果策略是forward-mode
,dtype 将是输出的 dtype;否则,是输入的 dtype。- 返回类型
雅可比矩阵 (张量 或嵌套的张量元组)
示例
>>> 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.]]))