Cython--Python加速工具的使用

 

前言

纯Python做计算，相比Numpy、PyTorch、C、C++，性能低很多。Numpy适合单机版CPU矩阵计算，PyTorch适用于单机、分布式CPU、GPU矩阵计算，C、C++的性能不做多说；

另外，针对大规模分布式数据集矩阵计算，有Numba、Dask、JAX（Google）、Mars（阿里）等；

如果你的项目对性能要求很高，且数据不易组织为矩阵，又不想直接写C、C++代码，Cython是一个不错的选择。

Cython官网：https://cython.org/

Cython官方文档：http://docs.cython.org/en/latest/

接下来，用几个栗子来描述Cython的使用；

测试环境：Centos7.6 + python3.6.8 + Cython0.29.7

1. 纯Python转为.so可执行文件

使用Cython将纯Python代码转化为.so可执行文件，可加速计算；

首先，创建文件test.py，添加如下内容：

def f(x):
    return x ** 2 - x


def integrate_f(a, b, N):
    s = 0
    dx = (b - a) / N
    for i in range(N):
        s += f(a + i * dx)
    return s * dx

创建文件main.py，然后命令行执行 python3.6 main.py

import test
import time

start = time.time()
for _ in range(10000):
    test.integrate_f(12,10,100)
print(time.time() - start)

执行10000次计算，纯Python耗时0.297s；

创建test_cython.pyx，

def f(x):
    return x ** 2 - x


def integrate_f(a, b, N):
    s = 0
    dx = (b - a) / N
    for i in range(N):
        s += f(a + i * dx)
    return s * dx

下面将.pyx使用Cython编译成.so文件，编译过程中，会先将.pyx文件编译成.c文件，然后将.c文件编译成.so可执行文件；

操作如下：

创建文件setup.py，

from distutils.core import setup
from Cython.Build import cythonize

setup(
    name='test_module',
    ext_modules=cythonize('test.pyx'),
)

然后，在命令行执行编译过程；

python3.6 setup.py build_ext --inplace

操作成功后，在当前文件夹中会出现2个文件：test.c和test.cpython-36m-x86_64-linux-gnu.so，这个.so是可以直接在Python直接import的，但是.pyx是不可以的；

接下来，测试优化后的性能；main.py不需要做任何修改，将之前的test.py修改为test.py.bak；

然后命令行执行，python3.6 main.py，运行时间如下：

执行10000次计算，纯Python转为.so文件后耗时0.18s；性能提高了1.6倍；

2. 使用静态类型加速

给Python变量添加上C的类型，可以进一步加速；

栗子如下：

创建文件test_static.pyx，

def f(double x):
    return x ** 2 - x


def integrate_f(double a, double b, int N):
    cdef int i
    cdef double s, dx
    s = 0
    dx = (b - a) / N
    for i in range(N):
        s += f(a + i * dx)
    return s * dx

脚本中出现的变量都添加了C的类型，变量使用cdef关键字做声明；

下一步同上，将.pyx编译为.so文件：

修改setup.py文件，

from distutils.core import setup
from Cython.Build import cythonize

setup(
    name='test_module_2',
    ext_modules=cythonize('test_static.pyx'),
)

然后，命令行执行下面代码，会多出来一个.c文件，一个.so文件，这个.so文件也是可以直接在Python脚本中import；

python3.6 setup.py build_ext --inplace

接下来，见证奇迹的时刻到了；

修改main.py，并执行；

import time
import test_static

start = time.time()
for _ in range(10000):
    test_static.integrate_f(12,10,100)
print(time.time() - start)

执行10000次计算，静态类型加速后耗时0.04s；性能进一步提高了4.5倍；

3. 使用c函数加速

计算主函数，使用cdef替代def，修改test_static.pyx

cdef f(double x):
    return x ** 2 - x


def integrate_f(double a, double b, int N):
    cdef int i
    cdef double s, dx
    s = 0
    dx = (b - a) / N
    for i in range(N):
        s += f(a + i * dx)
    return s * dx

重新编译运行，耗时0.0249s；性能进一步提高了1.6倍。