如何在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 for循環索引一個numpy數組非常慢。

我覺得可能有某種numpy功能可以更快地解決這個問題，任何人都知道什么？

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))')

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)

中型陣列（相等）

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函數，或者如果它是一個數字，則將它作為參數傳遞給jitted函數。

其他選擇

Python科學生態系統是巨大的，我只提到了一些其他等效的選項來加速： Cython ， Nuitka ， Pythran ， bottleneck和許多其他。 也許你對gpu computing感興趣，但這實際上是另一個故事。

計時

在我的電腦上，不幸的是舊的，時間是：

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)版本

%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快那么你的初步解決方案，同時也24.5x速度比numexpr解決方案（也許你需要英特爾的矢量數學庫，以加速它）。 內圈展開的簡單方法與初始版本相比，速度提高了60倍。 我的規格是：

英特爾（R）酷睿（TM）2四核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倍。 當然，這些結果必須帶有一定的反諷:)

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擴展的單核上大約是x10加速。