用Cython來提高Python代碼速度 [二]

這篇文章將通過一個簡單的例子來展示如何使用Cython來提升函數的速度。

目標函數：

def is_subset(a, b): """ 判斷a是否為b的子集 :return boolean """ pass

大多數人知道在python裡面實現這個目標非常簡單：

def py_is_subset(a, b): return set(a).issubset(set(b))

然而我們要做的就是，假設我們知道a和b都是long類型的整數，並且知道他們各自的長度，如何實現最快速的演算法。注意：這些假設對上面的Python函數沒有任何作用。

當然這個演算法本身毫無亮點可言，任何會計算機演算法的人都會不屑花時間研究這個。然而為了本文行文方便，演算法描述如下 (用python句式）：

for i in range(len(a)): j = 0 while j < len(b): if a[i] == b[j]: break j += 1 if j == len(b): return Falsereturn True

第一步：演算法初稿：

我們根據上述演算法寫出cython函數。

以下演算法中：

len_a = len(a), len_b = len(b),

def is_subset_bad(a, b, len_a, len_b): for i in range(len_a): j = 0 while j < len_b and a[i] != b[j]: j += 1 if j == len_b: return False return True

看起來不錯，然而benchmark之後：（compile和benchmark的過程暫時省略，請參考有關資料）

Average runtime of 200 runs is:[Cython] 0.0018010350004260545 ms[Python] 0.00029243499966469244 ms

簡直辣雞，到底怎麼回事？

第二步：分析問題：

Cython提供了一個非常好用的功能: annotation，如果你的源文件是is_subset.pyx，可以使用：

cython is_subset.pyx -a

來生成一個html文件，打開後發現：

注意那些深淺不一的黃色，顏色越深，調用的Python命令越多。點開+號會看到：

非常可怕的各種檢查。也就是說，我們並沒有做到太多的優化，這跟python自帶的issubset命令比較起來，高下立判。那麼問題在哪裡？

第三步：static typing

這是最重要，也是最常用的cython優化步驟。我們需要將所有變數的類型申明，這樣cython在編譯的時候會跳過檢查變數類型，從而節省大量時間。修改之後的代碼如下：

def is_subset_bad(list a, list b, size_t len_a, size_t len_b): cdef size_t i, j for i in range(len_a): j = 0 while j < len_b and a[i] != b[j]: j += 1 if j == len_b: return False return True

benchmark 結果如下：

Average runtime of 200 runs is:[Cython] 0.0005020649996367865 ms[Python] 0.00029244499955893843 ms

一下比以前提高了大約3倍！！

然而依然不如python自帶函數。問題在哪裡？

注意函數變數類型：

def is_subset_bad(list a, list b, size_t len_a, size_t len_b):

其中 size_t 簡單來說是cython用來代表列表長度的變數，不能進行優化了。然而list是一個非常含糊的類型，cython提供了這一個類型來方便python用戶在不知道列表元素類型時候使用，在這裡就並沒有用到我們已知的「a,b都是長整數類型的列表」這一知識。

這裡需要介紹一個非常重要的cython類型：array，在python中也有對應的庫。簡而言之，這是一個overhead非常少的，儲存統一類型變數的類。

於是我們修改函數如下：

def is_subset_bad(array.array a, array.array b, size_t len_a, size_t len_b): cdef size_t i, j for i in range(len_a): j = 0 while j < len_b and a[i] != b[j]: j += 1 if j == len_b: return False return True

Benchmark如下：

Average runtime of 200 runs is:[Cython] 0.0007023650005066884 ms[Python] 0.00029687000051126233 ms

等一下，怎麼更慢了？讀者可以自行使用cython -a來確定細節。簡而言之，我們仍然沒有把「長整數」這個知識使用進去！（hint: a[i] != b[j] 仍然使用大量python檢查)，這裡繼續進行究級改進：

cimport cythonfrom cpython cimport arrayfrom libc.string cimport memcmpdef is_subset_bad(array.array a, array.array b, size_t len_a, size_t len_b): cdef size_t i, j cdef size_t long_size = sizeof(long) for i in range(len_a): j = 0 while j < len_b and memcmp(a.data.as_longs + i, b.data.as_longs + j, long_size): j += 1 if j == len_b: return False return True

這裡需要解釋一些函數和methods:

1. memcmp是c的函數，純粹比較內存位置，相等的話返回0（True)，參考C library function - memcmp()
2. a.data.as_longs直接獲得array a對應的指針，並將其認為是長整形指針，這就完全避免了Python的類型檢查！
3. sizeof同樣是c的函數，用來獲得長整型類型變數長度，用於memcmp中

benchmark結果如下：

Average runtime of 200 runs is:[Cython] 0.00014715500128659187 ms[Python] 0.0002981149987135723 ms

成功了！我們的函數現在比內置的python函數快了100%！

小結：

這是一個非常簡單的列子，然而基本上組成了使用cython來提速python函數的基本要素。本文暫時略過了如何設置編譯源代碼，如何導入編譯完的函數，以及如何進行benchmark的方法。以後的文章中有機會再進行補充。此外，值得一提的是，如果a和b都是排序過的，還可以簡單修改cython函數（兩個循環）來進一步節省計算次數。這樣的雙循環在python中可能會由於太多的類型檢查而並不能獲得顯著的速度提高。在Cython中則完全不用擔心！

參考文章：

1. 用Cython來提高Python代碼速度 [一]
2. 【學習心得】寫出快速python程序【可能是一】
3. Cythons Documentation
4. Cython: A Guide for Python Programmers

推薦閱讀：