Question

我想知道以下代码如何更快。目前，这似乎过慢，并且我怀疑我可能使用了错误的autograd API。我期望的输出是在得到f的雅可比时对timeline的每个元素求值的结果，但是我花了很长时间：

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

我期望以下几点：

jacobian(f)返回代表梯度向量w.r.t的函数。参数。
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，该函数可以使用快进模式计算jacobian。

链接状态：

给定一个函数f，向量f和f的域中的向量make_jvp(f)(x)(v)计算f（x）和在x处求值的f的雅可比行列，并乘以向量v。

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

根据您的示例：

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

提高autograd jacobian的性能

1 个答案: