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

Question

我需要通过阈值化3D数据阵列来创建布尔掩码:数据小于可接受下限的位置处的掩码或大于可接受上限的数据必须设置为True(否则False).简洁:

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

我有两个版本的代码用于执行此操作:一个直接使用整个3D数组,numpy而另一个方法循环遍历数组的切片.与我的期望相反,第二种方法似乎比第一种方法更快.为什么？？？

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

在我自己的测试中，性能差异比您的问题更加明显。增加数据的第二维和第三维后，差异仍然可以清晰地观察到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 中创建数组（我的实现）与执行所以原生（原始实现）