為什么NumPy有時比NumPy +普通Python循環慢？

Question

這是基於2018-10問的這個問題 。

請考慮以下代碼。 三個簡單函數用於計算NumPy 3D陣列（1000×1000×1000）中的非零元素。

import numpy as np

def f_1(arr):
    return np.sum(arr > 0)

def f_2(arr):
    ans = 0
    for val in range(arr.shape[0]):
        ans += np.sum(arr[val, :, :] > 0)
    return ans

def f_3(arr):
    return np.count_nonzero(arr)

if __name__ == '__main__':

    data = np.random.randint(0, 10, (1_000, 1_000, 1_000))
    print(f_1(data))
    print(f_2(data))
    print(f_3(data))

我機器上的運行時（Python 3.7。？，Windows 10，NumPy 1.16。？）：

%timeit f_1(data)
1.73 s ± 21.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_2(data)
1.4 s ± 1.36 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_3(data)
2.38 s ± 956 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

因此， f_2()比f_1()和f_3()工作得更快。 但是，尺寸較小的data並非如此。 問題是 - 為什么這樣？ 它是NumPy，Python還是其他什么？

Answer 1

這是由於內存訪問和緩存。 這些函數中的每一個都做了兩件事，以第一個代碼為例：

np.sum(arr > 0)

它首先進行比較，找出arr大於零的位置（或非零，因為arr包含非負整數）。 這將創建一個與arr形狀相同的中間數組。 然后，它將此數組相加。

直截了當，對吧？ 好吧，當使用np.sum(arr > 0)這是一個大數組。 當它足夠大以至於不適合高速緩存時，性能將降低，因為當處理器開始執行總和時，大多數數組元素將從內存中逐出並需要重新加載。

由於f_2在第一維上迭代，因此它處理較小的子陣列。 完成相同的副本和總和，但這次中間數組適合內存。 無需留下記憶就可以創建，使用和銷毀它。 這要快得多。

現在，您會認為f_3最快（使用內置方法和所有方法），但查看源代碼顯示它使用以下操作：

a_bool = a.astype(np.bool_, copy=False)
return a_bool.sum(axis=axis, dtype=np.intp

a_bool只是查找非零條目的另一種方法，並創建一個大型中間數組。

結論

經驗法則就是這樣，並且經常是錯誤的。 如果你想要更快的代碼，請對其進行分析並查看問題所在（這里的工作很好）。

Python做得很好。 在優化的情況下，它可能比numpy更快。 不要害怕將普通的舊python代碼或數據類型與numpy結合使用。

如果您經常發現自己手動編寫for循環以獲得更好的性能，您可能需要查看numexpr - 它會自動執行其中的一些操作。 我自己並沒有太多使用它，但如果中間數組正在減慢程序的速度，它應該提供一個很好的加速。

Answer 2

這完全取決於數據如何在內存中布局以及代碼如何訪問它。 本質上，數據是以塊為單位從內存中提取的，然后緩存; 如果算法設法使用來自緩存中的塊的數據，則無需再次從內存中讀取數據。 這可以節省大量時間，特別是當緩存比您正在處理的數據小得多時。

考慮這些變化，它們只在我們迭代的軸上有所不同：

def f_2_0(arr):
    ans = 0
    for val in range(arr.shape[0]):
        ans += np.sum(arr[val, :, :] > 0)
    return ans

def f_2_1(arr):
    ans = 0
    for val in range(arr.shape[1]):
        ans += np.sum(arr[:, val, :] > 0)
    return ans

def f_2_2(arr):
    ans = 0
    for val in range(arr.shape[2]):
        ans += np.sum(arr[:, :, val] > 0)
    return ans

我的筆記本電腦上的結果：

%timeit f_1(data)
2.31 s ± 47.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_2_0(data)
1.88 s ± 60 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_2_1(data)
2.65 s ± 142 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_2_2(data)
12.8 s ± 650 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

您可以看到f_2_1幾乎與f_1一樣快， 這讓我覺得numpy沒有使用最佳訪問模式（ f_2_0使用的f_2_0 ） 。 另一個答案是關於精確緩存如何影響時間的解釋 。

Answer 3

讓我們完全刪除臨時數組

正如@ user2699在他的回答中已經提到的那樣，分配和寫入不適合緩存的大型數組可能會大大減慢這個過程。 為了顯示這種行為，我使用Numba（JIT-Compiler）編寫了兩個小函數。

在編譯語言（C，Fortran，..）中，您通常會避免使用臨時數組。 在解釋Py​​thon（不使用Cython或Numba）中，您經常需要在較大的數據塊（向量化）上調用已編譯的函數，因為解釋代碼中的循環非常慢。 但這也可能有一個缺點（如臨時數組，錯誤的緩存使用）

沒有臨時數組分配的功能

@nb.njit(fastmath=True,parallel=False)
def f_4(arr):
    sum=0
    for i in nb.prange(arr.shape[0]):
        for j in range(arr.shape[1]):
            for k in range(arr.shape[2]):
                if arr[i,j,k]>0:
                    sum+=1
    return sum

用臨時數組

請注意，如果打開parallelization parallel=True ，編譯器不僅會嘗試並行化代碼，還會啟用其他優化（如循環融合）。

@nb.njit(fastmath=True,parallel=False)
def f_5(arr):
    return np.sum(arr>0)

計時

%timeit f_1(data)
1.65 s ± 48.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f_2(data)
1.27 s ± 5.66 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f_3(data)
1.99 s ± 6.11 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit f_4(data) #parallel=false
216 ms ± 5.45 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f_4(data) #parallel=true
121 ms ± 4.85 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f_5(data) #parallel=False
1.12 s ± 19 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f_5(data) #parallel=true Temp-Array is automatically optimized away
146 ms ± 12.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)