提升Python性能 - Numba与Cython

背景

这篇文章的背景是在最近的工作中涉及到了一些计算密集型任务,这些计算密集型任务或多或少触发了一些之前几乎没有关心过的Python性能问题,所以写下这篇文章分析Python的性能问题,并调研了一些对应的改善方案(Numba、Cython)。

坦白地说,在过往用Objective-C写iOS应用的经历中,除了一些面试和工作中常见的关于一些UI组件渲染的性能问题外,几乎没有关心过代码执行的效率(性能)问题。这次也正好是一个契机使我有机会复习一些本科学过的知识。

二维数组求和

首先让我们看一段简单的Python代码,这段代码定义了一个函数,其功能是对一个np.ndarray类型的二维数组求和,并返回结果:

def arr_sum(src_arr):
    res = 0
    shape = src_arr.shape
    for r in range(0, shape[0]):
        for c in range(0, shape[1]):
            res += src_arr[r][c]
    return res

这段代码没什么特别的,如果我们的src_arr.shape128x128,上面这个代码片段执行1000次,在我的机器上(i7 6700K、32G),大概需要3.7857s,同样的,没有对比就没有伤害,同样的代码,如果我们用C++重写一下,大概会是这个样子:

double sum(double arr[][128], int row) {
    double res = 0;
    for (int i = 0; i < row; ++i) {
        for (int j = 0; j < 128; ++j) {
            res += arr[i][j];
        }
    }
    return res;
}

这段代码的执行时间将会是0.035s,即大约35ms,可以看出,在这个场景下(当然,在实际的项目或者研究中,根据问题规模的不同,某个代码片段的实现也会不尽相同),还是可以有一个粗糙的结论:大概Python比C++慢了100倍。

虽然这个结论不是非常严谨,例如,我们出于某种面向对象的考虑,希望这段代码不是非常的胶水,可能会用std::vector或者某些容器替换一个C风格的二维数组:

template<typename T>
T sum(vector<vector<T>>& arr, int row, int col) {
    T res = 0;
    for (int i = 0; i < row; ++i) {
        for (int j = 0; j < col; ++j) {
            res += arr[row][col];
        }
    }
    return res;
}

而上面这段粗糙的实现对应的执行时间将会是0.060s,即60ms,但是仍然要比Python快出两个数量级,依然没有问题。

甚至如果你不幸将:

T sum(vector<vector<T>>& arr, int row, int col)

写成了:

T sum(vector<vector<T>> arr, int row, int col)

而导致函数调用时复制整个二维数组,也只需要4.3s,而对比Python的3.75s,反而会感觉还没慢多少。

为什么Python慢?

那么为什么Python会显得慢呢?首先,Python通常被称作解释型语言,是相对于像C++这样的编译型语言来说的。

事实上py文件也会被编译,但是并不像C++,或者是其他静态强类型编译型语言那样,通过预处理、编译、汇编、链接这样的过程最终得到机器码。py文件,即Python的源代码通常会在运行时被解释器先解释为字节码,然后交由虚拟机将字节码翻译成机器码执行,而这一步就很尬了。

事实上,也正是因为这样,我们往往才因此获得Python在运行时一些非常强大的特性,例如generator,利用generator我们可以做一些非常神奇的事情,例如协程等。

但是另一方面,Python的解释器和虚拟机翻译并执行字节码的过程带来了很大的性能开销,一个直觉的解释是:由于没有原生的编译时类型检查,所有的类型的检查都被移交给了运行时,执行一行Python代码很可能需要做不只一行的类型检查、边界检查等等。

这里其实并不打算详细探究Python字节码的编译与执行,只是简单的通过一个例子大致说明一行Python代码是如何被解释和执行的,

考虑源文件test.py,他们的实现很简单,其中test.py的实现大概是这样的:

def add(x, y):
    res = x + y
    print('Res: ' + res)

如你所见的,计算两个值的和,然后将结果打印到标准输出。

我们通过dis模块将test.py翻译成可读的字节码指令,将会是这样:

