vlambda博客
学习文章列表

高性能Python(三) 通过编译优化性能


我们此前谈到了性能分析的基本思路和代码优化的编程习惯,众所周知习惯是需要慢慢养成的。总有一些时候来不及做这么多分析,也来不及逐字逐句做代码优化。又或者你不愿意花费时间了解性能优化的种种方式,而是希望套用模版实现"一键优化"。这种“懒人版”的性能优化方式真的存在吗?我推荐以下的两种方式:

  1. 使用Cython将代码编译为 so(Linux平台)或者 dll(windows平台)
  2. 如果程序的计算、io部分可以并行,则可以使用多进程调用多核执行任务

使用支持GPU计算的库也满足这种懒人优化的要求,希望了解的读者可以搜索相关资料自行了解。

Part1编译代码

python是一种动态的高级语言。python语言虽然接近人类的自然语言,但是让机器难以理解,解释器为了理解python需要做更多(相较其他语言来说)的工作。python因为其逐行解释执行的特征,天生在执行效率上具有劣势。显然,要提高python编译器执行代码的效率,就是让它做尽量少的工作,直接将python代码固定为机器码,可以大大提高执行的效率。这个过程,就叫做编译。

这种优化,适合整个代码中使用原生python的地方比较多的情形,尤其是使用嵌套循环比较多的情景。如果你的代码大量使用numpy进行了向量化处理,那么编译代码带来的提升将不会很明显,这时候采用Numba重新实现你的计算可能是一种比较好的选择。并且Numba也支持GPU计算。

另外,编译代码会将你可读的py文件转换为二进制的so(Linux平台)或者dll(windows平台)文件,还能间接实现代码内容隐蔽的功能。假如你希望将开发的工具交给朋友使用,但又不希望直白地将代码内容(所表明的心迹与情感)明文展示给他,而你的朋友恰好又不是一位反编译专家,那么编译代码就是你的不二选择。你只用将编译好且不可阅读的二进制文件传输给他,既能够给他调用代码的方法,同时又保护了自己的知识产权。

Part2编译代码的工具

为python程序提供编译功能的工具有很多,根据编译方式可以分为提前编译(AOT)和即时编译(JIT)两种。AOT是在执行代码之前,将代码固定为机器码,执行过程中没有编译过程。JIT是告诉python解释器,在解释执行的过程中只编译对应的代码(每次执行都会产生额外的编译开销),对不需要编译的代码则保留解释执行的特性。如果你想在你的程序中使用JIT来进行优化,最好不要使用在频繁启动但是运行并不耗时的程序上,因为每次执行的固定编译开销可能让你的代码比不使用JIT更加耗时,这时候使用AOT是更好的选择。

  • AOT编译工具: CythonShedskinPythran
  • JIT编译工具: NumbaPyPyPyston

1Cython编译

如何编译

首先需要注意的是Cython不是一个只能用来编译python代码的工具,我们这里仅是使用Cython编译代码的功能来对python程序进行快速性能优化。使用Cython进行编译的代码并不复杂,可以直接套用如下的模版来实现:

# setup.py
import glob
import os
import setuptools
import sys

from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext
from Cython.Build import cythonize

if __name__ == '__main__':

    params = sys.argv[1]
    sys.argv.remove(params)

    if os.path.isdir(params):
        target = glob.glob(os.path.join(params,'*.py'))
    elif os.path.isfile(params):
        target = [params]

    try:
        print(f"cythonlizing {target}")
        setup(ext_modules=cythonize(target))
    except Exception as e:
        print(repr(e))

如果编译单个文件,直接执行:

python setup.py target_python_file.py build_ext

编译的时候将需要编译的代码都统一放到一个文件夹中folder即可:

python setup.py folder build_ext

windows平台下,会在对应的目录下生成build文件夹,.\build\lib.win-amd64-3.8内会存放对应py的pyd二进制文件,此pyd文件可以直接import使用,和使用py文件是一样的,不同的是pyd是二进制dll文件,一般人无法知道里面的具体内容。

在linux下也不用修改setup.py的代码,只需修改相应路径即可,编译好的so文件可以在对应目录的.\build\文件夹下找到。

