快捷導(dǎo)航

利用Rust實(shí)現(xiàn)Python加速的技巧分享

更新時(shí)間：2023年09月12日 11:22:05 作者：shelgi

這篇文章主要想來和大家一起探討一下關(guān)于使用Rust對(duì)Python計(jì)算進(jìn)行加速的問題,文中的示例代碼講解詳細(xì),感興趣的小伙伴可以跟隨小編一起學(xué)習(xí)一下

背景

長期以來,Python由于易上手,有GC且生態(tài)強(qiáng)大等特點(diǎn)被廣泛使用,可是漸漸的人們也發(fā)現(xiàn)了它的不足,解釋型語言的運(yùn)行速度終究比不過編譯型,況且由于Python設(shè)計(jì)時(shí)的動(dòng)態(tài)數(shù)據(jù)類型一切皆對(duì)象(內(nèi)存都分配在堆上)等思想,也導(dǎo)致了運(yùn)行速度緩慢.

隨著實(shí)時(shí)性要求的不斷提升,在一些計(jì)算量大要求快速響應(yīng)的場景傳統(tǒng)的Python就很難滿足要求,所以隨之慢慢有了各種解決辦法:

用更高效的解釋器
用jit即時(shí)編譯加速
改寫成Cython加速
對(duì)GIL動(dòng)手,提高多線程性能
…

其中目前使用最廣泛,最有效的應(yīng)該是jit與Cython這兩種方案,jit即時(shí)編譯可以將部分需要解釋的代碼直接轉(zhuǎn)為機(jī)器碼從而實(shí)現(xiàn)加速(減少解釋時(shí)開銷);而Cython更絕直接將Python原地升級(jí),得到一個(gè)Cython這個(gè)Python與C的混血,可以通過Cython將代碼翻譯成C/C++的代碼再編譯成動(dòng)態(tài)庫文件供使用.可是這樣就會(huì)造成Python原本的語法被改的“四不像”,這些后面慢慢再談.

既然都用上動(dòng)態(tài)庫了,為什么不直接使用C++或者其他高效語言實(shí)現(xiàn),然后供Python調(diào)用呢?說到底,Cython不也是翻譯成C/C++的代碼編譯使用,只是為了方便Python開發(fā)人員才設(shè)計(jì)了這種類似于Python的語法.如果熟悉其他語言的話完全可以直接使用其他語言實(shí)現(xiàn)而不影響Python的基本語法.所以今天就來討論一下關(guān)于使用Rust對(duì)Python計(jì)算進(jìn)行加速的問題.

Rust加速Python計(jì)算

首先,來看看Rust實(shí)現(xiàn)和Python實(shí)現(xiàn)基本的速度對(duì)比,目標(biāo)是求斐波那契數(shù)列第n項(xiàng)的值.其中實(shí)現(xiàn)均采用遞歸調(diào)用,為了突出時(shí)間差異這里求第30項(xiàng)的值并重復(fù)50次

Python的實(shí)現(xiàn)與耗時(shí)如下:

import time
?
def fib(n:int) ->int:
 ? ?assert n>=0
 ? ?if n <= 1:
 ? ? ? ?return n
 ? ?return fib(n-1)+fib(n-2)
?
def main(test_times=50):
 ? ?start = time.time()
 ? ?for _ in range(test_times):
 ? ? ? ?fib(30)
 ? ?print(f"time cost {time.time()-start} s")
?
if __name__ == '__main__':
 ? ?main()

Rust的實(shí)現(xiàn)與耗時(shí)如下:

use std::time;
?
fn fib(n:i32)->u64{
 ? ?if n<=0{
 ? ? ? ?panic!("{} must be a postive number!",n);
 ?  }
 ? ?match n{
 ? ? ? ?1|2 => 1,
 ? ? ? ?_ => fib(n-1) + fib(n-2)
 ?  }
}
?
fn main() {
 ? ?let test_times = 50;
 ? ?let start = time::Instant::now();
 ? ?for i in 0..test_times{
 ? ? ? ?fib(30);
 ?  }
 ? ?println!("time cost {:?}",start.elapsed())
}
?

