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

时间:2018-02-20 14:25:59

标签: python numpy

我需要通过对3D数据数组进行阈值处理来创建布尔掩码:数据小于可接受下限的位置处的掩码或大于可接受上限的数据必须设置为True(否则为{{1} }})。简洁地:

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

这里发生了什么?

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

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)

2 个答案:

答案 0 :(得分:3)

超越两种方法

在方法一中,您要访问数组两次。如果它不适合缓存,则数据将从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))在我的电脑上(Core i7-4771,python 3.5,windows),你的方法1优于你的方法1,你的方法_2优于你的百分之二十五

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

答案 1 :(得分:0)

在我自己的测试中,性能上的差异比你的问题更明显。在增加arr数据的第二和第三维之后,差异仍然可以清楚地观察到。在评论出两个比较函数之一(greater_equallogical_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

我认为,当提供highlow作为那些numpy函数的标量时,它们可能会在内部首先创建一个填充了该标量的正确形状的numpy数组。当我们在功能之外手动执行此操作时,在两种情况下仅对整个形状执行一次,性能差异变得不那么明显。这意味着,无论出于何种原因(可能是缓存?),创建一个填充了相同常量的大型数组可能效率低于创建具有相同常量的k个较小数组(由{的实现自动完成)原始问题中的{1}}。

注意:除了缩小性能差距外,还会使两种方法的性能更差(比第一种方法更严重地影响第二种方法)。因此,虽然这可能会显示出问题的位置,但它似乎并不能解释所有问题。

修改

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

method2

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

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