学习笔记 —— 基于C加速的Python高效计算 (Cython pybind11)
创始人
2024-05-28 20:12:45

目录

  • 引言
  • Cython
    • 示例介绍
    • 第一阶段优化
    • 第二阶段优化
      • Cython Annotation tool
      • 优化方法
    • 第三阶段优化
    • 比对下 JIT的Numba
    • 总结
  • pybind11
    • Links
    • Introduction
    • Implementation
    • Implementation Cmake
    • cytpes
  • Cython & pybind11 性能比较

TODO
Implementation Cmake --pybind11
Cython & pybind11 性能比较

引言

Python与C++,是目前比较流行的两种计算机。Python 易用但效率低,C++复杂但效率高。 所以有很多利用C来给python加速的工具。
这些工具一般分为两类: AOT(Ahead of Time) 和 JIT(Just in Time)
但是JIT有"cold start" 的问题,并且即便过了"cold start" ,其效率依旧不如经过人工精致优化的AOT,唯一的优点就是方便易用。 比如Numba,函数上面加个“@jit()”就完事儿了。

个人偏好更有水平的AOT方法,所以本文来介绍两个主流的AOT工具:Cython 和 pybind11

安装很简单:

conda install numpy
conda install cython
conda install pybind11

Cython

示例介绍

使用Cython一般需要三个文件:cython格式编写的库文件cythonfn.pyx, 用于编译cythonfn.pyx的setup.py, 以及调用cython库的python代码 main.py
考虑一种运算:

...
def calculate_z_serial_purepython(maxiter, zs, cs):"""Calculate output list using Julia update rule"""output = [0] * len(zs)for i in range(len(zs)):n = 0z = zs[i]c = cs[i]while abs(z) < 2 and n < maxiter:z = z * z + cn += 1output[i] = nreturn outputdef calc_pure_python(desired_width, max_iterations):"""Create a list of complex coordinates (zs) and complex parameters (cs),build Julia set"""x_step = (x2 - x1) / desired_widthy_step = (y1 - y2) / desired_widthx = []y = []ycoord = y2while ycoord > y1:y.append(ycoord)ycoord += y_stepxcoord = x1while xcoord < x2:x.append(xcoord)xcoord += x_step# build a list of coordinates and the initial condition for each cell.# Note that our initial condition is a constant and could easily be removed,# we use it to simulate a real-world scenario with several inputs to our# functionzs = []cs = []for ycoord in y:for xcoord in x:zs.append(complex(xcoord, ycoord))cs.append(complex(c_real, c_imag))output = [0] * len(zs)start_time = time.time()output = calculate_z_serial_purepython(max_iterations, zs, cs)    ## attention!!!end_time = time.time()secs = end_time - start_timeprint(calculate_z_serial_purepython.__name__ + " took", secs, "seconds")...
if __name__ == "__main__":secs = calc_pure_python(desired_width=1000, max_iterations=300)

经过分析不难发现(学习笔记 —— python代码耗时及内存占用测试方法 以及一些零碎的python小工具),进行循环计算的如下函数是耗时大户,也是我们最需要优化的对象。
在这里插入图片描述
由于*.pyx本身是基于python的,因此利用Cython我们可以看到python到C的演变过程,以及各阶段的优化力度。
现在优化前,我的耗时是这样的:
在这里插入图片描述

第一阶段优化

编写cythonfn.pyx文件, 把calculate_z_serial_purepython()函数原封不动copy过来:

def calculate_z_serial_purepython(maxiter, zs, cs):"""Calculate output list using Julia update rule"""output = [0] * len(zs)for i in range(len(zs)):n = 0z = zs[i]c = cs[i]while abs(z) < 2 and n < maxiter:z = z * z + cn += 1output[i] = nreturn output

编写setup.py

from distutils.core import setup
from Cython.Build import cythonizesetup(ext_modules=cythonize("cythonfn.pyx",compiler_directives={'language_level': "3"}))

