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

Question

假设向量\theta是神经网络中的所有参数，我想知道如何在pytorch中计算\theta的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 .

我试过在pytorch 1.5中使用torch.autograd.functional.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 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

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

考虑到以下几点：

首先，关于torch.autograd.functional.hessian()第一个参数必须是一个函数，第二个参数应该是一个元组或张量列表。这意味着我们不能直接将标量损失传递给它。（我不知道为什么，因为我认为标量损失或 returns 标量的函数之间没有太大区别）
其次，我想得到一个完整的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 相同。注意参数a, b, c, d都被展开为one-dimensional向量，所以h中张量的形状都是二维的，顺序很好，很容易组合成一个大的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中的张量固定到相应的位置即可。

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

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

How to compute hessian matrix for all parameters in a network in pytorch?

machine-learning

hessian

pytorch