表现各种numpy花式索引方法,也与numba

Question

因为对于我的程序来说,快速索引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-mode中的jitting ,caching和nogilI装饰了索引的方式,这些方式受以下支持numba:

@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很明显,小数组最好用布尔掩码索引(比较大约2倍ndarray.take(idx)),因为较大的数组ndarray.take(idx)表现最佳,在这种情况下比布尔索引快6倍.盈亏平衡点位于1000单元格周围的数组大小,单元格周围的索引数组大小20.
对于具有1e5元素和5e3索引数组大小的数组,ndarray.take(idx)将比布尔掩码索引快10倍.因此,似乎布尔索引似乎随着数组大小而显着减慢,但在达到某个数组大小阈值后会略微增加.

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

对于我的情况,有许多不同大小的数组,使用numba进行花式索引是最好的方法.也许其他人也可以在这篇相当冗长的帖子中找到一些有用的信息.

编辑:
对不起,我忘了问我的问题,我实际上已经问过了.我只是在工作日结束时快速输入这个并完全忘记了......嗯,你知道比我测试的更好更快的方法吗？使用Cython我的时间是在Numba和Python之间.
由于索引数组是预定义的一次并且在长迭代中使用而没有改变,因此任何预定义索引过程的方式都会很好.为此我想到了使用步幅.但我无法预先定义一组自定义的步幅.是否可以使用步幅将预定义视图添加到内存中？

编辑2:
我想我将移动我关于预定义常量索引数组的问题,这些数组将在相同的值数组(其中只有值改变而不是形状)上使用几百万次迭代到一个新的更具体的问题.这个问题太笼统了,或许我也提出了一个有点误导性的问题.一打开新问题,我就会在这里发布链接!
以下是后续问题的链接.

Answer 1

您的摘要不完全正确,您已经使用不同大小的数组进行了测试,但您没有做的一件事是更改索引元素的数量.

我将它限制为纯索引并省略take(实际上是整数数组索引)compress和extract(因为它们实际上是布尔数组索引).这些的唯一区别是不变因素.对这些方法的常数因子take和compress将小于开销为numpy的功能np.take和np.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"值.这意味着它应该对阵列大致"恒定".

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

这就是它变得更加复杂的地方.在这种情况下,布尔数组具有True随机位置,这意味着它将受到分支预测失败的影响.如果True并且False将会发生相同的事件,但是在随机的地方,这些将更有可能发生.这就是为什么布尔数组索引越来越慢-因为比True来False得到更加平等,从而更加"随机".如果有更多的Trues也会消耗更多的时间,结果数组也会更大.

作为这个分支预测的一个例子,使用它作为例子(可能因不同的系统/编译器而不同):

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)

因此,分布True和False会严重影响运行时用,即使它们包含相同数量的布尔口罩True小号!对于compress-functions,可以看到相同的效果.

对于整数数组索引(以及同样np.take),将显示另一个效果:缓存局部性.您的案例中的索引是随机分布的,因此您的计算机必须执行大量"RAM"到"处理器缓存"加载,因为两个索引不太可能彼此靠近.

比较一下:

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用途不同实现这些功能,但是当你使用NumPy的特征numba它可能是更慢或更快 - 但即使它的速度更快也不会快很多(可能除了小数组).因为如果它可以更快,NumPy开发人员也会实现它.我的经验法则是:如果你能用NumPy做(向量化),不要打扰numba.只有你不能用矢量化的NumPy函数或者NumPy才会使用太多的临时数组,那么numba就会闪耀!