运行setup.py 编译cythonfn.pyx
在这里插入图片描述
你会发现多了两个文件:
在这里插入图片描述
现在我们用这个cython编译后的calculate_z_serial_purepython()替代之前的:

import cythonfndef calc_pure_python(desired_width, max_iterations):...start_time = time.time()output = cythonfn.calculate_z_serial_purepython(max_iterations, zs, cs)   # cythonend_time = time.time()secs = end_time - start_timeprint(calculate_z_serial_purepython.__name__ + " took", secs, "seconds")

在这里插入图片描述

效率快提高一倍了,厉害吧?但如标题所说,这才刚刚开始!
如果你图懒,觉得这就够了,建议你不要考虑AOT方法了,去找JIT的吧,比如Numba。

第二阶段优化

Cython Annotation tool

在第二阶段开始前,我们先介绍一个Cython自己的可视化工具,可以把Cython编译后的底层C代码展示给我们:

运行命令:

cython -a cythonfn.pyx

会生成一个 cythonfn.html的文件,打开:
在这里插入图片描述
现在我们通过点击原代码对应行左边的加号便可展开看到被编译后的底层C代码(我只知道你看不懂哈哈)。
我们需要关注的是颜色的深浅,颜色越深表示对Python virtual machine的调用越多,也就越慢; 纯白色表示就可以认为是C代码。

好了,现在我们的目标就是:让这个界面越白越好!

优化方法

Cython中可以通过 cdef 命令直接定义C层次的变量,避免Python的type识别:

def calculate_z_serial_purepython(maxiter, zs, cs):"""Calculate output list using Julia update rule"""cdef unsigned int i,ncdef double complex z,coutput = [0] * len(zs)for i in range(len(zs)):n = 0z = zs[i]c = cs[i]while abs(z) < 2 and n < maxiter:z = z * z + cn += 1output[i] = nreturn output

现在看是不是白了许多?尤其是最耗时的第14行,纯白色!这意味着完全脱离的Python到达了C的层次!
在这里插入图片描述

现在重新用 setup.py 编译后的结果如下:
在这里插入图片描述
比起最开始提高了一个数量级不止,有没有种恐怖如斯的感觉?

第三阶段优化

我们都知道python中常用numpy来进行科学计算,会加速很多。但numpy是一种封装好的类,而C喜欢直接对raw_data,如double[] 操作,因此np.array直接传递给C有点臃肿。

好在python使用了一种缓存协议(buffer protoco),可以直接把缓存中的数据传给C。

import numpy as np
cimport numpy as npdef calculate_z_serial_purepython(int maxiter, double complex[:] zs, double complex[:] cs):"""Calculate output list using Julia update rule"""cdef unsigned int i,ncdef double complex z,ccdef int[:] output = np.empty(len(zs), dtype=np.int32)# output = [0] * len(zs)for i in range(len(zs)):n = 0z = zs[i]c = cs[i]# while abs(z) < 2 and n < maxiter:while  n < maxiter and z.real*z.real+z.imag*z.imag < 4:z = z * z + cn += 1output[i] = nreturn output

因为用了numpy,所以setup.py也得做相应改动:

from distutils.core import setup
from Cython.Build import cythonize
import numpysetup(ext_modules=cythonize("cythonfn.pyx",compiler_directives={'language_level': "3"}), include_dirs=[numpy.get_include()])           

现在相比之前是不是又白了很多?
在这里插入图片描述
在这里插入图片描述
又在第二阶段的基础上快了一倍。
没有很大提高是因为最耗时的第18行已经优化到极致了(不能更白了),针对这个函数只是优化了数据存取而已。

好了,经过三个阶段的优化,我们把原本的5.49s耗时缩减到了 0.17s,足足32倍!

比对下 JIT的Numba

函数上方加个“@jit()”

from numba import jit
@jit()
def calculate_z_serial_purepython_numba(maxiter, zs, cs, output):"""Calculate output list using Julia update rule"""for i in range(len(zs)):n = 0z = zs[i]c = cs[i]while abs(z) < 2 and n < maxiter:z = z * z + cn += 1output[i] = n

