詳解如何利用Cython為Python代碼加速

2020-01-04 16:08:55

字體：大中小

供稿：網(wǎng)友

引言

通常，在 Python 中寫(xiě)循環(huán)（特別是多重循環(huán)）非常的慢，我們的元胞自動(dòng)機(jī)的狀態(tài)更新函數(shù) update_state 使用了兩重循環(huán)，所以我們嘗試用 Cython 重構(gòu)該方法。

代碼

我們?cè)谕募A下新建一個(gè) update.pyx 文件，寫(xiě)入如下內(nèi)容

import numpy as np cimport numpy as np cimport cythonDTYPE = np.floatctypedef np.float_t DTYPE_tdef update_state(np.ndarray[DTYPE_t, ndim=2] cells):  return update_state_c(cells)@cython.boundscheck(False)@cython.wraparound(False)cdef np.ndarray[DTYPE_t, ndim=2] update_state_c(np.ndarray[DTYPE_t, ndim=2] cells):  """更新一次狀態(tài)"""  cdef unsigned int i  cdef unsigned int j  cdef np.ndarray[DTYPE_t, ndim=2] buf = np.zeros((cells.shape[0], cells.shape[1]), dtype=DTYPE)  cdef DTYPE_t neighbor_num  for i in range(1, cells.shape[0] - 1):    for j in range(1, cells.shape[0] - 1):      # 計(jì)算該細(xì)胞周?chē)拇婊罴?xì)胞數(shù)            neighbor_num = cells[i, j-1] + cells[i, j+1] + cells[i+1, j] + cells[i-1, j] +/              cells[i-1, j-1] + cells[i-1, j+1] +/              cells[i+1, j-1] + cells[i+1, j+1]            if neighbor_num == 3:        buf[i, j] = 1      elif neighbor_num == 2:        buf[i, j] = cells[i, j]      else:        buf[i, j] = 0  return buf

update_state_c 函數(shù)上的兩個(gè)裝飾器是用來(lái)關(guān)閉 Cython 的邊界檢查的。

在同文件下新建一個(gè) setup.py 文件

import numpy as npfrom distutils.core import setupfrom Cython.Build import cythonizesetup(  name="Cython Update State",  ext_modules=cythonize("update.pyx"),  include_dirs=[np.get_include()])

因?yàn)樵?Cython 文件中使用了 NumPy 的頭文件，所以我們需要在 setup.py 將其包含進(jìn)去。

執(zhí)行 python setup.py build_ext --inplace 后，同文件夾下會(huì)生成一個(gè) update.cp36-win_amd64.pyd 的文件，這就是編譯好的 C 擴(kuò)展。

我們修改原始的代碼，首先在文件頭部加入 import update as cupdate，然后修改更新方法如下

def update_state(self):  """更新一次狀態(tài)"""  self.cells = cupdate.update_state(self.cells)  self.timer += 1

將原方法名就改為 update_state_py 即可，運(yùn)行腳本，無(wú)異常。

測(cè)速

我們編寫(xiě)一個(gè)方法來(lái)測(cè)試一下使用 Cython 可以帶來(lái)多少速度的提升

def test_time():  import time  game = GameOfLife(cells_shape=(60, 60))  t1 = time.time()  for _ in range(300):    game.update_state()  t2 = time.time()  print("Cython Use Time:", t2 - t1)  del game  game = GameOfLife(cells_shape=(60, 60))  t1 = time.time()  for _ in range(300):    game.update_state_py()  t2 = time.time()  print("Native Python Use Time:", t2 - t1)

運(yùn)行該方法，在我的電腦上輸出如下