Улучшить производительность автограда Якобиана - PullRequest
Новая
       4

Улучшить производительность автограда Якобиана

0 голосов
/ 02 февраля 2019

Мне интересно, как следующий код может быть быстрее.На данный момент это кажется неоправданно медленным, и я подозреваю, что я неправильно использую API Autograd.Я ожидаю вывод, что каждый элемент timeline вычисляется по якобиану f, который я получаю, но это занимает много времени: 

import numpy as np
from autograd import jacobian


def f(params):
    mu_, log_sigma_ = params
    Z = timeline * mu_ / log_sigma_
    return Z


timeline = np.linspace(1, 100, 40000)

gradient_at_mle = jacobian(f)(np.array([1.0, 1.0]))

Я ожидал бы следующее:

  1. jacobian(f) возвращает функцию, которая представляет вектор градиента по параметрам.
  2. jacobian(f)(np.array([1.0, 1.0])) - Якобиан, оцененный в точке (1, 1).Для меня это должно быть похоже на векторизованную функцию numpy, поэтому она должна выполняться очень быстро, даже для массивов длиной 40 КБ.Тем не менее, это не то, что происходит.

Даже что-то вроде следующего имеет такую ​​же низкую производительность: 

import numpy as np
from autograd import jacobian


def f(params, t):
    mu_, log_sigma_ = params
    Z = t * mu_ / log_sigma_
    return Z


timeline = np.linspace(1, 100, 40000)

gradient_at_mle = jacobian(f)(np.array([1.0, 1.0]), timeline)

1 Ответ

0 голосов
/ 04 февраля 2019

С https://github.com/HIPS/autograd/issues/439 Я понял, что есть недокументированная функция autograd.make_jvp, которая вычисляет якобиан в режиме ускоренной перемотки вперед.

Состояние ссылки:

Для данной функции f, векторов x и v в области f, make_jvp(f)(x)(v) вычисляет как f (x), так и якобиан из f, вычисленный в точке x, умноженный справа на вектор v.

Чтобы получить полноеJacobian of f, вам просто нужно написать цикл для оценки make_jvp(f)(x)(v) для каждого v в стандартном базисе области f.Наш оператор Якобиана обратного режима работает аналогичным образом.

Из вашего примера: 

import autograd.numpy as np
from autograd import make_jvp

def f(params):
    mu_, log_sigma_ = params
    Z = timeline * mu_ / log_sigma_
    return Z

timeline = np.linspace(1, 100, 40000)

gradient_at_mle = make_jvp(f)(np.array([1.0, 1.0]))

# loop through each basis
# [1, 0] evaluates (f(0), first column of jacobian)
# [0, 1] evaluates (f(0), second column of jacobian)
for basis in (np.array([1, 0]), np.array([0, 1])):
    val_of_f, col_of_jacobian = gradient_at_mle(basis)
    print(col_of_jacobian)

Вывод: 

[  1.           1.00247506   1.00495012 ...  99.99504988  99.99752494
 100.        ]
[  -1.           -1.00247506   -1.00495012 ...  -99.99504988  -99.99752494
 -100.        ]

Это выполняется за ~ 0,005 секундына Google Collab.

Редактировать:

Функции, подобные cdf, еще не определены для обычной jvp, но вы можете использовать другую недокументированную функцию make_jvp_reversemode там, где она определена.Использование аналогично, за исключением того, что выводом является только столбец, а не значение функции: 

import autograd.numpy as np
from autograd.scipy.stats.norm import cdf
from autograd.differential_operators import make_jvp_reversemode


def f(params):
    mu_, log_sigma_ = params
    Z = timeline * cdf(mu_ / log_sigma_)
    return Z

timeline = np.linspace(1, 100, 40000)

gradient_at_mle = make_jvp_reversemode(f)(np.array([1.0, 1.0]))

# loop through each basis
# [1, 0] evaluates first column of jacobian
# [0, 1] evaluates second column of jacobian
for basis in (np.array([1, 0]), np.array([0, 1])):
    col_of_jacobian = gradient_at_mle(basis)
    print(col_of_jacobian)

Выход: 

[0.05399097 0.0541246  0.05425823 ... 5.39882939 5.39896302 5.39909665]
[-0.05399097 -0.0541246  -0.05425823 ... -5.39882939 -5.39896302 -5.39909665]

Обратите внимание, что make_jvp_reversemode будет немного быстрее, чем make_jvp постоянным фактором из-за использования кэширования.

...