了解 Numba 性能差异

Understanding Numba Performance Differences

我试图通过使用算法的各种 numba 实现来了解我所看到的性能差异。特别是,我希望下面的 func1d 是最快的实现,因为它是唯一不复制数据的算法,但是从我的时间来看 func1b 似乎是最快的。

import numpy
import numba


def func1a(data, a, b, c):
    # pure numpy
    return a * (1 + numpy.tanh((data / b) - c))


@numba.njit(fastmath=True)
def func1b(data, a, b, c):
    new_data = a * (1 + numpy.tanh((data / b) - c))
    return new_data


@numba.njit(fastmath=True)
def func1c(data, a, b, c):
    new_data = numpy.empty(data.shape)
    for i in range(new_data.shape[0]):
        for j in range(new_data.shape[1]):
            new_data[i, j] = a * (1 + numpy.tanh((data[i, j] / b) - c)) 
    return new_data


@numba.njit(fastmath=True)
def func1d(data, a, b, c):
    for i in range(data.shape[0]):
        for j in range(data.shape[1]):
            data[i, j] = a * (1 + numpy.tanh((data[i, j] / b) - c)) 
    return data

用于测试内存复制的辅助函数

def get_data_base(arr):
    """For a given NumPy array, find the base array
    that owns the actual data.
    
    https://ipython-books.github.io/45-understanding-the-internals-of-numpy-to-avoid-unnecessary-array-copying/
    """
    base = arr
    while isinstance(base.base, numpy.ndarray):
        base = base.base
    return base


def arrays_share_data(x, y):
    return get_data_base(x) is get_data_base(y)


def test_share(func):
    data = data = numpy.random.randn(100, 3)
    print(arrays_share_data(data, func(data, 0.5, 2.5, 2.5)))

时间

# force compiling
data = numpy.random.randn(10_000, 300)
_ = func1a(data, 0.5, 2.5, 2.5)
_ = func1b(data, 0.5, 2.5, 2.5)
_ = func1c(data, 0.5, 2.5, 2.5)
_ = func1d(data, 0.5, 2.5, 2.5)

data = numpy.random.randn(10_000, 300)
%timeit func1a(data, 0.5, 2.5, 2.5)
%timeit func1b(data, 0.5, 2.5, 2.5)
%timeit func1c(data, 0.5, 2.5, 2.5)
%timeit func1d(data, 0.5, 2.5, 2.5)
67.2 ms ± 230 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
13 ms ± 10.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
69.8 ms ± 60.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
69.8 ms ± 105 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

测试哪些实现复制内存

test_share(func1a)
test_share(func1b)
test_share(func1c)
test_share(func1d)
False
False
False
True

在这里,数据的复制并没有起到很大的作用:瓶颈是如何评估 tanh-函数的速度很快。算法有很多种:有的快,有的慢,有的精确,有的不精确。

不同的 numpy 发行版使用不同的 tanh 函数实现,例如它可以是 mkl/vml 中的一个或 gnu-math-library 中的一个。

根据 numba 版本,也可以使用 mkl/svml 实现或 gnu-math-library。

查看内部的最简单方法是使用探查器,例如 perf

对于我机器上的 numpy 版本,我得到:

>>> perf record python run.py
>>> perf report
Overhead  Command  Shared Object                                      Symbol                                  
  46,73%  python   libm-2.23.so                                       [.] __expm1
  24,24%  python   libm-2.23.so                                       [.] __tanh
   4,89%  python   _multiarray_umath.cpython-37m-x86_64-linux-gnu.so  [.] sse2_binary_scalar2_divide_DOUBLE
   3,59%  python   [unknown]                                          [k] 0xffffffff8140290c

正如您所见,numpy 使用慢速 gnu-math-library (libm) 功能。

对于 numba 函数,我得到:

 53,98%  python   libsvml.so                                         [.] __svml_tanh4_e9
   3,60%  python   [unknown]                                          [k] 0xffffffff81831c57
   2,79%  python   python3.7                                          [.] _PyEval_EvalFrameDefault

这意味着使用了快速 mkl/svml 功能。

这就是(几乎)全部内容。


正如@user2640045 所言,numpy 的性能将因创建临时数组而导致额外的缓存未命中而受到损害。

然而,缓存未命中并没有像tanh的计算那样发挥这么大的作用:

%timeit func1a(data, 0.5, 2.5, 2.5)  # 91.5 ms ± 2.88 ms per loop 
%timeit numpy.tanh(data)             # 76.1 ms ± 539 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

即创建临时对象约占 运行 时间的 20%。


FWIW,对于带有手写循环的版本,我的 numba 版本 (0.50.1) 能够向量化并调用 mkl/svml 功能。如果对于其他版本没有发生这种情况 - numba 将回退到 gnu-math-library 功能,这似乎发生在你的机器上。


列表run.py

import numpy

# TODO: define func1b for checking numba
def func1a(data, a, b, c):
    # pure numpy
    return a * (1 + numpy.tanh((data / b) - c))


data = numpy.random.randn(10_000, 300)

for _ in range(100):
    func1a(data, 0.5, 2.5, 2.5)

性能差异不在 tanh 函数的评估中

我必须不同意@ead。让我们暂时假设

the main performance difference is in the evaluation of the tanh-function

然后人们会期望 运行宁 tanhnumpynumba 与快速数学将显示速度差异。

def func_a(data):
    return np.tanh(data)

@nb.njit(fastmath=True)
def func_b(data):
    new_data = np.tanh(data)
    return new_data

data = np.random.randn(10_000, 300)
%timeit func_a(data)
%timeit func_b(data)

然而在我的机器上,上面的代码在性能上几乎没有差异。

15.7 ms ± 129 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
15.8 ms ± 82 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

NumExpr

上绕道而行

我尝试了您代码的 NumExpr 版本。但在惊讶于它 运行 几乎快 7 倍之前,您应该记住它使用了我机器上所有可用的 10 个内核。在允许 numba 到 运行 并行并优化一点点之后,性能优势很小但仍然存在 2.56 ms3.87 ms。请参阅下面的代码。

@nb.njit(fastmath=True)
def func_a(data):
    new_data = a * (1 + np.tanh((data / b) - c))
    return new_data

@nb.njit(fastmath=True, parallel=True)
def func_b(data):
    new_data = a * (1 + np.tanh((data / b) - c))
    return new_data

@nb.njit(fastmath=True, parallel=True)
def func_c(data):
    for i in nb.prange(data.shape[0]):
        for j in range(data.shape[1]):
            data[i, j] = a * (1 + np.tanh((data[i, j] / b) - c)) 
    return data

def func_d(data):
    return ne.evaluate('a * (1 + tanh((data / b) - c))')

data = np.random.randn(10_000, 300)
%timeit func_a(data)
%timeit func_b(data)
%timeit func_c(data)
%timeit func_d(data)
17.4 ms ± 146 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.31 ms ± 193 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
3.87 ms ± 152 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
2.56 ms ± 104 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

实际解释

numba 相比,NumExpr 节省了约 34% 的时间,这很好,但更好的是,他们有一个简明的解释为什么它们比 numpy 快。我很确定这也适用于 numba

来自NumExpr github page

The main reason why NumExpr achieves better performance than NumPy is that it avoids allocating memory for intermediate results. This results in better cache utilization and reduces memory access in general.

所以

a * (1 + numpy.tanh((data / b) - c))

速度较慢,因为它执行了很多步骤来生成中间结果。