torch.func.hessian¶
- torch.func.hessian(func, argnums=0)¶
通过前向-反向策略计算
func
相对于索引argnum
处的参数的 Hessian。前向-反向策略(组合
jacfwd(jacrev(func))
)是获得良好性能的良好默认值。可以通过jacfwd()
和jacrev()
的其他组合(如jacfwd(jacfwd(func))
或jacrev(jacrev(func))
)来计算 Hessian。- 参数
- 返回
返回一个函数,该函数接受与
func
相同的输入,并返回func
相对于argnums
处的参数的 Hessian。
注意
您可能会看到此 API 因“运算符 X 未实现前向模式 AD”而出错。如果是这样,请提交错误报告,我们将优先处理它。另一种选择是使用
jacrev(jacrev(func))
,它具有更好的运算符覆盖范围。使用 R^N -> R^1 函数的基本用法会得到一个 N x N Hessian
>>> from torch.func import hessian >>> def f(x): >>> return x.sin().sum() >>> >>> x = torch.randn(5) >>> hess = hessian(f)(x) # equivalent to jacfwd(jacrev(f))(x) >>> assert torch.allclose(hess, torch.diag(-x.sin()))