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

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

这里计算的积分如下:
∫ b a e − x 2 \int_b^a e^{-x^2} baex2
未作优化的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=False)
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左右。如果再使用 --release 参数去run,时间会被优化到 1.44 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秒。如果再使用 --release 参数去run,时间会被优化到 0.14 s。

总结(单位:秒)

pythongogo(goroutine)rustrust(–release)rust(rayon)rust(rayon)(–release)pypynumba(no-fastmath)numba(fastmath)cython
22.055.830.433.561.440.750.141.981.240.630.96

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

  • 1
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 2
    评论
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值