看看效果:
在这里插入图片描述
只能说和第一阶段的cython等同。
差的那几百毫秒怕是“cold start”造成的。

总结

方法Time Usage
无优化5.4992077350616455 seconds
cython第一阶段3.16896915435791 seconds
cython第二阶段0.37267589569091797 seconds
cython第三阶段0.17330312728881836 seconds
Numba3.286240577697754 seconds

pybind11

Links

Pybind11 github
Pybind11 doc

Introduction

Cython使用了用C编写的独立共享库,并将它们转换为Python模块。实际还是在Python上写库。
PyBind11相反,它提供C++ API,写的是C++代码库,然后转成Python模块。
Pybind11只支持C++,使用C++ 11特性。
此外,Pybind11的另一个优点是可以轻松处理来自Python的底层数据和类型(如Numpy array)
小知识:TensorFlow也用了Pybind11。

Pybind11的编译可以用g++或者Cmake,二者在VS上都要有很多配置,个人觉得很麻烦,所以建议摆脱win,到linux上去用

Implementation

考虑一个AXPY形式 的Python 矩阵运算:
y←αx+yy\leftarrow \alpha x +y y←αx+y
AXPY 有两种变体:SAXPY 和 DAXPY。 分别对应 Single Precision (32-bit float) 和 Double Precision (64-bit double)
这意味着我们需要区分double和float类型才能正常使用pybind。

以SAXPY 为例,编写核函数:

void axpy(size_t n, float a, float *x, float *y)
{for (size_t i = 0; i < n; i++)y[i] = a * x[i] + y[i];
}

前面我们说Pybind11可以轻松处理Numpy array,这还是float* , 没看到Numpy啊,难道这就行了,可以自动转换?当然不行,自动就需要多余的识别,自然就慢了。

别着急,我们再写个 wrapper to handle the Numpy arrays:(使用模板类型 py::array_t 来存取Numpy数据)

小知识:Python supports a very general and convenient way for exchanging data between libraries through a buffer protocol. In this sense, the py::array_t can be considered a buffer object that is restricted to Numpy arrays. The information of the buffer, such as the dimension and size, are stored in an information object called py::buffer_info.

#include 
using namespace py = pybind11;
void axpy_wrapper(float a, py::array_t py_x, py::array_t py_y)
{// Request for buffer informationpy::buffer_info x_buffer = py_x.request();py::buffer_info y_buffer = py_y.request();// check dimif (x_buffer.ndim != 1 || y_buffer.ndim != 1) {throw std::runtime_error("Error: dimension of vector should be 1");}// check shape matchif (x_buffer.shape[0] != y_buffer.shape[0]) {throw std::runtime_error("Error: size of X and Y do not match");}// extract raw pointerfloat *x = (float*)x_buffer.ptr;float *y = (float*)y_buffer.ptr;return axpy(x_buffer.shape[0], a, x, y);
}

For more information on format checking and performance optimization, refer to here and here

完整example.cpp代码如下:

#include 
#include      // Include header that supports Numpy arraysnamespace py = pybind11;void axpy(size_t n, float a, float *x, float *y)
{for (size_t i = 0; i < n; i++)y[i] = a * x[i] + y[i];
}void axpy_wrapper(float a, py::array_t py_x, py::array_t py_y)
{// Request for buffer informationpy::buffer_info x_buffer = py_x.request();py::buffer_info y_buffer = py_y.request();// check dimif (x_buffer.ndim != 1 || y_buffer.ndim != 1) {throw std::runtime_error("Error: dimension of vector should be 1");}// check shape matchif (x_buffer.shape[0] != y_buffer.shape[0]) {throw std::runtime_error("Error: size of X and Y do not match");}// extract raw pointerfloat *x = (float*)x_buffer.ptr;float *y = (float*)y_buffer.ptr;return axpy(x_buffer.shape[0], a, x, y);
}PYBIND11_MODULE(math_kernels, m)
{m.def("axpy", &axpy_wrapper);
}

使用g++编译:

