提高 autograd jacobian 的性能

Question

我想知道以下代码如何更快。目前，它似乎异常缓慢，我怀疑我可能使用了 autograd API 错误。我期望的输出是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 函数，所以它应该执行得非常快，即使对于 40k 长度的数组也是如此。然而，这不是正在发生的事情。

即使像下面这样的东西也有同样糟糕的表现：

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)

Answer 1

从https://github.com/HIPS/autograd/issues/439我收集到有一个未记录的函数autograd.make_jvp可以用快进模式计算雅可比。

该链接指出：

给定函数 f，f 域中的向量 x 和 v，make_jvp(f)(x)(v)计算 f(x) 和 f 在 x 处评估的雅可比矩阵，右乘向量 v。

要获得 f 的完整雅可比行列式，您只需要编写一个循环来评估make_jvp(f)(x)(v)f 域的标准基中的每个 v。我们的反向模式雅可比运算符以相同的方式工作。

从你的例子：

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

这在 google collab 上运行约 0.005 秒。

编辑：

尚未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常数因子稍快。