快捷導(dǎo)航

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

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

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

背景

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

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

用更高效的解釋器
用jit即時編譯加速
改寫成Cython加速
對GIL動手,提高多線程性能
…

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

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

Rust加速Python計算

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

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

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的實現(xiàn)與耗時如下:

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)用,這個時候可以使用pyo3,首先安裝一下maturin工具pip install maturin,然后配置一下項目的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實現(xiàn)Python module的代碼,這里的module name必須和toml中設(shè)置的name保持一致,否則也會無法導(dǎo)入.

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

加速對比

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

其中Cython實現(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進行編譯

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

運行python setup.py build_ext --inplace編譯,最后整體對比

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()