如何在Python中优化嵌套的for循环

Question

我正在尝试编写一个Python函数来返回一个称为Mielke-Berry R值的指标。该指标的计算方式如下：

我已经编写了当前可用的代码，但是由于方程中涉及到两个求和操作，唯一想到的解决方法是在Python中使用嵌套for循环，这非常慢...

以下是我的代码：

def mb_r(forecasted_array, observed_array):
    """Returns the Mielke-Berry R value."""
    assert len(observed_array) == len(forecasted_array)
    y = forecasted_array.tolist()
    x = observed_array.tolist()
    total = 0
    for i in range(len(y)):
        for j in range(len(y)):
            total = total + abs(y[j] - x[i])
    total = np.array([total])
    return 1 - (mae(forecasted_array, observed_array) * forecasted_array.size ** 2 / total[0])

我将输入的数组转换为列表是因为我听说（尚未测试）使用Python循环索引numpy数组非常慢。

我觉得可能有某种numpy函数可以更快地解决这个问题，有人知道吗？

- pythonweb

---

除非已经有某种类似numpy的东西，否则我猜你可以尝试“Cythonize”它。看起来像是你会获得很大收益的事情。 - SuperStew

看起来这可以在NumPy中完成，而不需要使用Cython。 - Brad Solomon

我感觉自己可能忽略了什么，但这不就等于sum(abs(n*y - sum(x)))吗？如果是这样的话，你可以使用np.sum(np.absolute(n*forecasted_array - np.sum(observed_array)))。 - Engineero

Answer 1

以下是一种利用broadcasting的向量化方法，以获得total -

np.abs(forecasted_array[:,None] - observed_array).sum()

为了接受列表和数组，我们可以使用NumPy内置的外部减法，如下所示 -

np.abs(np.subtract.outer(forecasted_array, observed_array)).sum()

我们还可以利用 numexpr 模块 来更快地进行 absolute 计算，并在一次 numexpr evaluate 调用中执行 summation-reductions，从而更加节省内存，代码如下 -

import numexpr as ne

forecasted_array2D = forecasted_array[:,None]
total = ne.evaluate('sum(abs(forecasted_array2D - observed_array))')

- Divakar

---

我想补充一下，广播方法涉及到分配一个大小为 (N, N) 的数组，而 numexpr 的主要优点在于无需分配这个额外的内存。 - jakevdp

我刚想说，我用来测试的数组有12,000个值，所以仅使用numpy就出现了内存错误。而numexpr模块方法有效且速度可能快100-500倍。谢谢！ - pythonweb

Answer 2

使用numpy进行广播

如果没有内存限制，优化在numpy中的嵌套循环的第一步是使用广播并以矢量化的方式执行操作：

import numpy as np    

def mb_r(forecasted_array, observed_array):
        """Returns the Mielke-Berry R value."""
        assert len(observed_array) == len(forecasted_array)
        total = np.abs(forecasted_array[:, np.newaxis] - observed_array).sum() # Broadcasting
        return 1 - (mae(forecasted_array, observed_array) * forecasted_array.size ** 2 / total[0])

但是，尽管在此情况下循环发生在C中而不是Python中，但它涉及到一个大小为(N，N)的数组分配。

广播并非万能之策，请尝试展开内部循环。

正如前面所提到的，广播意味着巨大的内存开销。因此，应谨慎使用，并不总是正确的方式。虽然您可能第一时间想要在任何地方使用它-请不要。不久之前，我也被这个事实困惑，可以看看我的问题 Numpy ufuncs speed vs for loop speed。为避免过于冗长，我将使用您的示例说明：

import numpy as np

# Broadcast version
def mb_r_bcast(forecasted_array, observed_array):
    return np.abs(forecasted_array[:, np.newaxis] - observed_array).sum()

# Inner loop unrolled version
def mb_r_unroll(forecasted_array, observed_array):
    size = len(observed_array)
    total = 0.
    for i in range(size):  # There is only one loop
        total += np.abs(forecasted_array - observed_array[i]).sum()
    return total

小型数组（广播更快）

forecasted = np.random.rand(100)
observed = np.random.rand(100)

%timeit mb_r_bcast(forecasted, observed)
57.5 µs ± 359 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit mb_r_unroll(forecasted, observed)
1.17 ms ± 2.53 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Medium-size arrays (equal)

forecasted = np.random.rand(1000)
observed = np.random.rand(1000)

%timeit mb_r_bcast(forecasted, observed)
15.6 ms ± 208 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit mb_r_unroll(forecasted, observed)
16.4 ms ± 13.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

大型数组（广播速度较慢）

forecasted = np.random.rand(10000)
observed = np.random.rand(10000)

