torch.func.jacfwd¶

torch.func.jacfwd(func, argnums=0, has_aux=False, *, randomness='error')¶

计算func相对于索引为argnum的参数的前向模式自动微分的雅可比矩阵

Parameters

func (函数) – 一个Python函数，接受一个或多个参数，其中必须有一个是Tensor，并返回一个或多个Tensor
argnums (int 或 Tuple[int]) – 可选，整数或整数元组，指示要获取关于哪些参数的雅可比矩阵。默认值：0。
has_aux (bool) – 标志，指示 func 返回一个 (output, aux) 元组，其中第一个元素是要微分的函数的输出，第二个元素是不会被微分的辅助对象。默认值：False。
随机性 (str) – 标志，指示使用哪种类型的随机性。参见 vmap() 了解更多详情。允许的值：“different”、“same”、“error”。默认值：“error”

Returns

返回一个函数，该函数接受与func相同的输入，并返回func关于argnums的雅可比矩阵。如果has_aux is True，则返回的函数将返回一个(jacobian, aux)元组，其中jacobian是雅可比矩阵，aux是由func返回的辅助对象。

注意

您可能会看到此API错误显示“未为运算符X实现前向模式自动微分”。如果是这样，请提交错误报告，我们将优先处理。另一种选择是使用jacrev()，它具有更好的运算符覆盖率。

使用点对点的一元操作的基本用法将给出对角矩阵作为雅可比矩阵

>>> from torch.func import jacfwd
>>> x = torch.randn(5)
>>> jacobian = jacfwd(torch.sin)(x)
>>> expected = torch.diag(torch.cos(x))
>>> assert torch.allclose(jacobian, expected)

jacfwd() 可以与 vmap 组合以生成批量雅可比矩阵：

>>> from torch.func import jacfwd, vmap
>>> x = torch.randn(64, 5)
>>> jacobian = vmap(jacfwd(torch.sin))(x)
>>> assert jacobian.shape == (64, 5, 5)

如果您希望同时计算函数的输出以及函数的雅可比矩阵，请使用 has_aux 标志将输出作为辅助对象返回：

>>> from torch.func import jacfwd
>>> x = torch.randn(5)
>>>
>>> def f(x):
>>>   return x.sin()
>>>
>>> def g(x):
>>>   result = f(x)
>>>   return result, result
>>>
>>> jacobian_f, f_x = jacfwd(g, has_aux=True)(x)
>>> assert torch.allclose(f_x, f(x))

此外，jacrev() 可以与自身或 jacrev() 组合以生成 Hessian 矩阵

>>> from torch.func import jacfwd, jacrev
>>> def f(x):
>>>   return x.sin().sum()
>>>
>>> x = torch.randn(5)
>>> hessian = jacfwd(jacrev(f))(x)
>>> assert torch.allclose(hessian, torch.diag(-x.sin()))

默认情况下，jacfwd() 计算关于第一个输入的雅可比矩阵。然而，它可以通过使用 argnums 来计算关于不同参数的雅可比矩阵：

>>> from torch.func import jacfwd
>>> def f(x, y):
>>>   return x + y ** 2
>>>
>>> x, y = torch.randn(5), torch.randn(5)
>>> jacobian = jacfwd(f, argnums=1)(x, y)
>>> expected = torch.diag(2 * y)
>>> assert torch.allclose(jacobian, expected)

此外，将一个元组传递给 argnums 将计算相对于多个参数的雅可比矩阵

>>> from torch.func import jacfwd
>>> def f(x, y):
>>>   return x + y ** 2
>>>
>>> x, y = torch.randn(5), torch.randn(5)
>>> jacobian = jacfwd(f, argnums=(0, 1))(x, y)
>>> expectedX = torch.diag(torch.ones_like(x))
>>> expectedY = torch.diag(2 * y)
>>> assert torch.allclose(jacobian[0], expectedX)
>>> assert torch.allclose(jacobian[1], expectedY)