Нападающий Якобиан Нейронной Сети в Pytorch - Медленный - PullRequest
Новая
       5

Нападающий Якобиан Нейронной Сети в Pytorch - Медленный

0 голосов
/ 27 января 2019

Я вычисляю прямой якобиан (производная выходов по входам) двухслойной нейронной сети с прямой связью в pytorch, и мои результаты верны, но относительно медленны. Учитывая характер расчета, я ожидал бы, что он будет примерно таким же быстрым, как прямой проход по сети (или, может быть, в 2-3 раза дольше), но для выполнения шага оптимизации в этой подпрограмме требуется ~ 12x (в моем случае Тестовый пример Я просто хочу, чтобы jacobian = 1 во всех точках) по сравнению со стандартной среднеквадратичной ошибкой, поэтому я предполагаю, что делаю что-то неоптимальным образом. Мне просто интересно, если бы кто-нибудь знал более быстрый способ кодирования этого. Моя тестовая сеть имеет 2 входных узла, за которыми следуют 2 скрытых слоя по 5 узлов каждый и выходной слой из 2 узлов, и использует скрытые функции активации на скрытых слоях с линейным выходным слоем.

Якобианские вычисления основаны на статье Ограничения глубокого обучения в состязательных настройках , которая дает базовое рекурсивное определение прямой производной (в основном вы в конечном итоге умножаете производную ваших функций активации на весовые коэффициенты). и предыдущие частные производные каждого слоя). Это очень похоже на прямое распространение, поэтому я ожидаю, что оно будет быстрее, чем оно есть. Тогда определитель якобиана 2x2 в конце довольно прост.

Ниже приведен код сети и якобиана 

class Network(torch.nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        self.h_1_1 = torch.nn.Linear(input_1, hidden_1)
        self.h_1_2 = torch.nn.Linear(hidden_1, hidden_2)
        self.out = torch.nn.Linear(hidden_2, out_1)


    def forward(self, x):
        x = F.tanh(self.h_1_1(x))
        x = F.tanh(self.h_1_2(x))
        x = (self.out(x))
        return x

    def jacobian(self, x):
        a = self.h_1_1.weight
        x = F.tanh(self.h_1_1(x))
        tanh_deriv_tensor = 1 - (x ** 2)
        expanded_deriv = tanh_deriv_tensor.unsqueeze(-1).expand(-1, -1, input_1)
        partials = expanded_deriv * a.expand_as(expanded_deriv)

        a = torch.matmul(self.h_1_2.weight, partials)
        x = F.tanh(self.h_1_2(x))
        tanh_deriv_tensor = 1 - (x ** 2)
        expanded_deriv = tanh_deriv_tensor.unsqueeze(-1).expand(-1, -1, out_1)
        partials = expanded_deriv*a

        partials = torch.matmul(self.out.weight, partials)

        determinant = partials[:, 0, 0] * partials[:, 1, 1] - partials[:, 0, 1] * partials[:, 1, 0]
        return determinant

и вот две сравниваемые функции ошибок. Обратите внимание, что первый требует дополнительного прямого вызова через сеть, чтобы получить выходные значения (помеченные как действие), в то время как вторая функция этого не делает, поскольку она работает с входными значениями. 

def actor_loss_fcn1(action, target):
    loss = ((action-target)**2).mean()
    return loss

def actor_loss_fcn2(input): # 12x slower
    jacob = model.jacobian(input)
    loss = ((jacob-1)**2).mean()
    return loss

Любая идея по этому вопросу будет принята с благодарностью

1 Ответ

0 голосов
/ 15 июня 2019

Второе вычисление 'a' занимает больше всего времени на моей машине (CPU). 

# Here you increase the size of the matrix with a factor of "input_1"
expanded_deriv = tanh_deriv_tensor.unsqueeze(-1).expand(-1, -1, input_1)
partials = expanded_deriv * a.expand_as(expanded_deriv)

# Here your torch.matmul() needs to handle "input_1" times more computations than in a normal forward call
a = torch.matmul(self.h_1_2.weight, partials)

На моей машине время вычисления якобиана примерно равно времени, которое требуется факелу для вычисления 

a = torch.rand(hidden_1, hidden_2)
b = torch.rand(n_inputs, hidden_1, input_1)
%timeit torch.matmul(a,b)

Я не думаю, что это можно ускорить в вычислительном отношении. Если только вы не можете перейти с CPU на GPU, потому что GPU становится лучше на больших матрицах.

...