python - 为什么 Python 在 numpy 数组切片上循环比完全矢量化操作更快

Question

我需要通过对 3D 数据数组进行阈值处理来创建 bool 掩码:数据小于可接受下限或数据大于可接受上限的位置的掩码必须设置为 True(否则错误)。简明扼要:

mask = (data < low) or (data > high)

我有两个版本的代码来执行此操作:一个直接使用 numpy 中的整个 3D 数组，而另一个方法循环遍历数组的切片。出乎我的意料，第二种方法似乎比第一种方法更快。为什么？？？

In [1]: import numpy as np

In [2]: import sys

In [3]: print(sys.version)
3.6.2 |Continuum Analytics, Inc.| (default, Jul 20 2017, 13:14:59) 
[GCC 4.2.1 Compatible Apple LLVM 6.0 (clang-600.0.57)]

In [4]: print(np.__version__)
1.14.0

In [5]: arr = np.random.random((10, 1000, 1000))

In [6]: def method1(arr, low, high):
   ...:     """ Fully vectorized computations """
   ...:     out = np.empty(arr.shape, dtype=np.bool)
   ...:     np.greater_equal(arr, high, out)
   ...:     np.logical_or(out, arr < low, out)
   ...:     return out
   ...: 

In [7]: def method2(arr, low, high):
   ...:     """ Partially vectorized computations """
   ...:     out = np.empty(arr.shape, dtype=np.bool)
   ...:     for k in range(arr.shape[0]):
   ...:         a = arr[k]
   ...:         o = out[k]
   ...:         np.greater_equal(a, high, o)
   ...:         np.logical_or(o, a < low, o)
   ...:     return out
   ...: 

首先，让我们确保这两种方法产生相同的结果:

In [8]: np.all(method1(arr, 0.2, 0.8) == method2(arr, 0.2, 0.8))
Out[8]: True

现在进行一些计时测试:

In [9]: %timeit method1(arr, 0.2, 0.8)
14.4 ms ± 111 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [10]: %timeit method2(arr, 0.2, 0.8)
11.5 ms ± 241 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

这是怎么回事？

---

编辑 1:在旧环境中观察到类似的行为:

In [3]: print(sys.version)
2.7.13 |Continuum Analytics, Inc.| (default, Dec 20 2016, 23:05:08) 
[GCC 4.2.1 Compatible Apple LLVM 6.0 (clang-600.0.57)]

In [4]: print(np.__version__)
1.11.3

In [9]: %timeit method1(arr, 0.2, 0.8)
100 loops, best of 3: 14.3 ms per loop

In [10]: %timeit method2(arr, 0.2, 0.8)
100 loops, best of 3: 13 ms per loop

Answer 1

优于两种方法

在方法一中，您访问数组两次。如果它不适合缓存，数据将从 RAM 中读取两次，从而降低性能。此外，可能会创建临时数组，如评论中所述。

方法二对缓存更友好，因为您两次访问数组的较小部分，这很可能适合缓存。缺点是循环慢，函数调用多，速度也比较慢。

为了在这里获得良好的性能，建议编译代码，这可以使用 cython 或 numba 来完成。由于 cython 版本需要更多工作(注释，需要单独的编译器)，我将展示如何使用 Numba 完成此操作。

import numba as nb
@nb.njit(fastmath=True, cache=True)
def method3(arr, low, high):
  out = np.empty(arr.shape, dtype=nb.boolean)
  for i in range(arr.shape[0]):
    for j in range(arr.shape[1]):
      for k in range(arr.shape[2]):
        out[i,j,k]=arr[i,j,k] < low or arr[i,j,k] > high
  return out

使用 arr = np.random.random((10, 1000, 1000)) 这比你的 method_1 好两倍，你的 method_2 在我的电脑上好 50%(Core i7-4771 , python 3.5, windows)

这只是一个简单的例子，在更复杂的代码上，你可以使用 SIMD，并行处理也很容易使用，性能增益可以大很多。在非编译代码上，向量化通常但并不总是(如图所示)是你能做的最好的，但它总是会导致不良的缓存行为，如果你正在访问的数据 block 至少不适合，这可能会导致性能不佳在 L3 缓存中。在其他一些问题上，如果数据无法容纳在更小的 L1 或 L2 缓存中，也会影响性能。另一个优势是在调用此函数的 njited 函数中自动内联小型 njited 函数。