引言
通常, 在 Python 中写循环 (特别是多重循环) 非常的慢, 在文章 /article/18/0204/363226.html 中, 我们的元胞自动机的状态更新函数 update_state 使用了两重循环, 所以我们尝试用 Cython 重构该方法
代码
我们在同文件夹下新建一个 update.pyx 文件, 写入如下内容
- import numpy as np
- cimport numpy as np
- cimport cython
- DTYPE = np.float
- ctypedef np.float_t DTYPE_t
- def 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):
- """更新一次状态"""
- 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):
- # 计算该细胞周围的存活细胞数
- 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 函数上的两个装饰器是用来关闭 Cython 的边界检查的
在同文件下新建一个 setup.py 文件
- import numpy as np
- from distutils.core import setup
- from Cython.Build import cythonize
- setup(
- name="Cython Update State",
- ext_modules=cythonize("update.pyx"),
- include_dirs=[np.get_include()]
- )
因为在 Cython 文件中使用了 NumPy 的头文件, 所以我们需要在 setup.py 将其包含进去
执行 python setup.py build_ext --inplace 后, 同文件夹下会生成一个 update.cp36-win_amd64.pyd 的文件, 这就是编译好的 C 扩展
我们修改原始的代码, 首先在文件头部加入 import update as cupdate, 然后修改更新方法如下
- def update_state(self):
- """更新一次状态"""
- self.cells = cupdate.update_state(self.cells)
- self.timer += 1
将原方法名就改为 update_state_py 即可, 运行脚本, 无异常
测速
我们编写一个方法来测试一下使用 Cython 可以带来多少速度的提升
- 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)
运行该方法, 在我的电脑上输出如下
- Cython Use Time: 0.007000446319580078
- Native Python Use Time: 4.342248439788818
速度提升了 600 多倍
来源: http://www.phperz.com/article/18/0225/363225.html