提高自组雅可比矩阵的性能

Question

我想知道下面的代码如何更快。目前，它似乎慢得离谱，我怀疑我可能用错了autogradapi。我期望的输出是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)返回一个表示梯度向量w.r.t.参数的函数。在
2. jacobian(f)(np.array([1.0, 1.0]))是在点（1，1）处求值的雅可比。对我来说，这应该像一个矢量化的numpy函数，所以它应该执行得非常快，即使对于40k长度的数组也是如此。然而，这并不是正在发生的事情。在

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

^{pr2}$

Answer 1

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

链接说明：

Given a function f, vectors x and v in the domain of f, make_jvp(f)(x)(v) computes both f(x) and the Jacobian of f evaluated at x, right multiplied by the vector v.

To get the full Jacobian of f you just need to write a loop to evaluate make_jvp(f)(x)(v) for each v in the standard basis of f's domain. Our reverse mode Jacobian operator works in the same way.

从你的例子来看：

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)

输出：

^{pr2}$

在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稍微快一点，因为它使用了缓存。在