用pypy、numba、cython分别对python的数学计算做性能优化[附带其他语言的版本]

本文只是一个简单的记录和比较，引用一个积分程序为例，尝试用几种方法优化Python性能，并将优化的结果做个比较。

这里计算的积分如下:
${\int }_b^a{e}^{-x^2}$
未作优化的Python代码如下:

import time

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


s = time.time()
print(integrate_f(1.0, 10.0, 100000000))
print("Elapsed: {} s".format(time.time() - s))

输出结果:

$ python main.py
0.13940280918664258
Elapsed: 22.051971435546875 s

可以看到最原始的python运行时间大概要22秒。

1. pypy运行

一般而言的Python都是指CPython解释器，CPython是广泛接受的Python标准。然而PyPy是另一个解释器，使用了JIT编译，和CPython高度兼容。不过PyPy的缺点是不支持C扩展模块，所以如果程序中用到Numpy，Scipy，就没法用PyPy优化了。

$ pypy main.py
0.139402809187
Elapsed: 1.98751616478 s

用PyPy执行上面的python程序，耗时1.98s。

2. numba运行

Numba是一个加速Python执行的库，可以用其中的JIT编译加速代码的执行。使用Numba JIT的代码如下:

• 不使用fastmath

import time

from numba import njit


@njit(fastmath=True)
def integrate_f(a, b, N):
    s = 0
    dx = (b - a) / N
    for i in range(N):
        s += 2.71828182846 ** (-((a + i * dx) ** 2))
    return s * dx


s = time.time()
print(integrate_f(1.0, 10.0, 100000000))
print("Elapsed: {} s".format(time.time() - s))

输出结果:

$ python main.py 
0.13940280918664258
Elapsed: 1.2390804290771484 s

可以看到使用numba的加速效果和pypy基本相当，但还是快一些。耗时1.24 s左右。

• 使用fastmath进一步加速

import time

from numba import njit


@njit(fastmath=True)
def integrate_f(a, b, N):
    s = 0
    dx = (b - a) / N
    for i in range(N):
        s += 2.71828182846 ** (-((a + i * dx) ** 2))
    return s * dx


s = time.time()
print(integrate_f(1.0, 10.0, 100000000))
print("Elapsed: {} s".format(time.time() - s))

输出结果:

$ python main.py 
0.13940280919433573
Elapsed: 0.6259806156158447 s

进一步加速，开启了fastmath功能，加速效果比不开启更又压缩了一半时间，达到了0.63 s左右。这速度达到了最原始python运行的36倍，相当恐怖。

3. cython运行

Cython将Python代码编译成C源码，再把C源码转换成Python扩展模块。用Cython改写Python代码，将动态类型用Cython中的静态类型声明后，可以大大提升执行的效率。

不过用Cython优化的步骤有点复杂。需要先生成Python扩展模块，然后在另外一个程序里import这个模块并调用模块中的方法。

Cython改写的代码如下：

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

然后需要在同级目录创建一个 setup.py 文件，然后生成扩展，假设上面的cython文件为 test.pyx，则 setup.py 文件写入:

from Cython.Build import cythonize
from setuptools import setup

setup(
    name="demo",
    ext_modules=cythonize("test.pyx"),
    zip_safe=False,
)

然后执行:

python setup.py build_ext --inplace

会在同级目录下生成 .c 文件和一些 .so 动态链接库，即编译好的可以供其他程序使用的代码和数据。
然后我们只需要new一个新的py文件，然后导入cython里面的那个函数，直接调用就行了，比如 main.py:

import time
from test import integrate_f_cython

s = time.time()
print(integrate_f_cython(1.0, 10.0, 100000000))
print(f"Elapsed: {time.time() - s} s")

输出结果:

$ python main.py 
0.13940280918664258
Elapsed: 0.9620242118835449 s

可以看到使用cython运行的时间也在1s之内，基本是原始python的20倍左右，而且对原始python的改动也不是很大。

4. go写一下

• 原始go代码

从python代码直接翻译到go代码:

package main

import (
	"fmt"
	"math"
	"time"
)

func integrate_f(a, b float64, N int64) float64 {
	s := 0.0
	dx := (b - a) / float64(N)

	for i := int64(0); i < N; i++ {
		_r := -math.Pow((a + float64(i)*dx), 2)
		s += math.Pow(2.71828182846, _r)
	}

	return s * dx
}

