2

我有一个函数来计算我正在装饰的log gamma 函数numba.njit。

import numpy as np
from numpy import log
from scipy.special import gammaln
from numba import njit

coefs = np.array([
    57.1562356658629235, -59.5979603554754912,
    14.1360979747417471, -0.491913816097620199,
    .339946499848118887e-4, .465236289270485756e-4,
    -.983744753048795646e-4, .158088703224912494e-3,
    -.210264441724104883e-3, .217439618115212643e-3,
    -.164318106536763890e-3, .844182239838527433e-4,
    -.261908384015814087e-4, .368991826595316234e-5
])

@njit(fastmath=True)
def gammaln_nr(z):
    """Numerical Recipes 6.1"""
    y = z
    tmp = z + 5.24218750000000000
    tmp = (z + 0.5) * log(tmp) - tmp
    ser = np.ones_like(y) * 0.999999999999997092

    n = coefs.shape[0]
    for j in range(n):
        y = y + 1
        ser = ser + coefs[j] / y

    out = tmp + log(2.5066282746310005 * ser / z)
    return out

例如，当我gammaln_nr用于大型数组时np.linspace(0.001, 100, 10**7)，我的运行时间大约比 scipy 慢 7 倍（参见下面附录中的代码）。但是，如果我为任何单个值运行，我的 numba 函数总是快 2 倍左右。这是怎么回事？

z = 11.67
%timeit gammaln_nr(z)
%timeit gammaln(z)
>>> 470 ns ± 29.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
>>> 1.22 µs ± 28.3 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

我的直觉是，如果我的函数对于一个值更快，那么对于一组值应该更快。当然，情况可能并非如此，因为我不知道 numba 是使用 SIMD 指令还是其他类型的矢量化，而 scipy 可能是。

附录

import matplotlib.pyplot as plt
import seaborn as sns

n_trials = 8
scipy_times = np.zeros(n_trials)
fastats_times = np.zeros(n_trials)

for i in range(n_trials):
    zs = np.linspace(0.001, 100, 10**i) # evaluate gammaln over this range

    # dont take first timing - this is just compilation
    start = time.time()
    gammaln_nr(zs)
    end = time.time()

    start = time.time()
    gammaln_nr(zs)
    end = time.time()
    fastats_times[i] = end - start

    start = time.time()
    gammaln(zs)
    end = time.time()
    scipy_times[i] = end - start

fig, ax = plt.subplots(figsize=(12,8))
sns.lineplot(np.logspace(0, n_trials-1, n_trials), fastats_times, label="numba");
sns.lineplot(np.logspace(0, n_trials-1, n_trials), scipy_times, label="scipy");
ax.set(xscale="log");
ax.set_xlabel("Array Size", fontsize=15);
ax.set_ylabel("Execution Time (s)", fontsize=15);
ax.set_title("Execution Time of Log Gamma");

在此处输入图像描述

4

1 回答 1

4

在 Numba 中实现 gammaln

重新实现一些经常使用的函数可能需要相当多的工作，不仅要达到性能，还要获得明确定义的精度水平。所以直接的方法是简单地包装一个工作实现。

在gammalnscipy- 调用此函数的C 实现的情况下。因此，scipy 实现的速度还取决于编译 scipy 依赖项时使用的编译器和编译器标志。

一个值的性能结果可能与较大数组的结果有很大不同，这也不足为奇。在第一种情况下，调用开销（包括类型转换、输入检查等）占主导地位，在第二种情况下，实现的性能变得越来越重要。

改进您的实施

  • 编写显式循环。在 Numba 中，向量化操作被扩展为循环，然后 Numba 尝试加入循环。通常最好手动写出并加入这个循环。
  • 想想基本算术实现的差异。Python 总是检查除以 0 并在这种情况下引发异常，这是非常昂贵的。默认情况下，Numba 也使用此行为，但您也可以切换到 Numpy 错误检查。在这种情况下，除以 0 会导致 NaN。在进一步计算中处理 NaN 和 Inf -0/+0 的方式也受快速数学标志的影响。

代码