$ g++ -O3 -Wall -shared -std=c++11 -fPIC $(python3 -m pybind11 --includes) example.cpp -o example$(python3-config --extension-suffix)

然后python中调用试试:

$ python
Python 3.8.0 (default, Nov  6 2019, 21:49:08) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import numpy as np
>>> import example
>>> x = np.random.random(5)
>>> y = np.random.random(5)
>>> a = 0.5
>>> z = a * x + y
>>> example.axpy(a, x, y)
>>> z
array([0.72109651, 0.80642238, 1.3581158 , 0.68677673, 0.47000476])
>>> y
array([0.67480229, 0.67799412, 0.97670182, 0.20975148, 0.35173144])
>>> z - y
array([0.04629421, 0.12842827, 0.38141398, 0.47702525, 0.11827332])
>>> 

哈哈哈,pybind11的结果完全是错的!!!
还是数据类型的问题呀! Numpy arrays的默认数据类型是 np.float64,也就是 double, 咱们上面C库里写的可是float哦~
这下你知道即便是 floatdouble 甚至常常在C代码中都会被忽略数据类型冲突的两种类型,在pybind11里搞错后有多大影响了吧?

但是你有没有发现,数据类型不一样它不给你报错,然后给了个错的离谱的结果! 这点是很坑的!因为pybind11会做自动类型转换,还转不对!所以大家养好习惯,在PYBIND11_MODULE中做一些声明,不让他自动转:

/* name of module-        ---a variable with type py::module_ to create the binding */
/* (这个才是python里   |       |                                                                                                */
/*  import的库名)     v       v                                                                                               */
PYBIND11_MODULE(example, m) {                            // Create a module using the PYBIND11_MODULE macrom.doc() = "pybind11 math module";    // optional module docstringm.def("sum", &sum, "sum two numbers in double"); // calls module::def() to create generate binding code that exposes sum()m.def("axpy", &axpy_wrapper, py::arg("a"), py::arg("x").noconvert(true), py::arg("y").noconvert(true)); // 第一个参数是python调用时用的函数名
}

现在我们编译再试:

$ python
Python 3.8.0 (default, Nov  6 2019, 21:49:08) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import numpy as np
>>> import example
>>> x = np.random.random(5).astype(np.float32)
>>> y = np.random.random(5)
>>> example.axpy(0.5, x, y)
Traceback (most recent call last):File "", line 1, in 
TypeError: axpy(): incompatible function arguments. The following argument types are supported:1. (a: float, x: numpy.ndarray[float32], y: numpy.ndarray[float32]) -> NoneInvoked with: 0.5, array([0.583261  , 0.4329178 , 0.1882612 , 0.94220173, 0.56544894],dtype=float32), array([0.65019744, 0.12314781, 0.04364232, 0.54704202, 0.18679054])
>>> 

这次你数据类型不对就别想跑通了嘿嘿。

但是我们C里定义的是模板T呀,我们加个double的声明不就完了:

#include 
#include     // Include header that supports Numpy arraysnamespace py = pybind11;template 
void axpy(size_t n, T a, T *x, T *y)
{for (size_t i = 0; i < n; i++)y[i] = a * x[i] + y[i];
}template 
void axpy_wrapper(T a, py::array_t py_x, py::array_t py_y)
{// Request for buffer informationpy::buffer_info x_buffer = py_x.request();py::buffer_info y_buffer = py_y.request();// check dimif (x_buffer.ndim != 1 || y_buffer.ndim != 1) {throw std::runtime_error("Error: dimension of vector should be 1");}// check shape matchif (x_buffer.shape[0] != y_buffer.shape[0]) {throw std::runtime_error("Error: size of X and Y not match");}// extract raw pointerT *x = (T*)x_buffer.ptr;T *y = (T*)y_buffer.ptr;return axpy(x_buffer.shape[0], a, x, y);
}PYBIND11_MODULE(math_kernels, m)
{m.def("axpy", &axpy_wrapper,  py::arg("a"), py::arg("x").noconvert(true), py::arg("y").noconvert(true));m.def("axpy", &axpy_wrapper, py::arg("a"), py::arg("x").noconvert(true), py::arg("y").noconvert(true));
}

