各种 numpy 花式索引方法的性能,也与 numba

Performance of various numpy fancy indexing methods, also with numba

因为对于我的程序来说,Numpy 数组的快速索引是非常必要的,而考虑到性能,花哨的索引并没有很好的声誉,我决定进行一些测试。特别是因为 Numba 发展得非常快,我尝试了哪些方法适用于 numba。

作为输入,我一直在使用以下数组进行小型数组测试:

import numpy as np
import numba as nb

x = np.arange(0, 100, dtype=np.float64)  # array to be indexed
idx = np.array((0, 4, 55, -1), dtype=np.int32)  # fancy indexing array
bool_mask = np.zeros(x.shape, dtype=np.bool)  # boolean indexing mask
bool_mask[idx] = True  # set same elements as in idx True
y = np.zeros(idx.shape, dtype=np.float64)  # output array
y_bool = np.zeros(bool_mask[bool_mask == True].shape, dtype=np.float64)  #bool output array (only for convenience)

以及以下用于我的大型数组测试的数组(此处需要 y_bool 来处理来自 randint 的欺骗数字):

x = np.arange(0, 1000000, dtype=np.float64)
idx = np.random.randint(0, 1000000, size=int(1000000/50))
bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[idx] = True
y = np.zeros(idx.shape, dtype=np.float64)
y_bool = np.zeros(bool_mask[bool_mask == True].shape, dtype=np.float64)

这会在不使用 numba 的情况下产生以下时间:

%timeit x[idx]
#1.08 µs ± 21 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
#large arrays: 129 µs ± 3.45 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit x[bool_mask]
#482 ns ± 18.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
#large arrays: 621 µs ± 15.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.take(x, idx)
#2.27 µs ± 104 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
# large arrays: 112 µs ± 5.76 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit np.take(x, idx, out=y)
#2.65 µs ± 134 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
# large arrays: 134 µs ± 4.47 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit x.take(idx)
#919 ns ± 21.3 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 108 µs ± 1.71 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit x.take(idx, out=y)
#1.79 µs ± 40.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# larg arrays: 131 µs ± 2.92 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit np.compress(bool_mask, x)
#1.93 µs ± 95.8 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 618 µs ± 15.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.compress(bool_mask, x, out=y_bool)
#2.58 µs ± 167 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
# large arrays: 637 µs ± 9.88 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit x.compress(bool_mask)
#900 ns ± 82.4 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 628 µs ± 17.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit x.compress(bool_mask, out=y_bool)
#1.78 µs ± 59.8 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 628 µs ± 13.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.extract(bool_mask, x)
#5.29 µs ± 194 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
# large arrays: 641 µs ± 13 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

并且使用 numba,在 nopython 模式下使用 jitting,caching 和 nogil 我修饰了索引方式,[=19 支持这些方式=]:

@nb.jit(nopython=True, cache=True, nogil=True)
def fancy(x, idx):
    x[idx]

@nb.jit(nopython=True, cache=True, nogil=True)
def fancy_bool(x, bool_mask):
    x[bool_mask]

@nb.jit(nopython=True, cache=True, nogil=True)
def taker(x, idx):
    np.take(x, idx)

@nb.jit(nopython=True, cache=True, nogil=True)
def ndtaker(x, idx):
    x.take(idx)

对于小型和大型阵列,这会产生以下结果:

%timeit fancy(x, idx)
#686 ns ± 25.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 84.7 µs ± 1.82 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit fancy_bool(x, bool_mask)
#845 ns ± 31 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 843 µs ± 14.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit taker(x, idx)
#814 ns ± 21.1 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 87 µs ± 1.52 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit ndtaker(x, idx)
#831 ns ± 24.5 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# large arrays: 85.4 µs ± 2.69 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

总结

虽然对于没有 numba 的 numpy 来说,很明显小数组最好用布尔掩码索引(与 ndarray.take(idx) 相比大约是一个因子 2),对于较大的数组 ndarray.take(idx) 将表现最好,在这种情况下,比布尔索引快 6 倍左右。盈亏平衡点的数组大小约为 1000 个单元格,索引数组大小约为 20 个单元格。
对于具有 1e5 个元素和 5e3 索引数组大小的数组,ndarray.take(idx) 将比布尔掩码索引快 10 倍。因此,布尔索引似乎随着数组大小的增加而显着减慢,但在达到某个数组大小阈值后会稍微赶上。

对于 numba jitted 函数,除了布尔掩码索引之外,所有索引函数都有一个小的加速。简单的花式索引在这里效果最好,但仍然比没有 jitting 的布尔掩码慢。
对于较大的数组,布尔掩码索引比其他方法慢很多,甚至比非 jitted 版本慢。其他三种方法都表现得很好,比非 jitted 版本快大约 15%。

对于我有许多不同大小的数组的情况,使用 numba 进行花式索引是最好的方法。或许其他人也能从这篇冗长的文章中找到一些有用的信息post.

编辑:
对不起,我忘了问我的问题,事实上我已经问了。我只是在工作日结束时快速输入这个,然后完全忘记了它...... 好吧,你知道比我测试过的那些更好更快的方法吗?使用 Cython,我的时间安排在 Numba 和 Python.
之间 由于索引数组被预定义一次并且在长时间迭代中无需更改即可使用,因此任何预定义索引过程的方法都会很棒。为此,我考虑使用步幅。但我无法预先定义一组自定义步幅。是否可以使用 strides 获得内存的预定义视图?