import numpy as np
from numpy import log
from scipy.special import gammaln
from numba import njit
import numba as nb

@njit(fastmath=True,error_model='numpy')
def gammaln_nr(z):
    """Numerical Recipes 6.1"""
    #Don't use global variables.. (They only can be changed if you recompile the function)
    coefs = np.array([
    57.1562356658629235, -59.5979603554754912,
    14.1360979747417471, -0.491913816097620199,
    .339946499848118887e-4, .465236289270485756e-4,
    -.983744753048795646e-4, .158088703224912494e-3,
    -.210264441724104883e-3, .217439618115212643e-3,
    -.164318106536763890e-3, .844182239838527433e-4,
    -.261908384015814087e-4, .368991826595316234e-5])

    out=np.empty(z.shape[0])


    for i in range(z.shape[0]):
      y = z[i]
      tmp = z[i] + 5.24218750000000000
      tmp = (z[i] + 0.5) * np.log(tmp) - tmp
      ser = 0.999999999999997092

      n = coefs.shape[0]
      for j in range(n):
          y = y + 1.
          ser = ser + coefs[j] / y

      out[i] = tmp + log(2.5066282746310005 * ser / z[i])
    return out

@njit(fastmath=True,error_model='numpy',parallel=True)
def gammaln_nr_p(z):
    """Numerical Recipes 6.1"""
    #Don't use global variables.. (They only can be changed if you recompile the function)
    coefs = np.array([
    57.1562356658629235, -59.5979603554754912,
    14.1360979747417471, -0.491913816097620199,
    .339946499848118887e-4, .465236289270485756e-4,
    -.983744753048795646e-4, .158088703224912494e-3,
    -.210264441724104883e-3, .217439618115212643e-3,
    -.164318106536763890e-3, .844182239838527433e-4,
    -.261908384015814087e-4, .368991826595316234e-5])

    out=np.empty(z.shape[0])


    for i in nb.prange(z.shape[0]):
      y = z[i]
      tmp = z[i] + 5.24218750000000000
      tmp = (z[i] + 0.5) * np.log(tmp) - tmp
      ser = 0.999999999999997092

      n = coefs.shape[0]
      for j in range(n):
          y = y + 1.
          ser = ser + coefs[j] / y

      out[i] = tmp + log(2.5066282746310005 * ser / z[i])
    return out


import matplotlib.pyplot as plt
import seaborn as sns
import time

n_trials = 8
scipy_times = np.zeros(n_trials)
fastats_times = np.zeros(n_trials)
fastats_times_p = np.zeros(n_trials)

for i in range(n_trials):
    zs = np.linspace(0.001, 100, 10**i) # evaluate gammaln over this range

    # dont take first timing - this is just compilation
    start = time.time()
    arr_1=gammaln_nr(zs)
    end = time.time()

    start = time.time()
    arr_1=gammaln_nr(zs)
    end = time.time()
    fastats_times[i] = end - start

    start = time.time()
    arr_3=gammaln_nr_p(zs)
    end = time.time()
    fastats_times_p[i] = end - start
    start = time.time()

    start = time.time()
    arr_3=gammaln_nr_p(zs)
    end = time.time()
    fastats_times_p[i] = end - start
    start = time.time()

    arr_2=gammaln(zs)
    end = time.time()
    scipy_times[i] = end - start
    print(np.allclose(arr_1,arr_2))
    print(np.allclose(arr_1,arr_3))

fig, ax = plt.subplots(figsize=(12,8))
sns.lineplot(np.logspace(0, n_trials-1, n_trials), fastats_times, label="numba");
sns.lineplot(np.logspace(0, n_trials-1, n_trials), fastats_times_p, label="numba_parallel");
sns.lineplot(np.logspace(0, n_trials-1, n_trials), scipy_times, label="scipy");
ax.set(xscale="log");
ax.set_xlabel("Array Size", fontsize=15);
ax.set_ylabel("Execution Time (s)", fontsize=15);
ax.set_title("Execution Time of Log Gamma");
fig.show()
于 2019-03-08T09:49:07.927 回答