Comment calculer la matrice hessienne pour tous les paramètres d'un réseau dans pytorch ?

Question

Supposons que le vecteur \theta est l'ensemble des paramètres d'un réseau neuronal, je me demande comment calculer la matrice hessienne pour \theta en pytorch.

Supposons que le réseau soit le suivant :

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

Je sais que la dérivée seconde peut être calculée en appelant torch.autograd.grad() deux fois, mais les paramètres dans pytorch sont organisés par net.parameters() et je ne sais pas comment calculer la hessienne pour tous les paramètres.

J'ai essayé d'utiliser torch.autograd.functional.hessian() dans pytorch 1.5 comme suit :

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 est un tuple, et les résultats ont une forme étrange. Par exemple, la forme de \frac{\delta Loss^2}{\delta c1.weight^2} est [32,1,3,3,32,1,3,3] . Il semble que je puisse les combiner en une H mais je ne sais pas de quelle partie il s'agit dans l'ensemble de la matrice hessienne et l'ordre correspondant.

Answer 1

Voici une solution, je pense qu'elle est un peu trop complexe mais elle pourrait être instructive.

Réfléchissez à ces points :

1. Tout d'abord, à propos de torch.autograd.functional.hessian() le premier argument doit être une fonction et le second un tuple ou une liste de tenseurs. Cela signifie que nous ne pouvons pas lui passer directement une perte scalaire. (Je ne sais pas pourquoi, car je pense qu'il n'y a pas de grande différence entre une perte scalaire et une fonction qui retourne un scalaire).
2. Deuxièmement, je veux obtenir une matrice hessienne complète, qui est la dérivée seconde de tous les paramètres, et elle doit être dans un ordre approprié.

Voici donc la solution :

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

Je crée une nouvelle fonction haha qui fonctionne de la même manière que le réseau neuronal Net . Remarquez que les arguments a, b, c, d sont tous développés en vecteurs unidimensionnels, de sorte que les formes des tenseurs en h sont tous bidimensionnels, en bon ordre et faciles à combiner dans une grande matrice de Hessian.

Dans mon exemple, les formes des tenseurs en h es

# 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]
...

Il est donc facile de comprendre la signification des tenseurs et de savoir de quelle partie il s'agit. Tout ce que nous avons à faire, c'est d'allouer un (288+32+23040+5)*(288+32+23040+5) et fixer les tenseurs dans h dans les emplacements correspondants.

Je pense que la solution peut encore être améliorée, par exemple nous n'avons pas besoin de construire une fonction qui fonctionne de la même manière que le réseau neuronal, et de transformer la forme des paramètres deux fois. Mais pour l'instant je n'ai pas de meilleures idées, s'il y a une meilleure solution, n'hésitez pas à me le faire savoir.