如何在pytorch中计算网络中所有参数的hessian矩阵？

Question

假设向量\theta是神经网络中的所有参数，我想知道如何\theta在 pytorch 中计算Hessian矩阵。

假设网络如下：

class Net(Module):
    def __init__(self, h, w):
        super(Net, self).__init__()
        self.c1 = torch.nn.Conv2d(1, 32, 3, 1, 1)
        self.f2 = torch.nn.Linear(32 * h * w, 5)

    def forward(self, x):
        x = self.c1(x)
        x = x.view(x.size(0), -1)
        x = self.f2(x)
        return x

我知道可以通过调用torch.autograd.grad()两次来计算二阶导数，但是pytorch中的参数是由 组织的net.parameters()，我不知道如何计算所有参数的hessian。

我尝试torch.autograd.functional.hessian()在 pytorch 1.5 中使用如下：

import torch
import numpy as np
from torch.nn import Module
import torch.nn.functional as F


class Net(Module):
    def __init__(self, h, w):
        super(Net, self).__init__()
        self.c1 = torch.nn.Conv2d(1, 32, 3, 1, 1)
        self.f2 = torch.nn.Linear(32 * h * w, 5)

    def forward(self, x):
        x = self.c1(x)
        x = x.view(x.size(0), -1)
        x = self.f2(x)
        return x


def func_(a, b c, d):
    p = [a, b, c, d]
    x = torch.randn(size=[8, 1, 12, 12], dtype=torch.float32)
    y = torch.randint(0, 5, [8])
    x = F.conv2d(x, p[0], p[1], 1, 1)
    x = x.view(x.size(0), -1)
    x = F.linear(x, p[2], p[3])
    loss = F.cross_entropy(x, y)
    return loss


if __name__ == '__main__':
    net = Net(12, 12)

    h = torch.autograd.functional.hessian(func_, tuple([_ for _ in net.parameters()]))
    print(type(h), len(h))

h是一个元组，结果是奇怪的形状。例如，形状\frac{\delta Loss^2}{\delta c1.weight^2}为[32,1,3,3,32,1,3,3]。好像可以把它们组合成一个完整的H，但是不知道在整个Hessian Matrix中是哪个部分以及对应的顺序。

Answer 1

这是一个解决方案，我认为它有点过于复杂，但可能具有指导意义。

考虑以下几点：

1. 首先，关于torch.autograd.functional.hessian()第一个参数必须是一个函数，第二个参数应该是一个元组或张量列表。这意味着我们不能直接将标量损失传递给它。（我不知道为什么，因为我认为标量损失或返回标量的函数之间没有太大区别）
2. 其次，我想得到一个完整的Hessian矩阵，它是所有参数的二阶导数，并且应该是一个合适的顺序。

所以这里是解决方案：

import torch
import numpy as np
from torch.nn import Module
import torch.nn.functional as F

class Net(Module):
    def __init__(self, h, w):
        super(Net, self).__init__()
        self.c1 = torch.nn.Conv2d(1, 32, 3, 1, 1)
        self.f2 = torch.nn.Linear(32 * h * w, 5)

    def forward(self, x):
        x = self.c1(x)
        x = x.view(x.size(0), -1)
        x = self.f2(x)
        return x

def haha(a, b, c, d):
    p = [a.view(32, 1, 3, 3), b, c.view(5, 32 * 12 * 12), d]
    x = torch.randn(size=[8, 1, 12, 12], dtype=torch.float32)
    y = torch.randint(0, 5, [8])
    x = F.conv2d(x, p[0], p[1], 1, 1)
    x = x.view(x.size(0), -1)
    x = F.linear(x, p[2], p[3])
    loss = F.cross_entropy(x, y)
    return loss


if __name__ == '__main__':
    net = Net(12, 12)

    h = torch.autograd.functional.hessian(haha, tuple([_.view(-1) for _ in net.parameters()]))
    
    # Then we just need to fix tensors in h into a big matrix

我构建了一个haha与神经网络以相同方式工作的新函数Net。请注意，argumentsa, b, c, d都扩展为一维向量，因此 inh中的张量形状都是二维的，有序且易于组合成一个大的 Hessian 矩阵。

在我的例子中，张量的形状h是

# with relation to c1.weight and c1.weight, c1.bias, f2.weight, f2.bias
[288,288]
[288,32]
[288,23040]
[288,5]

# with relation to c2.bias and c1.weight, c1.bias, f2.weight, f2.bias
[32, 288]
[32, 32]
[32, 23040]
[32, 5]
...

所以很容易看出张量的含义以及它是哪一部分。我们需要做的就是分配一个(288+32+23040+5)*(288+32+23040+5)矩阵并将张量固定h到相应的位置。

我认为解决方案仍然可以改进，就像我们不需要构建一个与神经网络一样工作的函数，并两次转换参数的形状。但目前我没有更好的想法，如果有更好的解决方案，请告诉我。