初探: 通过pyo3用rust为python写扩展

众所周知，python性能比较差，尤其在计算密集型的任务当中，所以机器学习领域的算法开发，大多是将python做胶水来用，他们会在项目中写大量的C/C++代码然后编译为so动态文件供python加载使用。那么时至今日，对于不想学习c/c++的朋友们，rust可以是一个不错的替代品，它有着现代化语言的设计和并肩c/c++语言的运行效率。

本文简单介绍使用rust为python计算性质的代码做一个优化，使用pyo3库为python写一个扩展供其调用，咱们下面开始，来看看具体的过程和效率的提升。（PS：本文只是抛砖引玉，初级教程）

我的台式机环境:

设备名称 DESKTOP
处理器 12th Gen Intel? Core? i7-12700 2.10 GHz
机带 RAM 32.0 GB (31.8 GB 可用)
系统类型 64 位操作系统, 基于 x64 的处理器

1. python代码

首先给出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, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))

执行这段代码花费了: Elapsed: 62.90621304512024 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;
}

执行这段代码花费了: Elapsed: 12.19 s

3. 通过pyo3写扩展

首先创建一个项目，并安装 maturin 库:

# (replace string_sum with the desired package name)
$ mkdir demo
$ cd demo
$ pip install maturin

然后初始化一个pyo3项目:

$ maturin init
? 🤷 What kind of bindings to use? · pyo3
  ? Done! New project created demo

在src/lib.rs 下写入:

use pyo3::prelude::*;

/// Formats the sum of two numbers as string.
#[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 string_sum(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(integrate_f, m)?)?;
    Ok(())
}

然后我们通过两种途径来使用它：

3.1 将扩展安装为python包

$ maturin develop

3.2 编译成动态文件从python加载

$ maturin develop --skip-install

--skip-install 命令会产生一个 pyd 文件而不是将其安装为python的包 - demo.cp312-win_amd64.pyd 文件在当前目录下，然后python可以直接导入使用。

另外还有一个指令 --release 会生成一个 xxxx.whl 文件，也就是Python pip安装的包源文件。

然后我们写一个python文件 src/main.py，下面是扩展的执行效果：

import time

import demo

s = time.time()
print(demo.integrate_f(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))

花费时间为：Elapsed: 12.638906717300415 s

可以看到python的执行时间是rust和rust扩展的5倍时长，如果我们将一些关键的性能通过rust重写，可以节省的时间成本是十分可观的。

整体的使用过程相当简洁，难点就是rust的学习曲线高，使用起来需要花费精力，但是还是可以慢慢尝试去使用它优化已有的项目性能。