各种Numpy花式索引方法的性能表现，还包括使用Numba进行优化。

Question

各种Numpy花式索引方法的性能表现，还包括使用Numba进行优化。

32

由于我的程序需要快速索引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，缓存和nogil，我装饰了由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而言，很明显小数组最好使用布尔掩码进行索引（比ndarray.take(idx)快大约2倍），但对于较大的数组来说，ndarray.take(idx)表现最佳，在这种情况下，它比布尔索引快约6倍。当数组大小约为1000单元格，索引数组大小约为20时，两者性能相当。
对于具有1e5元素和5e3索引数组大小的数组，ndarray.take(idx)比布尔掩码索引快10倍左右。因此似乎布尔索引随着数组大小的增加而明显变慢，但在达到一定的数组大小阈值后会稍微追赶上来。

对于numba jitted函数而言，所有索引功能都有轻微的加速，除了布尔掩码索引。简单的fancy indexing在这里表现最佳，但仍然比没有jitting的布尔掩码索引慢。
对于较大的数组，布尔掩码索引比其他方法慢得多，甚至比非jitted版本还要慢。另外三种方法的性能都相当不错，比非jitted版本快约15%。

对于我这种有许多不同大小数组的情况，使用numba的fancy indexing是最好的方法。也许其他人也能在这篇相当冗长的文章中找到一些有用的信息。

编辑：
很抱歉我忘记了我的问题，实际上我有一个问题。我在快下班时匆忙打字，完全忘记了它... 那么，您知道是否有比我测试过的更好更快的方法吗？使用Cython，我的计时介于Numba和Python之间。
由于索引数组预定义一次并在长迭代中不进行修改，因此任何预定义索引过程的方法都将是很棒的。为此，我考虑使用步幅。但我无法预定义自定义步幅集。是否可能使用步幅获得预定义视图进入内存？

编辑2：
我想我会将关于预定义常数索引数组的问题移动到一个更具体的新问题中，该问题将在相同值数组上使用（其中只有值发生变化而形状不变）进行数百万次迭代。这个问题太笼统了，也许我还表达了问题，我将尽快发布新问题的链接！
这是后续问题的链接。

- JE_Muc

5

这里的问题是什么？提出一个真正的问题并进行自我回答不是更好吗？ - MSeifert

2

Scotty，请将您的问题转换为实际问题并将所有内容粘贴到自我答案中。如果您愿意，我可以通过社区Wiki进行粘贴，这样您就可以在此关闭（并删除）之前接受它，因为它被视为“不清楚您要问什么”。 - Daniel F

@DanielF 谢谢你的提示！我在最后加了一个问题！ - JE_Muc

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- MSeifert · Accepted Answer

你的总结并不完全正确，虽然你已经使用了不同大小的数组进行了测试，但你没有改变索引元素数量这一点。

我限制了纯索引并省略了 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的比率更加接近，因此更加“随机”。还有，如果有更多的True，结果数组将更大，这也会消耗更多时间。

关于这种分支预测的例子，请使用以下内容（在不同的系统/编译器中可能会有所不同）：

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函数也会有相同的效果。

对于整数数组索引（以及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如何有效地实现这些函数，但是当您在numba中使用NumPy功能时，它可能会更快或更慢-但即使它更快，它也不会快多少（除非对于小数组）。因为如果可以使其更快，则NumPy开发人员也会实现它。我的经验法则是：如果您可以使用NumPy进行（向量化）操作，请不要理会numba。只有当您无法使用向量化的NumPy函数或者NumPy将使用太多临时数组时，numba才会发挥作用！