這差異,足足一百多倍.那看來使用Rust提速是完全可行的,那怎么將Rust與Python相結(jié)合呢?或者如何把Rust的代碼編譯供Python調(diào)用,這個(gè)時(shí)候可以使用pyo3,首先安裝一下maturin工具pip install maturin,然后配置一下項(xiàng)目的Cargo.toml

[package]
name = "speedup_python"
version = "0.1.0"
edition = "2021"
?
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
?
[lib]
name = "speed_python"
crate-type = ["cdylib"]
?
[dependencies]
pyo3 = { version = "0.19.2", features = ["extension-module"] }

這里編譯類型就設(shè)置為lib,關(guān)于多種不同lib類型的區(qū)別可以去看看Rust專欄之前的內(nèi)容.這里的name就是未來Python中調(diào)用的名字,下面再編寫lib.rs

use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
?
#[pyfunction]
pub fn fib(n:i32)->u64{
 ? ?if n<=0{
 ? ? ? ?panic!("{} must be a postive number!",n);
 ?  }
 ? ?match n{
 ? ? ? ?1|2 => 1,
 ? ? ? ?_ => fib(n-1) + fib(n-2)
 ?  }
}
?
#[pymodule]
fn speed_python(_py:Python,m:&PyModule)->PyResult<()>{
 ? ?m.add_wrapped(wrap_pyfunction!(fib))?;
 ? ?Ok(())
}

邏輯代碼基本沒有更改,只是添加了Rust實(shí)現(xiàn)Python module的代碼,這里的module name必須和toml中設(shè)置的name保持一致,否則也會(huì)無法導(dǎo)入.

最后運(yùn)行maturin develop就可以實(shí)現(xiàn)編譯,給Python調(diào)用了.

加速對(duì)比

現(xiàn)在,我們已經(jīng)實(shí)現(xiàn)了Rust的加速,是不是非常簡單而且調(diào)用的時(shí)候可以使用原本的Python語法而不用進(jìn)行任何更改.下面就來對(duì)比一下原始,numba,Cython,Rust四種方式的速度對(duì)比.

其中Cython實(shí)現(xiàn)cpy_fib.pyx如下

cpdef int c_fib(n:int):
 ? ?assert n>0
 ? ?if n in [1, 2]:
 ? ? ? ?return 1
 ? ?else:
 ? ? ? ?return c_fib(n - 1) + c_fib(n - 2)

然后寫setup進(jìn)行編譯

from distutils.core import setup,Extension
from Cython.Build import cythonize
?
setup(
 ? ?ext_modules=cythonize(Extension(
 ? ? ? ?'cpy_fib',
 ? ? ? ?sources=['./cpy_fib.pyx'],
 ? ? ? ?language='c'
 ?  )),
)

運(yùn)行python setup.py build_ext --inplace編譯,最后整體對(duì)比

import speed_python
import time
from cpy_fib import c_fib
from numba import jit
?
?
@jit(nopython=True)
def fib_jit(n: int) -> int:
 ? ?assert n > 0
 ? ?if n in [1, 2]:
 ? ? ? ?return 1
 ? ?else:
 ? ? ? ?return fib_jit(n - 1) + fib_jit(n - 2)
?
def fib(n: int) -> int:
 ? ?assert n > 0
 ? ?if n in [1, 2]:
 ? ? ? ?return 1
 ? ?else:
 ? ? ? ?return fib(n - 1) + fib(n - 2)
?
?
def test_speed(func,func_name:str,test_times=50):
 ? ?start = time.time()
 ? ?for _ in range(test_times):
 ? ? ? ?func(30)
 ? ?print(f"{func_name} speed up time cost {time.time() - start} s")
?
?
def main(test_times=50):
 ? ?test_speed(fib,"origin python")
 ? ?test_speed(fib_jit,"numba python")
 ? ?test_speed(speed_python.fib,"rust")
 ? ?test_speed(c_fib,"Cython")
?
?
if __name__ == '__main__':
 ? ?main()