Numba - nopython 模式比对象模式慢？

Question

首先在这里使用 Numba。我读到 nopython 模式应该产生更快的代码，但是这个：

@jit(float64[:](int64[:], int64, float64), nopython=True)
def epsilon_bound(l, k, delta):
    return l/k+np.sqrt(np.log(1/delta))*np.sqrt(1/(2*k))

@jit(float64(float64, int64, float64), nopython=False)
def sim_bin(epsilon, sims, delta):
    k=10000
    s = np.random.binomial(k, epsilon, size=(sims,))
    print(nb.typeof(s))
    bound = epsilon_bound(s, k, delta)
    violations = np.greater(epsilon, bound)
    return np.sum(violations)/float(sims)

%%time
a = sim_bin(0.1, 1_000_000, 0.1)

运行得更快：

array(int64, 1d, C)
CPU times: user 66.7 ms, sys: 0 ns, total: 66.7 ms
Wall time: 65.7 ms

比这个：

@jit(float64[:](int64[:], int64, float64), nopython=True)
def epsilon_bound(l, k, delta):
    return l/k+np.sqrt(np.log(1/delta))*np.sqrt(1/(2*k))

@jit(float64(float64, int64, float64), nopython=True)
def sim_bin(epsilon, sims, delta):
    k=10000
    s = np.random.binomial(k, epsilon, size=(sims,))
    #print(nb.typeof(s))
    bound = epsilon_bound(s, k, delta)
    violations = np.greater(epsilon, bound)
    return np.sum(violations)/float(sims)

CPU times: user 4.94 s, sys: 8.02 ms, total: 4.95 s
Wall time: 4.93 s

运行 sim_bin.inspect_types() 表明第一个选项正在使用所有 pyobjects，而第二个选项正确推断所有类型。根据文档 (http://numba.pydata.org/numba-doc/0.31.0/glossary.html#term-nopython-mode)，nopython 模式应该会产生更快的代码。 有谁知道发生了什么？一定有充分的理由，但我不熟悉使用 Numba。是因为我主要使用向量化的 numpy 函数吗？

谢谢！！

Answer 1

该功能的一个主要瓶颈（至少对于 numba 0.31，Windows 10）似乎是np.random.binomial-call。当我测试它时：

@jit(nopython=True)
def nbbinom():
    return np.random.binomial(10000, 0.1, size=(1000000,))

nbbinom() # warmup
%timeit nbbinom()
# 1 loop, best of 3: 2.45 s per loop
%timeit np.random.binomial(10000, 0.1, size=(1000000,))
# 10 loops, best of 3: 23.1 ms per loop

但这可能取决于 numba 版本。 Numba（不幸的是）经常遭受（幸运的是）快速修复的性能回归。可能它只需要他们的错误跟踪器上的一个问题（如果尚未修复）。

但即便如此，您的代码仍包含大量矢量化操作。如果您使用矢量化操作，您将不会通过 numba 获得太多的速度提升。根据经验：如果您已经可以在不使用 python 循环的情况下使用 numpy 进行操作，则不需要 numba（也有例外。例如：我发现 numba 对于小型数组来说确实更快，其中numpy-ufunc 的开销很大）。

另一件事是创建随机数组比实际操作花费更长的时间（在 numpy 和 numba 中）：

import numpy as np
from numba import njit

@njit
def epsilon_bound1(l, k, delta):
    return l/k+np.sqrt(np.log(1/delta))*np.sqrt(1/(2*k))

def epsilon_bound2(l, k, delta):
    return l/k+np.sqrt(np.log(1/delta))*np.sqrt(1/(2*k))

def sim_bin(s, k, epsilon, sims, delta, func):
    bound = func(s, k, delta)
    violations = np.greater(epsilon, bound)
    return np.sum(violations)/float(sims)

epsilon = 0.1
sims = 1000000
delta = 0.1
k=10000
s = np.random.binomial(k, epsilon, size=(sims,))
%timeit np.random.binomial(k, epsilon, size=(sims,))
# 1 loop, best of 3: 232 ms per loop

sim_bin(s, k, 0.1, 1000000, 0.1, epsilon_bound1)  # warmup
%timeit sim_bin(s, k, 0.1, 1000000, 0.1, epsilon_bound1)
# 10 loops, best of 3: 28.5 ms per loop
%timeit sim_bin(s, k, 0.1, 1000000, 0.1, epsilon_bound2)
# 10 loops, best of 3: 37.6 ms per loop

因此，当您对 sim_bin 进行基准测试时，您实际上只对对 np.random.binomial 的调用进行基准测试，无论是相对快速的 numpy 还是（当前）相当慢的 numba 实现。