了解 Numba 性能差异。

Question

我试图理解使用不同的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

- mgilbert

Answer 1

在这里，复制数据并不起到很大的作用：瓶颈在于如何快速评估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的性能会受到影响。

然而，缓存未命中并不像计算函数那样重要：

%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％的运行时间。

---

顺便提一下，对于手写循环版本，我的 Numba 版本（0.50.1）能够矢量化并调用 MKL/SVML 功能。如果对于其他版本发生了这种情况，则 Numba 将返回 GNU 数学库功能，似乎这是在你的机器上发生的情况。

---

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)

- ead

Answer 2

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

我必须反对@ead。 让我们暂时假设:

主要性能差异在于tanh函数的评估

那么人们会期望，只运行带有快速数学运算的numpy和numba的tanh，将显示速度差异。

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并行运行和进行一些优化之后，性能提高虽然有些微小，但仍然存在。如下代码所示，NumExpr版本的运行时间为 2.56 ms，而之前的版本为 3.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)

实际解释

NumExpr相比于numba节约的约34%的时间很好，但更好的是它们有一个简明扼要的解释，为什么它们比numpy更快。我非常确定它也适用于numba。

来自NumExpr github页面：

NumExpr比NumPy性能更好的主要原因在于它避免了为中间结果分配内存。这提高了缓存利用率并减少了一般的内存访问。

因此，

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

由于需要进行许多步骤生成中间结果，因此速度较慢。

- Lukas S

---

感谢@user2640045也通过numexpr进行比较，我之前不知道这个选项。话虽如此，我的结果似乎与@ead上面提到的一致，即如果我放弃使用tanh，我的函数1b、1c和1d都显示出类似的性能。 - mgilbert

@mgilbert 请再次查看我的帖子。我刚刚比较了numpy和numba中的tanh，并得到了相同的结果。但是如果他的答案正确，那么numba版本应该更快。 - Lukas S

@user2640045 说得很有道理。然而，正如您的测量结果所示，numpy.tanh 占用了 17.4 毫秒中的 15.7 毫秒。因此，临时数组的成本并不大，无法解释差异。正如我所说，使用哪个版本的 tanh 取决于使用的 Python 发行版和版本。 - ead

虽然 numba 使用 svml，但 numexpr 将使用 tanh 的 vml 版本，这可能解释了 numba 和 numexpr 之间的差异。 - ead

我的猜测是你在使用Windows系统，因为在该系统上tanh实现比gcc更快。你的numpy没有使用vml，numba使用svml（在Windows上速度并不比较快），而numexpr使用vml，因此是最快的。 - ead