3

我想知道以下代码如何更快。目前，它似乎慢得不合理，我怀疑我可能错误地使用了 autograd API。我期望的输出是timeline在 f 的 jacobian 处评估的每个元素，我确实得到了，但这需要很长时间：

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)返回一个表示梯度向量 wrt 参数的函数。
  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)
4

1 回答 1

5

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

该链接指出：

给定一个函数 f，向量 x 和 v 在 f 的域中，make_jvp(f)(x)(v)计算 f(x) 和在 x 处求值的 f 的雅可比行列式，右乘向量 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

于 2019-02-04T02:29:22.937 回答