效果测试

还是蒙特卡洛的例子。将之前比较慢的蒙特卡洛测试代码编译好:

# cal_pi.py

import argparse
import numpy as np

def cal_pi(n_sims):
    in_circle = 0
    for isims in range(n_sims):
        x, y = np.random.rand(2)
        if x**2 + y**2 < 1:
            in_circle += 1
    pi = in_circle/n_sims*4
    return pi

在window平台,执行如下代码,得到编译成对应pyd文件cal_pi.cp38-win_amd64.pyd

python setup.py cal_pi.py build_ext

在任意文件中导入编译好的pyd,开始测试:

# pyd文件测试
from cal_pi import cal_pi
%timeit cal_pi(10000000)
4.65 s ± 263 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

直接使用未编译的py文件获得的效果:

# py文件测试
from cal_pi import cal_pi
%timeit cal_pi(10000000)
5.16 s ± 230 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

也就是说什么都不用做,只需要比平时多一步编译的步骤,就能获得大概10%的性能提升。并且这个提升还是在本例使用了numpy向量化计算的情况下进行的,如果是纯python的环境还能有更大的提升。本例提供了一个可以方便将任何py文件编译为二进制文件的通用方法,通过这样的操作能够为程序在很短的时间之内做出一定加速效果,在迫在眉睫的优化情况下也是非常值得尝试的一种方式。

当然Cython能做的远远不止这些,如果你愿意学习,可以直接采用cython约定的语法,将python和c直接在编码过程中混合编程,能够获得更好的性能。但这种混合编程的语法设计,会让维护人员的学习和维护成本大大增加,也是需要关注的短板。

2Numba即时编译

Numba是python科学计算加速非常重要的库。作为anaconda的产品,Numba的信赖度和活跃程度非常可靠。

先说明,Numba的优化使用了numpy的函数。有一定的限制。如果你恰好使用了numpy并且不想改成纯python的形式,可以看看你使用的numpy函数是不是被Numba所支持。如果支持,那么恭喜你,可以直接通过一个装饰器获得性能的提升。在本例中,只使用了numpy的随机函数,此函数是被numba所支持的,可以直接使用。

(可在官网查询最新版本的Numba支持的 numpy函数)

import numpy as np
import numba

@numba.jit
def cal_pi(n_sims):
    in_circle = 0
    for isims in range(n_sims):
        x, y = np.random.rand(2)
        if x**2 + y**2 < 1:
            in_circle += 1
    pi = in_circle/n_sims*4
    return pi

优化的效果令大惊小怪的人(也就是我)震惊:

%timeit cal_pi(10000000)
1.66 s ± 149 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

同样的主函数,通过Numba优化后,只需要1.66s。比Cython编译好的版本还要快上1.8倍左右。这里Numba作弊使用了多核进行处理,而我们之前无论是使用pyd还是py都只使用了单核。不过通过如此简单的方式就能将你的代码并行起来,我很欣慰。

如果把系列第二篇文章中提出的向量化方式结合起来会是怎样的效果:

@numba.jit
def cal_pi_vec(n_sims):
    x = np.random.rand(n_sims)
    y = np.random.rand(n_sims)
    pi = (x**2+y**2 < 1).sum()/n_sims*4
    return pi

%timeit cal_pi_vec(10000000)
272 ms ± 28.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

def cal_pi_vec(n_sims):
    x = np.random.rand(n_sims)
    y = np.random.rand(n_sims)
    pi = (x**2+y**2 < 1).sum()/n_sims*4
    return pi

%timeit cal_pi_vec(10000000)
487 ms ± 22.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

执行后发现Numba最终将1000w次的蒙特卡洛模拟任务的耗时降低到了272ms,比起我们上次提到的向量化方法还快了接近一倍。比原生的python5.16s的耗时来说大约实现了19倍的性能提升。

Numba对于很多简单的计算来说优化效果还是简单并且明显的,但在一些复杂的场景中,并不是那么轻松就可以优化成功.本文蜻蜓点水地介绍一下Numba最简单有效的应用,在具体的业务场景中使用Numba,还需要因地制宜,具体情况具体分析。