func main() {
	t := time.Now()
	res := integrate_f(1.0, 10.0, 100000000)
	fmt.Println(res)
	fmt.Printf("Elapsed: %v", time.Since(t))
}

运行结果:

$ go run .\main.go                                                                                                                                                             
0.13940280918664258
Elapsed: 5.8312043s

运行时间是5.83 s，将近 6s。虽然没有上面cython和numba加速的效果好，但是也比python原本快多了。

• 发挥go协程的能力

package main

import (
	"fmt"
	"math"
	"runtime"
	"time"
)

func integrate_f(a, b float64, N int64) float64 {
	sum := make(chan float64, runtime.NumCPU())
	dx := (b - a) / float64(N)
	for i := 0; i < runtime.NumCPU(); i++ {
		go func(id int) {
			start := int64(id) * (N / int64(runtime.NumCPU()))
			end := start + (N / int64(runtime.NumCPU()))
			if id == runtime.NumCPU()-1 {
				end = N
			}
			s := float64(0)
			for j := start; j < end; j++ {
				x := a + float64(j)*dx
				s += math.Exp(-math.Pow(x, 2))
			}
			sum <- s
		}(i)
	}
	total := float64(0)
	for i := 0; i < runtime.NumCPU(); i++ {
		total += <-sum
	}
	return total * dx
}

func main() {
	t := time.Now()
	res := integrate_f(1.0, 10.0, 100000000)
	fmt.Println(res)
	fmt.Printf("Elapsed: %v s", time.Since(t).Seconds())
}

运行结果:

$ go run .\main.go                                                                                                                                                             
0.13940280919491324
Elapsed: 0.4314864 s

当go开启了协程后，运算时间达到了 0.43 s。go的协程并发执行相当给力。

5. rust测验

• 翻译python的版本，无并行加速

use std::time::Instant;

fn main() {
    let now = Instant::now();
    let result = integrate_f(1.0, 10.0, 100000000);
    println!("{}", result);

    println!("Elapsed: {:.2} s", now.elapsed().as_secs_f32())
}

fn integrate_f(a: f64, b: f64, n: i32) -> f64 {
    let mut s: f64 = 0.0;
    let dx: f64 = (b - a) / (n as f64);

    for i in 0..n {
        let mut _tmp: f64 = (a + i as f64 * dx).powf(2.0);
        s += (2.71828182846_f64).powf(-_tmp);
    }

    return s * dx;
}

运行结果:

cargo run
   Compiling noob v0.1.0 (/home/runstone/work/code/daily/py3/accelerate/noob)
    Finished dev [unoptimized + debuginfo] target(s) in 0.44s
     Running `target/debug/noob`
0.13940280918664258
Elapsed: 3.56 s

直接运行时间是 3.56 s左右。

• 使用了rayon并行计算库:

use rayon::prelude::*;
use std::time::Instant;

fn main() {
    let now = Instant::now();
    let result = integrate_f(1.0, 10.0, 100000000);
    println!("{}", result);

    println!("Elapsed: {:.5} s", now.elapsed().as_secs_f32())
}

fn integrate_f(a: f64, b: f64, n: i32) -> f64 {
    let dx: f64 = (b - a) / (n as f64);

    let s: f64 = (0..n)
        .into_par_iter()
        .map(|i| {
            let x = a + i as f64 * dx;
            (2.71828182846_f64).powf(-(x.powf(2.0)))
        })
        .sum();

    return s * dx;
}

运行结果:

cargo run
    Finished dev [unoptimized + debuginfo] target(s) in 0.02s
     Running `target/debug/noob`
0.13940280919481815
Elapsed: 0.75 s

并行运算后的时间是0.75秒。

总结(单位:秒)

python	go	go(goroutine)	rust	rust(rayon)	pypy	numba(no-fastmath)	numba(fastmath)	cython
22.05	5.83	0.43	3.56	0.75	1.98	1.24	0.63	0.96

从表格上面可以看到，这几种优化方式对python的执行速度都有高倍的优化。不考虑换语言的操作，仅仅从python的方式下手，numba的效果最好，也最方便，但是局限也很多，如果函数里面涉及到其他非计算性语句或逻辑，numba可能会报错。cython虽然也不错，但是对于不熟悉cython和c语法的人来说，并不是很方便。