python & cython & numba speed comparison

这篇文章简介一下python、cython、以及numba库在计算上的效率对比

Basel problem(巴塞尔问题)

$\pi^2=6\sum\limits_{n=1}^\infty{\frac{1}{n^2}}$

这里用巴塞尔问题作为引例，用以上代码方式计算pi值，通过cProfile库比较计算效率。

三种编译方式

1、python

# calc_pi.py
def recip_square(i):
    return 1. / i ** 2


def approx_pi(n=10000000):
    val = 0
    for k in range(1, n + 1):
        val += recip_square(k)
    print((6 * val) ** .5)


# profile_cal.py
import cProfile
import pstats

from tutorial.calc_pi import approx_pi


cProfile.runctx("approx_pi()", globals(), locals(), "Profile.prof")
s = pstats.Stats("Profile.prof")
s.strip_dirs().sort_stats("time").print_stats()

2、cython

# calc_pi_cy.pyx

# cython: profile=True
# cython: language_level=2

cimport cython


@cython.profile(False)
cdef inline double recip_square(int i):
    return 1. / (i ** 2)


def approx_pi(int n=10000000):
    cdef double val = 0
    cdef int k
    for k in range(1, n + 1):
        val += recip_square(k)
    print((6 * val) ** .5)


# profile_cal_cy.py
import cProfile
import pstats
import pyximport
pyximport.install()

from tutorial.calc_pi_cy import approx_pi


cProfile.runctx("approx_pi()", globals(), locals(), "Profile.prof")
s = pstats.Stats("Profile.prof")
s.strip_dirs().sort_stats("time").print_stats()

3、jit【jit的编译需要安装numba☞官网】

# calc_pi_jit.py
import numba


@numba.jit(nopython=True)
def recip_square(i):
    return 1. / i ** 2


@numba.jit(nopython=True)
def approx_pi(n=10000000):
    val = 0
    for k in range(1, n + 1):
        val += recip_square(k)
    print((6 * val) ** .5)


# profile_cal_jit.py
import cProfile
import pstats

from tutorial.calc_pi_jit import approx_pi


cProfile.runctx("approx_pi()", globals(), locals(), "Profile.prof")
s = pstats.Stats("Profile.prof")
s.strip_dirs().sort_stats("time").print_stats()

接下来分别使得n为不同的数值来看执行时间的差距。

1、n=10000000

# python
3.1415925580959025
Sun May 10 13:45:40 2020    Profile.prof

         10000005 function calls in 4.701 seconds

# cython
3.1415925580959025
Sun May 10 13:48:25 2020    Profile.prof

         5 function calls in 0.011 seconds

# jit
3.1415925580959025
Sun May 10 13:49:02 2020    Profile.prof

         519799 function calls (490055 primitive calls) in 0.428 seconds

分析：可以看到在n为一千万(1000 0000)时，python方式比cython慢了470多倍，比jit方式慢了10倍左右；

而jit比cython又慢了40倍左右，不过jit和cython的对比在小数据方面表现的差异并不是很大，只是略大，numba的jit编译方式在计算次数越多表现出的优势越明显，看下面。

2、n=100000000

# python
3.14159264498239
Sun May 10 13:57:28 2020    Profile.prof

         100000005 function calls in 49.787 seconds

# cython
3.14159264498239
Sun May 10 13:56:53 2020    Profile.prof

         5 function calls in 0.106 seconds

# jit
3.14159264498239
Sun May 10 13:57:34 2020    Profile.prof

         519983 function calls (490217 primitive calls) in 0.501 seconds

分析：可以看到在n为一亿(1 0000 0000)时，python耗时近50s，是cython的500倍左右，jit的100倍左右；

而jit比cython慢了5倍左右，这与上面的那种方式对比看来，与cython之间缩减了将近8倍的时间。

jit相较于之前1000 0000次计算，虽然这次增加了10倍计算量，python与cython都明显也是成10倍的耗时增加，但jit却只是增加了0.08s而已，相较于之前只是增加了17%左右的耗时。

看起来好像jit的这种方式优势凸显了。

3、n=1000000000

# python
3.14159264498239
Sun May 10 14:08:18 2020    Profile.prof

         1000000005 function calls in 485.110 seconds

# cython
3.14159264498239
Sun May 10 14:11:16 2020    Profile.prof

         5 function calls in 1.059 seconds

# jit
3.14159264498239
Sun May 10 14:10:38 2020    Profile.prof

         519965 function calls (490199 primitive calls) in 1.402 seconds

分析：可以看到本次n为十亿(10 0000 0000)时，python耗时485s，是上一次一亿次元算的10倍关系，是cython(cython本身也在成10倍增加计算时间)的485倍左右，是jit的345倍左右；

而jit愈来愈趋近于cython的执行效率，是cython的1.32倍左右。

4、n=2000000000

# cython
3.14159264498239
Sun May 10 14:43:43 2020    Profile.prof

         5 function calls in 2.080 seconds

# jit
3.14159264498239
Sun May 10 14:44:32 2020    Profile.prof

         519904 function calls (490148 primitive calls) in 2.502 seconds

分析：鉴于python是单线程，且随着数据成倍增长，这次n设置为二十亿(20 0000 0000)次，只比较cython和jit的运行时间差，但是二者差别在毫秒级别，并不是很大。

cython相较于前一次增加1s的时间，jit也基本同样增加了1s的时间。

总结

可以看出在进行cpu密集型运算时，原始python(即默认cpython解释器)的计算效率很低，原因在于GIL的限制，使得只能使用单线程。

cython会把相应的cython书写的代码编译为c语言，这就大大提高了执行效率。

而numba是使用了jit(Just-in-time)即时编译器，他会在程序运行期间实时编译，有别于python的默认解释器，所以速度要快很多很多。