编辑 2:
我想我会把关于预定义常量索引数组的问题移到一个新的更具体的问题上,这些数组将在相同的值数组(其中只有值发生变化但形状不变)上迭代几百万次。这个问题太笼统了,也许我提出的问题也有点误导。我一打开新问题就post这里link!
Here is the link to the followup question.

您的总结不完全正确,您已经对不同大小的数组进行了测试,但您没有做的一件事是更改索引元素的数量。

我将其限制为纯索引并省略了 take(实际上是整数数组索引)以及 compressextract(因为它们实际上是布尔数组索引)。这些的唯一区别是常数因子。方法 takecompress 的常数因子将小于 numpy 函数 np.takenp.compress 的开销,但对于大小合理的数组,其影响可以忽略不计。

让我用不同的数字来表示:

# ~ every 500th element
x = np.arange(0, 1000000, dtype=np.float64)
idx = np.random.randint(0, 1000000, size=int(1000000/500))  # changed the ratio!
bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[idx] = True

%timeit x[idx]
# 51.6 µs ± 2.02 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit x[bool_mask]
# 1.03 ms ± 37.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)


# ~ every 50th element
idx = np.random.randint(0, 1000000, size=int(1000000/50))  # changed the ratio!
bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[idx] = True

%timeit x[idx]
# 1.46 ms ± 55.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit x[bool_mask]
# 2.69 ms ± 154 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)


# ~ every 5th element
idx = np.random.randint(0, 1000000, size=int(1000000/5))  # changed the ratio!
bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[idx] = True

%timeit x[idx]
# 14.9 ms ± 495 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit x[bool_mask]
# 8.31 ms ± 181 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

那么这里发生了什么?很简单:整数数组索引只需要访问与索引数组中的值一样多的元素。这意味着如果匹配项很少,它会非常快,但如果有很多索引,它就会变慢。但是,布尔数组索引始终需要遍历整个布尔数组并检查 "true" 值。这意味着数组应该大致为 "constant"。

但是,等等,对于布尔数组来说它并不是真正的常量,为什么整数数组索引比布尔数组索引花费的时间更长(最后一种情况),即使它必须处理的元素少了约 5 倍?

这就是事情变得更加复杂的地方。在这种情况下,布尔数组在随机位置有 True,这意味着它将受到 分支预测失败 的影响。如果 TrueFalse 出现的次数相同但在随机位置,则更有可能出现这些情况。这就是布尔数组索引变得更慢的原因——因为 TrueFalse 的比率变得更相等,因此 "random" 变得更多。如果有更多 Trues,结果数组也会更大,这也会消耗更多时间。

作为这个分支预测的例子使用这个作为例子(可能因不同的 system/compilers 而不同):

bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[:1000000//2] = True   # first half True, second half False
%timeit x[bool_mask]
# 5.92 ms ± 118 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[::2] = True   # True and False alternating
%timeit x[bool_mask]
# 16.6 ms ± 361 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

bool_mask = np.zeros(x.shape, dtype=np.bool)
bool_mask[::2] = True
np.random.shuffle(bool_mask)  # shuffled
%timeit x[bool_mask]
# 18.2 ms ± 325 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

因此 TrueFalse 的分布将严重影响布尔掩码的 运行 时间,即使它们包含相同数量的 Trues!对于 compress-函数,同样的效果将可见。

对于整数数组索引(以及 np.take),另一个效果将是可见的:缓存位置。您的情况下的索引是随机分布的,因此您的计算机必须执行大量 "RAM" 到 "processor cache" 负载,因为两个索引不太可能彼此靠近。

比较这个:

idx = np.random.randint(0, 1000000, size=int(1000000/5))
%timeit x[idx]
# 15.6 ms ± 703 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

idx = np.random.randint(0, 1000000, size=int(1000000/5))
idx = np.sort(idx)  # sort them
%timeit x[idx]
# 4.33 ms ± 366 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

通过对索引进行排序,下一个值已经在缓存中的可能性大大增加,这可以带来巨大的加速。这是一个非常重要的因素,如果你知道索引将被排序(例如,如果它们是由 np.where 创建的,它们被排序,这使得 np.where 的结果对于索引特别有效)。

所以,整数数组索引对于小数组来说不是更慢,而对于大数组来说更快,它取决于更多的因素。两者都有自己的用例,并且根据具体情况,一个可能(相当)快于另一个。


让我也谈谈 numba 函数。首先是一些一般性陈述:

  • cache 不会有什么不同,它只是避免重新编译函数。在交互式环境中,这基本上是无用的。如果将函数打包到模块中会更快。
  • nogil 本身不会提供任何速度提升。如果在不同的线程中调用它会更快,因为每个函数执行都可以释放 GIL,然后多个调用可以 运行 并行。

否则我不知道 numba 是如何有效地实现这些功能的,但是当你在 numba 中使用 NumPy 功能时,它可能会变慢或变快 - 但即使它更快也不会快得多(除了小数组)。因为如果它可以变得更快,NumPy 开发人员也会实现它。我的经验法则是:如果你能用 NumPy 做到(矢量化),就不要为 numba 而烦恼。只有当您不能使用矢量化 NumPy 函数或 NumPy 会使用太多临时数组时,numba 才会大放异彩!