Disassembly of add:
  2           0 LOAD_FAST                0 (x)
              3 LOAD_FAST                1 (y)
              6 BINARY_ADD
              7 STORE_FAST               2 (res)
  3          10 LOAD_GLOBAL              0 (print)
             13 LOAD_CONST               1 ('Res: ')
             16 LOAD_FAST                2 (res)
             19 BINARY_ADD
             20 CALL_FUNCTION            1 (1 positional, 0 keyword pair)
             23 POP_TOP
             24 LOAD_CONST               0 (None)
             27 RETURN_VALUE

即,首先Python的源代码py文件将会被解释器翻译成上类似上面的东西,但是具体可能会更加复杂一点,但是我想到这一步应该已经能说明一些问题了,然后这个字节码会被Python虚拟机尝试执行,而在执行时则会有很多运行时的检查,然后才是转化为真正的机器码交由CPU执行。

而C++就不同了,因为他没有这一步。

其实,字节码也不是每次都是从磁盘读py文件在运行时编译的,事实上,每个在运行时被编译的py文件将会产生PyCodeObject对象,这一步操作是在解释import语句时执行的,PyCodeObject对象创建后,一方面将会根据需要被虚拟机继续转化为PyFrameObject对象进行后续的机器码翻译和执行工作,另一方面会带着一个“最后修改日期”的字段被缓存在磁盘上,通常,这个位置是在源代码同级目录下__pycache__文件夹中,这些被缓存的PyCodeObject对象将与源文件同名,只是扩展名为pyc,以便下次运行时直接读取缓存,从而节约编译字节码的时间。

所以,一个粗糙的结论可能是这样的,Python由于要在运行时编译和解释执行字节码,而且这个过程中参与了很多类似运行时类型检查的操作等一系列其他操作,从而产生了很多额外开销,降低了性能。

如何提速?

那么如何提速?在本文我们调研了两种方案,分别是NumbaCython,接下来我们将分别简述它们的加速原理,并给出一些示例代码,并做一些简单的性能对比实验。

Numba

首先我们介绍Numba,先引一段官网文档的介绍:

Numba is a just-in-time compiler for Python that works best on code that uses NumPy arrays and functions, and loops. The most common way to use Numba is through its collection of decorators that can be applied to your functions to instruct Numba to compile them. When a call is made to a Numba decorated function it is compiled to machine code “just-in-time” for execution and all or part of your code can subsequently run at native machine code speed!

其中有两句话比较重要:

Numba is a just-in-time compiler for Python that works best on code that uses NumPy arrays and functions, and loops.

Numba是一个JIT编译器,它和Numpy的数组和函数以及循环一起用时,效果最佳。

另一句是:

When a call is made to a Numba decorated function it is compiled to machine code “just-in-time” for execution and all or part of your code can subsequently run at native machine code speed!

如果一个调用被Numba装饰器修饰,那么它将被JIT机制编译成机器码执行,性能堪比本地机器码的速度。

同样,它的原理简介也能在官网文档中找到:

Numba reads the Python bytecode for a decorated function and combines this with information about the types of the input arguments to the function. It analyzes and optimizes your code, and finally uses the LLVM compiler library to generate a machine code version of your function, tailored to your CPU capabilities. This compiled version is then used every time your function is called.

简要的概括即是,Numba通过一个装饰器读某些调用的字节码,并为它们的参数等添加类型信息,尝试优化代码后,通过LLVM编译器直接生成对应的机器码。

思想其实很明确,一方面是在字节码被执行前添加类型信息,然后省去原始调用的字节码被执行的过程,通过LLVM直接编译机器码,从而节省了性能开销,事实也确实如此。

通过引入numba模块,原始的Python代码将会被改写成这样:

import numba as nb

@nb.njit()
def arr_sum(src_arr):
    res = 0
    shape = src_arr.shape
    for r in range(0, shape[0]):
        for c in range(0, shape[1]):
            res += src_arr[r][c]
    return res

可以发现其实没什么变化,除了加了一个@nb.njit()的装饰器外,没有对原函数做任何改动,这其实也是numba的方便之处(与后续的Cython方案对比),事实上@nb.njit()会尽其所能去寻找能被numba的JIT机制添加类型信息并翻译成机器码的对象,如果失败了,运行起来的效果甚至会比原始的Python代码还慢。

在它的文档的开头也就提到,它和Numpy的数组和函数以及循环一起用时,效果最佳,同时文档也给出了一个暂时不支持pandas类型的例子。

