快捷方式

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 (函数) – 一个 Python 函数,它接受一个或多个参数,其中一个必须是张量,并返回一个或多个张量

  • argnums (intTuple[int]) – 可选,整数或整数元组,表示相对于哪个参数获取 Hessian。默认值:0。

返回

返回一个函数,该函数接受与 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()))

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源