初探: 通过pyo3用rust为python写扩展加速
眾所周知，python性能比較差，尤其在計(jì)算密集型的任務(wù)當(dāng)中，所以機(jī)器學(xué)習(xí)領(lǐng)域的算法開發(fā)，大多是將python做膠水來用，他們會(huì)在項(xiàng)目中寫大量的C/C++代碼然后編譯為so動(dòng)態(tài)文件供python加載使用。那么時(shí)至今日，對(duì)于不想學(xué)習(xí)c/c++的朋友們，rust可以是一個(gè)不錯(cuò)的替代品，它有著現(xiàn)代化語言的設(shè)計(jì)和并肩c/c++語言的運(yùn)行效率。
本文簡單介紹使用rust為python計(jì)算性質(zhì)的代碼做一個(gè)優(yōu)化，使用pyo3庫為python寫一個(gè)擴(kuò)展供其調(diào)用，咱們下面開始，來看看具體的過程和效率的提升。（PS：本文只是拋磚引玉，初級(jí)教程）
我的臺(tái)式機(jī)環(huán)境:
設(shè)備名稱 DESKTOP
處理器 12th Gen Intel(R) Core(TM) i7-12700 2.10 GHz
機(jī)帶 RAM 32.0 GB (31.8 GB 可用)
系統(tǒng)類型 64 位操作系統(tǒng), 基于 x64 的處理器
1. python代碼
首先給出python代碼，這是一個(gè)求積分的公式：
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, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))
執(zhí)行這段代碼花費(fèi)了: Elapsed: 32.59504199028015 s
2. rust
use std::time::Instant;
fn main() {
let now = Instant::now();
let result = integrate_f(1.0, 100.0, 200000000);
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;
}
執(zhí)行這段代碼花費(fèi)了: Elapsed: 10.80 s
3. 通過pyo3寫擴(kuò)展
首先創(chuàng)建一個(gè)項(xiàng)目，并安裝 maturin 庫:
# (replace demo with the desired package name)
$ mkdir demo
$ cd demo
$ pip install maturin
然后初始化一個(gè)pyo3項(xiàng)目:
$ maturin init
? ?? What kind of bindings to use? · pyo3
? Done! New project created demo
整體項(xiàng)目結(jié)構(gòu)如下:
Cargo.toml中的一些字段含義：https://doc.rust-lang.org/cargo/reference/manifest.html
.
├── Cargo.toml // rust包管理文件，會(huì)在[lib]中聲明目標(biāo)擴(kuò)展包的名稱
├── src // rust源文件目錄, 你的擴(kuò)展文件就寫在這里，這個(gè)目錄是maturin初始化的時(shí)候自動(dòng)創(chuàng)建
│ └── lib.rs // 擴(kuò)展文件
├── pyproject.toml // python包管理文件，里面有python的包名字定義
├── .gitignore
├── Cargo.lock
└── demo // 我們的目標(biāo)模塊名稱，需手動(dòng)創(chuàng)建
├── main.py // 用來測試的文件
└── demo.cp312-win_amd64.pyd // 編譯生成的動(dòng)態(tài)鏈接庫文件，供import給python使用
在src/lib.rs 下寫入:
use pyo3::prelude::*;
/// Caculate the integrate.
#[pyfunction]
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;
}
/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule]
fn demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(integrate_f, m)?)?;
Ok(())
}
然后我們通過兩種途徑來使用它：
3.1 將擴(kuò)展安裝為python包
$ maturin develop
這個(gè)命令會(huì)將rust代碼轉(zhuǎn)為python的包，并安裝在當(dāng)前python環(huán)境內(nèi)。通過 pip list 就能看到。
3.2 編譯成動(dòng)態(tài)文件從python加載
$ maturin develop --skip-install
--skip-install 命令會(huì)產(chǎn)生一個(gè) pyd 文件而不是將其安裝為python的包 - demo.cp312-win_amd64.pyd 文件在當(dāng)前目錄下，然后python可以直接導(dǎo)入使用。
另外還有一個(gè)指令替換
--skip-install為--release會(huì)生成一個(gè)xxxx.whl文件，也就是Python pip安裝的包源文件。
首先我們?cè)趓ust項(xiàng)目下，與 Cargo.toml 同級(jí)目錄下，創(chuàng)建一個(gè) demo 目錄，然后我們寫一個(gè)python文件 demo/main.py，下面是擴(kuò)展的執(zhí)行效果：
import time
import demo
s = time.time()
print(demo.integrate_f(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))
花費(fèi)時(shí)間為：Elapsed: 10.908721685409546 s
可以看到python的執(zhí)行時(shí)間是rust和rust擴(kuò)展的3倍時(shí)長，單進(jìn)程看著好像不太大是吧，下面還有并行版本。
4 并行加速
4.1 python多進(jìn)程效果
Python多進(jìn)程很神奇，你寫的不好的話，他比單進(jìn)程下還要慢。
import math
import os
import time
from functools import partial
from multiprocessing import Pool
def sum_s(i: int, dx: float, a: int):
return math.e ** (-((a + i * dx) ** 2))
def integrate_f_parallel(a, b, N):
s: float = 0.0
dx = (b - a) / N
sum_s_patrial = partial(sum_s, dx=dx, a=a)
with Pool(processes=os.cpu_count()) as pool:
tasks = pool.map_async(sum_s_patrial, range(N), chunksize=20000)
for t in tasks.get():
s += t
return s * dx
if __name__ == "__main__":
s = time.time()
print(integrate_f_parallel(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))
花費(fèi)時(shí)間: Elapsed: 18.86696743965149 s，比單進(jìn)程下時(shí)間少了不到一半。
4.2 rust多線程加速給python使用
如果我們使用rust的并行庫，將rust進(jìn)一步加速，速度效果更明顯:
將上面的 integrate_f 替換為下面的多線程版本:
use pyo3::prelude::*;
use rayon::prelude::*;
#[pyfunction]
fn integrate_f_parallel(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;
}
/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule]
fn demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(integrate_f_parallel, m)?)?;
Ok(())
}
執(zhí)行上一個(gè)標(biāo)題3.2的步驟，然后在引入python使用：
import time
import demo
s = time.time()
print(demo.integrate_f_parallel(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))
花費(fèi)時(shí)間為：Elapsed: 0.9684994220733643 s。這比原先的單線程rust版本又快了10倍。但是差不多是python并行版本的18倍左右，是python單進(jìn)程版本的32倍左右。如果我們將一些關(guān)鍵的性能通過rust重寫，可以節(jié)省的時(shí)間成本是十分可觀的。
總體來看，整體的使用過程相當(dāng)簡潔方便，難點(diǎn)就是rust的學(xué)習(xí)曲線高，使用起來需要花費(fèi)精力，但是還是可以慢慢嘗試去使用它優(yōu)化已有的項(xiàng)目性能，哪怕只是一個(gè)功能函數(shù)，熟能生巧，一切慢慢來。
以上數(shù)據(jù)比較僅供參考，不同機(jī)器可能差異也不同。
- 參考文章:
- https://saidvandeklundert.net/learn/2021-11-18-calling-rust-from-python-using-pyo3/
- https://pyo3.rs/
- https://www.maturin.rs/tutorial
- https://doc.rust-lang.org/cargo/
總結(jié)
以上是生活随笔為你收集整理的初探: 通过pyo3用rust为python写扩展加速的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： Treap（平衡树）
- 下一篇： 如何设计一个高并发系统？