Как вычислить матрицу гессиана для всех параметров в сети в pytorch?

Предположим, что вектор \theta - это все параметры в нейронной сети, мне интересно, как вычислить матрицу Гесси для \theta в pytorch.

Допустим, сеть выглядит следующим образом:

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(), и я не знаю, как вычислить гессиан для всех параметров.

Я пробовал использовать 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, но я не знаю, какая это часть во всей матрице Гессе и в соответствующем порядке.


pytorch machine-learning hessian
person david    schedule 23.09.2020    source источник
comment
Таким образом, мне нужно создать функцию, имеющую ту же функцию с классом Net, что слишком сложно. Есть ли решение попроще?   -  person david    schedule 23.09.2020
comment
Чем здесь полезна матрица Гессиана? Какую задачу он решит? Просто любопытно.   -  person Jodh Singh    schedule 25.09.2020
comment
@JodhSingh какие-то конкретные алгоритмы или анализ. Например, важность разных весов для потери тренированности.   -  person david    schedule 26.09.2020

Ответы (1)


arrow_upward
4
arrow_downward

Вот одно решение, я думаю, оно слишком сложное, но может быть поучительным.

Учитывая эти моменты:

  1. Во-первых, примерно torch.autograd.functional.hessian() первый аргумент должен быть функцией, а второй аргумент должен быть кортежем или списком тензоров. Это означает, что мы не можем напрямую передать ему скалярную потерю. (Я не знаю почему, потому что я думаю, что нет большой разницы между скалярной потерей и функцией, возвращающей скаляр)
  2. Во-вторых, я хочу получить полную матрицу Гессе, которая является второй производной всех параметров, и она должна быть в соответствующем порядке.

Итак, вот решение:

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 раскрываются в одномерные векторы, так что формы тензоров в h являются двумерными, в хорошем порядке и легко могут быть объединены в большую гессианскую матрицу.

В моем примере форма тензоров в 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 в соответствующих местах.

Я думаю, что решение все еще можно улучшить, например, нам не нужно создавать функцию, которая работает так же с нейронной сетью, и дважды трансформировать форму параметров. Но пока у меня нет лучших идей, если есть лучшее решение, дайте мне знать.

person david    schedule 23.09.2020