用 NumPy 向量化加速 Python：千萬級矩形面積計算變快 16~177 倍 | LKY 只有原創內容的 Blog

Python 簡單易用，可以很快速的驗證各種演算法。但是到了真的要實際應用時，Python 的毛病就會凸顯，往往就是慢！而且慢的很難受。

我也熟悉 Go 或 C# 這些靜態型別的語言，但是欠缺各式各樣方便好用的套件，難以快速的驗證各種演算法。

這篇文章我們來看看如何用 NumPy 來加速 Python 的效能，並且用一個簡單的例子來說明。

需求是：計算 1780 萬個矩形的面積

入門的人怎麼寫？速度定義為 1x

import NumPy as np
import time

# 隨機產生 1780 萬個矩形的兩個對角點座標，為什麼用這個數字呢？因為178是我的身高
np.random.seed(0) # 固定隨機種子，讓每次執行的結果都一樣，方便大家比較
rects = np.random.randint(0, 100, (17800000, 2, 2))

# 計算這些矩形的面積，使用 for 迴圈
start_time = time.time()
area = []
for rect in rects:
    area.append((abs(rect[0][0] - rect[0][1]) * abs(rect[1][0] - rect[1][1])))
elapsed_time_for_loop = time.time() - start_time
print("檢驗結果: %s" % (sum(area)))
print("Elapsed time: %s" % (elapsed_time_for_loop))

輸出如下：

檢驗結果: 19769742257
Elapsed time: 16.575732946395874

上面這個程式碼，就是一般初學者最直覺，用 for loop 寫出來的程式碼，但是效能很差，因為:

• Python 是動態型別，所以每次迴圈都要去判斷型別，並且要去做型別轉換。
• 而且 for 迴圈一次只能處理一個矩形，數據頻繁的進出往來 RAM 與 CPU，效能也很差。

---

用 NumPy 向量化計算(簡單版)，速度 16x

import NumPy as np
import time

# 隨機產生 1780 萬個矩形的兩個對角點座標，為什麼用這個數字呢？因為178是我的身高
rects = np.random.randint(0, 100, (17800000, 2, 2))

# 計算這些矩形的面積，使用最簡單的向量化
start_time = time.time()
min_x = np.min(rects[:, 0], axis=1)
max_x = np.max(rects[:, 0], axis=1)
min_y = np.min(rects[:, 1], axis=1)
max_y = np.max(rects[:, 1], axis=1)
area = (max_x - min_x) * (max_y - min_y)
elapsed_time_vectorlize = time.time() - start_time
print("檢驗結果: %s" % (sum(area)))
print("Elapsed time: %s" % (elapsed_time_vectorlize))

輸出如下：

檢驗結果: 19769742257
Elapsed time: 1.0222868919372559

上面這個程式碼，直接把矩形的兩個對角點座標的XY值，分別按照大小取出來，成為四個 Nx1 矩陣。然後只做1次矩陣化的面積運算，就算完了 N 個矩形的面積，效能提升了 16 倍。

為什麼快？

• 因為 NumPy 是靜態型別，所以不用每次迴圈都去判斷型別，並且不用去做型別轉換。
• 一次處理多個矩形，數據不用頻繁的進出 RAM 與 CPU，效能更好。

當演算法被正確地向量化時，CPU 僅需一條指令完成這行程式碼，而不是對每個 i 進行獨立操作。理想情況下，any(result)操作將只發生於 CPU 內部而不用將數據傳回 RAM。

---

用 NumPy 向量化計算(進階版，使用np.ptp)，速度 22x

# 以上重複部份程式碼省略
# 計算這些矩形的面積，使用np.ptp向量化
start_time = time.time()
area = np.ptp(rects[:, 0], axis=1) * np.ptp(rects[:, 1], axis=1)
elapsed_time_ptp = time.time() - start_time
print("檢驗結果: %s" % (sum(area)))
print("Elapsed time: %s" % (elapsed_time_ptp))

輸出如下：

檢驗結果: 19769742257
Elapsed time: 0.7305142879486084

上面這個程式碼，就是用 NumPy 向量化計算的方式，效能提升了 22 倍。
為什麼又更快？

因為 np.ptp 的意思就是「peak to peak」，把「最大值、最小值、取差值」的計算合併成一個函式，所以效能更好。

---

用 NumPy 向量化計算(進階版，直接取差值，然後絕對值相乘)，速度 177x

# 以上重複部份程式碼省略
# 計算這些矩形的面積，直接取絕對值相乘
start_time = time.time()
area = np.abs(rects[:, 0, 0] - rects[:, 0, 1]) * np.abs(rects[:, 1, 0] - rects[:, 1, 1])
# 下面這種方法更快一點，但可能會發生 Exception has occurred: _ArrayMemoryError
# area = np.abs(np.diff(rects[:, 0], axis=1)) * np.abs(np.diff(rects[:, 1], axis=1)).flatten()
elapsed_time_diff_abs = time.time() - start_time
print("檢驗結果: %s" % (sum(area)))
print("Elapsed time: %s" % (elapsed_time_diff_abs))

輸出如下：

檢驗結果: 19769742257
Elapsed time: 0.09360003471374512

上面這個程式碼，就是用 NumPy 向量化計算的方式，效能提升了 177 倍。

為什麼又更快？

老實說，到這裡已經到了我的知識盲區了。

不論是 slicing 或者進出 RAM 與 CPU 的次數，與前一種方法看起來都一樣，但是效能卻可以從 22 倍提升到 177 倍。我猜測可能是因為 np.ptp 有些限制導致速度比較慢。

用 NumPy 就是這樣，通常會比只用單純的 Python 快很多，但是各種不同用法的效能差異卻很大。

想達到最好的效能，有沒有一次到位的方法？除非有時間仔細鑽研 NumPy 的底層原理，不然我覺得沒有，只能一步一步的嘗試，並且不斷的優化。

更快的方法，要注意其他的風險

可以看到，在最後一段程式碼中，有一行被我註解的地方。

用 np.diff 也可以達到同樣的效果，還更快一點，而且我完全想不出原因是什麼？

但是我開了 Docker 之後，就會發生爆記憶體的錯誤：

Exception has occurred: _ArrayMemoryError
Unable to allocate 2.25 PiB for an array with shape (17800000, 17800000) and data type int64
  File "/Users/xxx/demo_vectorlize.py", line 44, in <module>
    area = np.abs(np.diff(rects[:, 0], axis=1)) * np.abs(np.diff(rects[:, 1], axis=1)).flatten()
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
numpy.core._exceptions._ArrayMemoryError: Unable to allocate 2.25 PiB for an array with shape (17800000, 17800000) and data type int64

可以看到 np.diff 跟記憶體要了一塊 N x N 的空間，而 N 高達 17800000，這對許多系統來說是無法承受的。而且你開發時可能不會出現這個問題，但是到了生產環境才發現，那就踩了大坑了。

總結

三種向量化計算的效能，與原本的 for-loop 效能比較如下：

Speed up for vectorlize: 16.214365142630854x
Speed up for ptp-vectorlize: 22.690497940763034x
Speed up for diff-abs-vectorlize: 177.09109827885285x

這篇文章主要是介紹了 NumPy 的向量化計算，用一個工作上實際遇過的簡單案例，來說明 NumPy 向量化計算的優點。

這篇文章的程式碼，可以在這裡找到：https://gist.github.com/mosdeo/743a6e2275ea7676dc2f98246e055b33