Note: 本文作为入门教程，抛砖引玉，帮你初步了解如何使用rust为python写扩展模块，涉及从头到尾的详细步骤，基于此，剩下的只有深入rust，才能做得更好。

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

本文简单介绍使用rust为python计算性质的代码做一个优化，使用pyo3库为python写一个扩展供其调用，咱们下面开始，来看看具体的过程和效率的提升。

我的台式机环境:

设备名称 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: 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;
}

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

3. 通过pyo3写扩展

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

# (replace demo 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

整体项目结构如下: Cargo.toml中的一些字段含义：doc.rust-lang.org/cargo/refer…

.
├── Cargo.toml	// rust包管理文件，会在[lib]中声明目标扩展包的名称
├── src			// rust源文件目录, 你的扩展文件就写在这里，这个目录是maturin初始化的时候自动创建
│   └── lib.rs	// 扩展文件
├── pyproject.toml	// python包管理文件，里面有python的包名字定义
├── .gitignore
├── Cargo.lock
└── demo	// 我们的目标模块名称，需手动创建
    ├── main.py	// 用来测试的文件
    └── demo.cp312-win_amd64.pyd	// 编译生成的动态链接库文件，供import给python使用

在src/lib.rs 下写入:

use pyo3::prelude::*;

// Caculate the integrate. This function will be exposed to python.
#[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 将扩展安装为python包

$ maturin develop

这个命令会将rust代码转为python的包，并安装在当前python环境内。通过 pip list 就能看到。

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

$ maturin develop --skip-install

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

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

首先我们在rust项目下，与 Cargo.toml 同级目录下，创建一个 demo 目录，然后我们写一个python文件 demo/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: 10.908721685409546 s

可以看到python的执行时间是rust和rust扩展的3倍时长，单进程看着好像不太大是吧，下面还有并行版本。

4. 并行加速

4.1 python多进程效果

Python多进程很神奇，你写的不好的话，他比单进程下还要慢。

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

花费时间: Elapsed: 18.86696743965149 s，比单进程下时间少了不到一半。

4.2 rust多线程加速给python使用

如果我们使用rust的并行库，将rust进一步加速，速度效果更明显: 将上面的 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(())
}

执行上一个标题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))

花费时间为：Elapsed: 0.9684994220733643 s。这比原先的单线程rust版本又快了10倍。但是差不多是python并行版本的18倍左右，是python单进程版本的32倍左右。如果我们将一些关键的性能通过rust重写，可以节省的时间成本是十分可观的。

总体来看，整体的使用过程相当简洁方便，难点就是rust的学习曲线高，使用起来需要花费精力，但是还是可以慢慢尝试去使用它优化已有的项目性能，哪怕只是一个功能函数，熟能生巧，一切慢慢来。

以上数据比较仅供参考，不同机器可能差异也不同。

如果你对Python感兴趣，想要学习python，这里给大家分享一份Python全套学习资料，都是我自己学习时整理的，希望可以帮到你，一起加油！

😝有需要的小伙伴，可以V扫描下方二维码免费领取🆓

​

1️⃣零基础入门

① 学习路线

对于从来没有接触过Python的同学，我们帮你准备了详细的学习成长路线图。可以说是最科学最系统的学习路线，你可以按照上面的知识点去找对应的学习资源，保证自己学得较为全面。

② 路线对应学习视频

还有很多适合0基础入门的学习视频，有了这些视频，轻轻松松上手Python~

③练习题

每节视频课后，都有对应的练习题哦，可以检验学习成果哈哈！

2️⃣国内外Python书籍、文档

① 文档和书籍资料

3️⃣Python工具包+项目源码合集

①Python工具包

学习Python常用的开发软件都在这里了！每个都有详细的安装教程，保证你可以安装成功哦！

②Python实战案例

光学理论是没用的，要学会跟着一起敲代码，动手实操，才能将自己的所学运用到实际当中去，这时候可以搞点实战案例来学习。100+实战案例源码等你来拿！

③Python小游戏源码

如果觉得上面的实战案例有点枯燥，可以试试自己用Python编写小游戏，让你的学习过程中增添一点趣味！

4️⃣Python面试题

我们学会了Python之后，有了技能就可以出去找工作啦！下面这些面试题是都来自阿里、腾讯、字节等一线互联网大厂，并且有阿里大佬给出了权威的解答，刷完这一套面试资料相信大家都能找到满意的工作。

上述所有资料 ⚡️ ，朋友们如果有需要的，可以扫描下方👇👇👇二维码免费领取🆓
​