详细的numba实现原理已经超出了本文做简单调研的范围,可能我们会有后续几篇文章讨论这个问题。

一个快速的入门文档可以参见:

https://numba.pydata.org/numba-doc/latest/user/5minguide.html

最后,添加numba装饰器后,代码片段对大小为128x128的二维数组求和,运行1000次时间为0.017122s,即17ms,比CXX还要快。

而JIT首次尝试编译求和函数代码的约为0.151606s,即150ms,所以,一个粗糙的结论是,如果这段代码确实性能开销较大,且被调用频率相对较高,那么一个短暂的编译时间还是可以被接受的。

Cython

接下来我们介绍Cython,

Cython是在Python中实现C-Extensions的一种方案,简单的理解是,Python提供了一些与CXX的Lib相互调用的机制,而能通过import作为模块来进行调用的C或者C++的Lib,就是C-Extensions,有很多方案可以用来实现C-Extensions,例如Swig等,而Cython就是其中一种。

同样引一段官方文档中关于Cython的介绍:

Cython is an optimising static compiler for both the Python programming language and the extended Cython programming language (based on Pyrex). It makes writing C extensions for Python as easy as Python itself.

它的核心精神是,Cython将Cython语言(一种基于Python的扩展语言)写的pyx文件直接编译成C extensions,从而获得近乎于写CXX语言的性能。

我们直接看一个Cython改写的二维数组求和代码片段,这个代码片段的文件名将会是func.pyx,可以注意到到与*.py结尾的文件不同,Cython的代码将是以pyx结尾。

cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)

def arr_sum(double[:, :] src_arr):

    cdef double res = 0
    cdef Py_ssize_t[:] shape = src_arr.shape

    cdef Py_ssize_t r = 0
    cdef Py_ssize_t c = 0

    for r in range(0, shape[0]):
        for c in range(0, shape[1]):
            res += src_arr[r][c]
    return res

我们用了一些cdef关键字,来在定义变量时指明它们的类型,同时,我们使用了形如double[:, :]这样的关键字,它代表了Python中的MemoryView,即内存视图。简而言之,内存视图可以快速索引值,通过内存视图,我们可以避开繁琐的Python对象引用流程,直接访问一个二维数组某个下标值,如果不经转置,它在内存上应该是连续的,永远是通过一个基地址加上一个偏移量。

Cython文档在Typed Memory Views一节详细的介绍了这个机制,这里就不在赘述了。

https://cython.readthedocs.io/en/latest/src/userguide/memoryviews.html

看上去我们在写带类型标识符的Python,实际上这些代码都会被Cython先解释称CXX,然后编译成.so(Linux),通过编写对应的setup.py

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

setup(
    ext_modules=cythonize(["playground/cyfunc/func.pyx"], annotate=True)
)

打开annotate=True,Cython会替我们生成一份源码分析,如下图,详细的展示了pyx文件是如何生成CXX代码的,同时,黄色对应的行说明这行有Python调用,可能会影响能:

image.png

通过这种方式,代码片段对大小为128x128的二维数组求和,运行1000次时间为0.0181s,约18ms

性能对比

最后我们给出了四组实验的结果,代码片段对大小为128x128的二维数组求和,运行1000次时间如下:

Total cost time for func: py_func, call 1000 times: 3.803216s.
Total cost time for func: np_func, call 1000 times: 0.343562s.
Total cost time for func: nb_func, call 1000 times: 0.017122s.
Total cost time for func: cy_func, call 1000 times: 0.018159s.

它们分别代表了原始Python、Numpy、Numba、Cython对应的性能。

需要注意的是,numba的编译时间约为150ms

Time cost for numba first call: 0.151606s.

希望在之后的几篇文章中,将讨论NumbaCython的实现细节。

以上。

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 212,222评论 6 493
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 90,455评论 3 385
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 157,720评论 0 348
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 56,568评论 1 284
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 65,696评论 6 386
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 49,879评论 1 290
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 39,028评论 3 409
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 37,773评论 0 268
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 44,220评论 1 303
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 36,550评论 2 327
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 38,697评论 1 341
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 34,360评论 4 332
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 40,002评论 3 315
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 30,782评论 0 21
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,010评论 1 266
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 46,433评论 2 360
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 43,587评论 2 350

推荐阅读更多精彩内容