如何在PyTorch中计算大型神经网络的Hessian值？

Question

如何在PyTorch中计算大型神经网络或变压器模型的Hessian矩阵？我知道torch.autograd.functional.hessian，但它似乎只计算函数的Hessian，而不是神经网络。我也在How to compute hessian matrix for all parameters in a network in pytorch?中看到了答案。问题是，我想计算关于权值的Hessian函数，但是对于大型神经网络来说，把它写成权值的函数是非常低效的。有更好的方法吗？如有任何建议，将不胜感激。谢谢。

Answer 1

过了一段时间，我终于找到了一个解决这个问题的新特性详细信息将在以下评论中描述：https://github.com/pytorch/pytorch/issues/49171#issuecomment-933814662。该解决方案使用函数torch.autograd.functional.hessian和新特性torch.nn.utils._stateless。注意，要使用这个新特性，您必须安装pytorch的夜间版本。