%timeit mb_r_bcast(forecasted, observed)
1.51 s ± 18 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit mb_r_unroll(forecasted, observed)
377 ms ± 994 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

您可以看到，对于小型数组而言，广播版本比展开版本快20倍，对于中型数组而言，它们相当，但对于大型数组而言，由于内存开销付出了代价，速度慢4倍。

Numba jit和并行化

另一种方法是使用numba及其强大的@jit函数装饰器。在这种情况下，只需要稍微修改您最初的代码即可。另外，要使循环并行化，您应该将range更改为prange，并提供parallel=True关键字参数。在下面的片段中，我使用了@njit装饰器，它与@jit(nopython=True)相同：

from numba import njit, prange

@njit(parallel=True)
def mb_r_njit(forecasted_array, observed_array):
    """Returns the Mielke-Berry R value."""
    assert len(observed_array) == len(forecasted_array)
    total = 0.
    size = len(forecasted_array)
    for i in prange(size):
        observed = observed_array[i]
        for j in prange(size):
            total += abs(forecasted_array[j] - observed)
    return 1 - (mae(forecasted_array, observed_array) * size ** 2 / total)

您没有提供“mae”函数，但是要在“njit”模式下运行代码，您还必须装饰“mae”函数，或者如果它是一个数字，请将其作为参数传递给被jit化的函数。

其他选项

Python科学生态系统非常庞大，我只提到了一些其他等效的加速选项：Cython、Nuitka、Pythran、bottleneck等等。也许您对gpu计算感兴趣，但这实际上是另一回事。

时间记录

在我的电脑上，不幸的是旧的那一台，时间记录如下：

import numpy as np
import numexpr as ne

forecasted_array = np.random.rand(10000)
observed_array   = np.random.rand(10000)

初始版本

%timeit mb_r(forecasted_array, observed_array)
23.4 s ± 430 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

numexpr

%%timeit
forecasted_array2d = forecasted_array[:, np.newaxis]
ne.evaluate('sum(abs(forecasted_array2d - observed_array))')[()]
784 ms ± 11.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

广播版本

%timeit mb_r_bcast(forecasted, observed)
1.47 s ± 4.13 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

内循环展开版本

%timeit mb_r_unroll(forecasted, observed)
389 ms ± 11.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

numba njit(parallel=True) version

%timeit mb_r_njit(forecasted_array, observed_array)
32 ms ± 4.05 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

可以看出，njit 方法比您最初的解决方案快 730x，也比 numexpr 解决方案快 24.5x（也许您需要英特尔的向量数学库来加速它）。此外，内部循环展开的简单方法使您的初始版本速度提高了 60x。我的规格如下：

Intel(R) Core(TM)2 Quad CPU Q9550 2.83GHz
Python 3.6.3
numpy 1.13.3
numba 0.36.1
numexpr 2.6.4

最终说明

我对你的一句话感到惊讶："我听说（尚未测试）使用python for循环索引numpy数组非常缓慢。" 所以我测试了一下：

arr = np.arange(1000)
ls = arr.tolistist()

%timeit for i in arr: pass
69.5 µs ± 282 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit for i in ls: pass
13.3 µs ± 81.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit for i in range(len(arr)): arr[i]
167 µs ± 997 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit for i in range(len(ls)): ls[i]
90.8 µs ± 1.07 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

事实证明，你是正确的。迭代列表的速度快2-5倍。当然，这些结果必须带有一定的幽默感 :)

- godaygo

---

所以我刚在我的电脑上运行了测试，使用 @njit 装饰器的 numba 模块确实快了大约17倍！谢谢！ - pythonweb

@Wade 比什么快五倍？ - godaygo

它的运行速度比numerexp快了大约17倍，这相当惊人！ - pythonweb

Answer 3

作为参考，以下代码：

#pythran export mb_r(float64[], float64[])
import numpy as np

def mb_r(forecasted_array, observed_array):
    return np.abs(forecasted_array[:,None] - observed_array).sum()

在纯CPython上以以下速度运行：

% python -m perf timeit -s 'import numpy as np; x = np.random.rand(400); y = np.random.rand(400); from mbr import mb_r' 'mb_r(x, y)' 
.....................
Mean +- std dev: 730 us +- 35 us

使用Pythran编译后，我得到了以下结果：

% pythran -march=native -DUSE_BOOST_SIMD mbr.py
% python -m perf timeit -s 'import numpy as np; x = np.random.rand(400); y = np.random.rand(400); from mbr import mb_r' 'mb_r(x, y)'
.....................
Mean +- std dev: 65.8 us +- 1.7 us

在具有AVX扩展的单个核心上，大致可以加快十倍的速度。

- serge-sans-paille