为什么Python循环过多的numpy数组比完全向量化的操作更快

Question

我需要通过对3D数据阵列进行阈值处理来创建布尔掩码：数据小于可接受下限的位置处的掩码或者大于可接受上限的数据必须设置为True （否则为False ）。 简洁：

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优于你的两倍，你的方法_2在我的电脑上优于50％（Core i7-4771，python 3.5，windows）

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

Answer 2

在我自己的测试中，性能差异比你的问题更明显。 在增加arr数据的第二和第三维之后，差异仍然可以清楚地观察到。 在评论出两个比较函数之一（ greater_equal或logical_or ）后，它仍然可以被观察到，这意味着我们可以排除两者之间的某种奇怪的相互作用。

通过将两种方法的实现更改为以下内容，我可以显着降低可观察到的性能差异（但不能完全消除它）：

def method1(arr, low, high):
    out = np.empty(arr.shape, dtype=np.bool)
    high = np.ones_like(arr) * high
    low = np.ones_like(arr) * low
    np.greater_equal(arr, high, out)
    np.logical_or(out, arr < low, out)
    return out

def method2(arr, low, high):
    out = np.empty(arr.shape, dtype=np.bool)
    high = np.ones_like(arr) * high
    low = np.ones_like(arr) * low
    for k in range(arr.shape[0]):
        a = arr[k]
        o = out[k]
        h = high[k]
        l = low[k]
        np.greater_equal(a, h, o)
        np.logical_or(o, a < l, o)
    return out

我认为，当向那些numpy函数提供high或low作为标量时，它们可能在内部首先创建一个填充了该标量的正确形状的numpy数组。 当我们在功能之外手动执行此操作时，在两种情况下仅对整个形状执行一次，性能差异变得不那么明显。 这意味着，无论出于何种原因（也许缓存？），创建这样一个大阵充满了相同的恒定一次可能比创造效率较低k用相同的常数较小的阵列（如通过实施自动完成的method2在原来的问题）。

---

注意：除了缩小性能差距外，它还会使两种方法的性能更差（比第一种方法更严重地影响第二种方法）。 因此，虽然这可能会给出问题可能出现的一些迹象，但似乎并未解释所有问题。

---

编辑

这是一个新版本的method2 ，我们现在每次都在循环中手动预先创建较小的数组，就像我怀疑在问题的原始实现中内部正在numpy中发生的那样：

def method2(arr, low, high):
    out = np.empty(arr.shape, dtype=np.bool)
    for k in range(arr.shape[0]):
        a = arr[k]
        o = out[k]
        h = np.full_like(a, high)
        l = np.full_like(a, low)
        np.greater_equal(a, h, o)
        np.logical_or(o, a < l, o)
    return out

这个版本确实比我上面的版本快得多（确认在循环内创建许多较小的数组比循环外的一个大数组更有效），但仍然比问题中的原始实现慢。

假设这些numpy函数确实首先将标量边界转换为这些类型的数组，这个最后一个函数与问题中的函数之间的性能差异可能是由于在Python中创建数组（我的实现）而不是原生（原始实施）