再试试:

$ python
Python 3.8.0 (default, Nov  6 2019, 21:49:08) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import numpy as np
>>> import example
>>> x = np.random.random(5).astype(np.float32)
>>> y = np.random.random(5).astype(np.float32)
>>> a = 0.5
>>> z = a * x + y
>>> example.axpy(a, x, y)
>>> y - z
array([0., 0., 0., 0., 0.], dtype=float32)
>>> x = np.random.random(5).astype(np.float64)
>>> y = np.random.random(5).astype(np.float64)
>>> z = a * x + y
>>> example.axpy(a, x, y)
>>> y - z
array([0., 0., 0., 0., 0.])
>>>

OK,现在你只要保证输入数据类型同时是 float 或者同时是 double 就行啦。

现在你可以尝试一下用pybind11实现上面Cython中提到的**calculate_z_serial_purepython()**看看有多少提高啦!

Implementation Cmake

如果只是简单测试,g++当然足够了。但大多数情况下,我们需要Cmake来管理一整个项目,因此如果使用Cmake编译pybind11变得很重要!

TODO

cytpes

还有一种方法与pybind11很像, ctypes,都是用g++编译后调用。
不过个人感觉pybind11更方便一点,效率暂且没做比较
下面是ctypes的一个简单示例:

定义C源文件 sum_src.c.

#include 
#include "sum.h"double sum(double a, double b)
{printf("hello world from sum()\n");return a + b;
}
```h
和头文件**sum.h**:
```cpp
#ifndef __SUM_H__
#define __SUM_H__
double sum(double, double);
#endif

gcc编译为 libsum.so:

gcc -fPIC -shared -o libsum.so sum_src.c

然后python调用:

Python 3.8.0 (default, Nov  6 2019, 21:49:08) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import ctypes
>>> _libsum = ctypes.CDLL('libsum.so')	# 加载共享库,定义函数签名,以便ctypes可以转换类型
>>> _libsum.sum.argtypes = [ctypes.c_double, ctypes.c_double]	# 指定参数类型
>>> _libsum.sum.restype = ctypes.c_double 	# 指定返回类型
>>> _libsum.sum(ctypes.c_double(3.1), ctypes.c_double(3.3)) # 函数调用
hello world from sum()
6.4
>>> 

python调的时候还需要再指定遍类型,是不是有点麻烦?另外,所有参数都必须转换为各自的基本数据类型, 也就是你用numpy的话还得转一下,而不是pybind那种可以直接调用。
ctypes的一个特性是可以定义回调函数以允许C程序调用Python函数。有关更多信息,请参阅 文档。

Cython & pybind11 性能比较

容易理解,Cython还是在pyhon上编写的,依旧需要Python的一些相对于C++比较多余的解码机制;而pybind11直接在C++写库,并且一定要指定数据类型,没有多余的解码,自然会快一些。

Pybind11 github
Pybind11 doc

相关内容

热门资讯

政策送到田埂上 景阳镇土关村宣讲组成员在田间地头宣讲农业产业化发展相关政策。通讯员 李旭 摄本报记者 李延平 通讯员...
深化投融资综合改革 不拘一格力...   证券时报记者 张淑贤  上交所党委书记、理事长邱勇近日在2026北外滩财富与文化论坛上表示,站在...
省委编办驻村工作队宣讲政策送温... 本报讯 (记者 郑思哲) 为深入学习贯彻2026年中央和省委一号文件精神,扎实推进万名干部下乡活动,...
青海将建设239处水旱灾害防御... 本报讯 (记者 王臻) 3月29日,记者从省水利厅获悉,今年全省水旱灾害防御工作将转变洪水传统治理管...
25家公募去年经营业绩曝光 2...   证券时报记者 余世鹏  随着上市公司年报持续披露,旗下公募基金公司的经营情况浮出水面。  据证券...