torch.func.jacfwd¶
- torch.func.jacfwd(func, argnums=0, has_aux=False, *, randomness='error')¶
使用前向模式自动微分计算
func
相对于索引为argnum
的参数的雅可比矩阵。- 参数
func (function) – 一个 Python 函数,它接受一个或多个参数,其中之一必须是一个张量,并返回一个或多个张量。
argnums (int 或 Tuple[int]) – 可选,整数或整数元组,表示要获取雅可比矩阵的相应参数。默认值:0。
has_aux (bool) – 标志,指示
func
是否返回一个(output, aux)
元组,其中第一个元素是函数要微分的输出,第二个元素是不会被微分的辅助对象。默认值:False。randomness (str) – 标志,指示要使用哪种类型的随机性。有关更多详细信息,请参见
vmap()
。允许的值:“different”、“same”、“error”。默认值:“error”。
- 返回值
返回一个函数,它接受与
func
相同的输入,并返回func
相对于索引为argnums
的参数的雅可比矩阵。如果has_aux is True
,则返回的函数将返回一个(jacobian, aux)
元组,其中jacobian
是雅可比矩阵,而aux
是func
返回的辅助对象。
注意
您可能会看到此 API 出现“forward-mode AD not implemented for operator 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)