如何使用 autograd.grad 计算 PyTorch 中参数损失的 Hessian 矩阵

Question

我知道有很多关于 pytorch 中“计算 Hessian 矩阵”的内容，但据我所知，我还没有发现任何对我有用的内容。因此，为了尽可能精确，我想要的 Hessian 矩阵是损失相对于网络参数的梯度的雅可比矩阵。也称为参数的二阶导数矩阵。

我发现一些代码可以以直观的方式工作，尽管速度不应该很快。显然，它只是计算参数与参数之间的损失梯度的梯度，并且一次只计算一个（梯度的）元素。我认为逻辑绝对是正确的，但我收到了一个错误，与requires_grad. 我是 pytorch 初学者，所以也许这是一件简单的事情，但错误似乎是说它不能接受变量的梯度env_grads，这是之前 grad 函数调用的输出。

任何对此的帮助将不胜感激。这是错误消息后面的代码。我还打印出了该env_grads[0]变量，因此我们可以看到它实际上是一个张量，这是上一次调用的正确输出grad。

env_loss = loss_fn(env_outputs, env_targets)
total_loss += env_loss
env_grads = torch.autograd.grad(env_loss, params,retain_graph=True)

print(env_grads[0])
hess_params = torch.zeros_like(env_grads[0])
for i in range(env_grads[0].size(0)):
    for j in range(env_grads[0].size(1)):
        hess_params[i, j] = torch.autograd.grad(env_grads[0][i][j], params, retain_graph=True)[0][i, j] #  <--- error here
print(hess_params)

输出：

tensor([[-6.4064e-03, -3.1738e-03,  1.7128e-02,  8.0391e-03],
        [ 7.1698e-03, -2.4640e-03, -2.2769e-03, -1.0687e-03],
        [-3.0390e-04, -2.4273e-03, -4.0799e-02, -1.9149e-02],
        ...,
        [ 1.1258e-02, -2.5911e-05, -9.8133e-02, -4.6059e-02],
        [ 8.1502e-04, -2.5814e-03,  4.1772e-02,  1.9606e-02],
        [-1.0075e-02,  6.6072e-03,  8.3118e-04,  3.9011e-04]], device='cuda:0')

错误：

Traceback (most recent call last):
  File "/home/jefferythewind/anaconda3/envs/rapids3/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/home/jefferythewind/anaconda3/envs/rapids3/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/jefferythewind/Projects/Irina/learning-explanations-hard-to-vary/and_mask/run_synthetic.py", line 258, in <module>
    main(args)
  File "/home/jefferythewind/Projects/Irina/learning-explanations-hard-to-vary/and_mask/run_synthetic.py", line 245, in main
    deep_mask=args.deep_mask
  File "/home/jefferythewind/Projects/Irina/learning-explanations-hard-to-vary/and_mask/run_synthetic.py", line 103, in train
    scale_grad_inverse_sparsity=scale_grad_inverse_sparsity
  File "/home/jefferythewind/Projects/Irina/learning-explanations-hard-to-vary/and_mask/and_mask_utils.py", line 154, in get_grads_deep
    hess_params[i, j] = torch.autograd.grad(env_grads[0][i][j], params, retain_graph=True)[0][i, j]
  File "/home/jefferythewind/anaconda3/envs/rapids3/lib/python3.7/site-packages/torch/autograd/__init__.py", line 157, in grad
    inputs, allow_unused)
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

Answer 1

PyTorch 最近添加了一个功能性更高级别的 API ，它torch.autograd提供torch.autograd.functional.hessian(func, inputs,...)直接评估标量函数func 相对于指定位置处的参数的hessian 矩阵inputs，该位置是与 的参数相对应的张量元组func。hessian我相信它本身不支持自动微分。

但请注意，截至 2021 年 3 月，它仍处于测试阶段。

torch.autograd.functional.hessian用于创建非零均值分数测试的完整示例（作为单样本 t 测试的（坏）替代方案）：

import numpy as np
import torch, torchvision
from torch.autograd import Variable, grad
import torch.distributions as td
import math
from torch.optim import Adam
import scipy.stats


x_data = torch.randn(100)+0.0 # observed data (here sampled under H0)

N = x_data.shape[0] # number of observations

mu_null = torch.zeros(1)
sigma_null_hat = Variable(torch.ones(1), requires_grad=True)

def log_lik(mu, sigma):
  return td.Normal(loc=mu, scale=sigma).log_prob(x_data).sum()

# Find theta_null_hat by some gradient descent algorithm (in this case an closed-form expression would be trivial to obtain (see below)):
opt = Adam([sigma_null_hat], lr=0.01)
for epoch in range(2000):
    opt.zero_grad() # reset gradient accumulator or optimizer
    loss = - log_lik(mu_null, sigma_null_hat) # compute log likelihood with current value of sigma_null_hat  (= Forward pass)
    loss.backward() # compute gradients (= Backward pass)
    opt.step()      # update sigma_null_hat
    
print(f'parameter fitted under null: sigma: {sigma_null_hat}, expected: {torch.sqrt((x_data**2).mean())}')
#> parameter fitted under null: sigma: tensor([0.9260], requires_grad=True), expected: 0.9259940385818481

theta_null_hat = (mu_null, sigma_null_hat)

U = torch.tensor(torch.autograd.functional.jacobian(log_lik, theta_null_hat)) # Jacobian (= vector of partial derivatives of log likelihood w.r.t. the parameters (of the full/alternative model)) = score
I = -torch.tensor(torch.autograd.functional.hessian(log_lik, theta_null_hat)) / N # estimate of the Fisher information matrix
S = torch.t(U) @ torch.inverse(I) @ U / N # test statistic, often named "LM" (as in Lagrange multiplier), would be zero at the maximum likelihood estimate

pval_score_test = 1 - scipy.stats.chi2(df = 1).cdf(S) # S asymptocially follows a chi^2 distribution with degrees of freedom equal to the number of parameters fixed under H0
print(f'p-value Chi^2-based score test: {pval_score_test}')
#> p-value Chi^2-based score test: 0.9203232752568568

# comparison with Student's t-test:
pval_t_test = scipy.stats.ttest_1samp(x_data, popmean = 0).pvalue
print(f'p-value Student\'s t-test: {pval_t_test}')
#> p-value Student's t-test: 0.9209265268946605

Answer 2

如果我在这里回答我自己的问题，请不要介意。我从 PyTorch 网站的这个线程中找到了我需要的提示，大约在一半的位置。

\n

\n

这不会\xe2\x80\x99t工作并且符合我提到的第三点。您需要\n需要使用可微分运算创建计算图，\n这将创建一个具有有效 grad_fn 的结果张量。

\n

\n

autograd.grad我注意到调用有一个参数create_graph，所以我True在第一次调用时将其设置为grad将其设置为，并最终解决了错误。

\n

修改后的工作代码：

\n

env_loss = loss_fn(env_outputs, env_targets)\ntotal_loss += env_loss\nenv_grads = torch.autograd.grad(env_loss, params, retain_graph=True, create_graph=True)\n\nprint(env_grads[0])\nhess_params = torch.zeros_like(env_grads[0])\nfor i in range(env_grads[0].size(0)):\n    for j in range(env_grads[0].size(1)):\n        hess_params[i, j] = torch.autograd.grad(env_grads[0][i][j], params, retain_graph=True)[0][i, j] #  <--- error here\nprint(hess_params)\n

\n