如何在Rust中调用AI模型 Rust使用ONNX运行时推理教程-人工智能-PHP中文网

如何在Rust中调用AI模型 Rust使用ONNX运行时推理教程

P粉602998670

发布： 2025-06-28 09:38:02

原创

284人浏览过

在rust中调用onnx模型进行推理的步骤包括：安装onnx runtime并配置依赖、加载模型创建会话、准备输入数据执行推理。1. 安装onnx runtime动态库并在cargo.toml中添加ort和ndarray依赖；2. 使用ort库加载.onnx模型文件并创建session对象；3. 构造与模型输入匹配的张量数据，通过session.run()执行推理并解析输出结果。此外需要注意内存管理、数据类型和维度匹配以及性能优化等细节问题。

如何在Rust中调用AI模型 Rust使用ONNX运行时推理教程

在 Rust 中调用 AI 模型，尤其是通过 ONNX Runtime 进行推理，是一个越来越常见的需求。虽然 Rust 在机器学习生态上不如 Python 成熟，但借助 ONNX Runtime 提供的 C 接口和一些封装库，我们可以在 Rust 项目中轻松集成训练好的模型进行推理。下面我们就来看看如何实际操作。

准备工作：安装 ONNX Runtime 和相关依赖

要在 Rust 中使用 ONNX Runtime，首先需要确保系统中已经安装了 ONNX Runtime 的动态库（DLL 或 .so 文件）。你可以从 ONNX Runtime 官网下载预编译版本，或者自己编译。

接着，在你的 Cargo.toml 中添加如下依赖：

[dependencies]
ort = "0.12"
ndarray = "0.15"

登录后复制

其中 ort 是一个对 ONNX Runtime 的 Rust 封装，简单易用；ndarray 则用于处理多维数组，方便数据准备和输出解析。

加载模型并创建会话

ONNX Runtime 在 Rust 中的核心是创建一个 Session，然后用它来运行模型。

use ort::{Environment, Session, SessionOutputs};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let env = Environment::create()?;
    let session = Session::from_file(&env, "model.onnx")?;

    // 查看输入输出信息（可选）
    for (i, input) in session.inputs.iter().enumerate() {
        println!("Input {}: {:?}", i, input);
    }

    Ok(())
}

登录后复制

这段代码加载了一个本地的 .onnx 模型文件，并创建了一个会话。接下来就可以准备输入数据了。

注意：模型路径要正确，且模型格式必须为 ONNX 格式。

准备输入数据与执行推理

假设你有一个图像分类模型，输入是一个形状为 [1, 3, 224, 224] 的浮点张量（NCHW 格式），你可以这样构造输入：

use ndarray::{Array, IxDyn};

let input_data = Array::random((1, 3, 224, 224), || rand::random::<f32>());
let inputs = vec![input_data.into_dyn()];
let outputs: SessionOutputs = session.run(inputs)?;

登录后复制

执行完 run() 后，outputs 就包含了模型的推理结果。你可以遍历查看每个输出张量的内容：

for (i, output) in outputs.iter().enumerate() {
    println!("Output {}: {:?}", i, output);
}

登录后复制

如果你知道输出的名字，也可以按名字提取：

if let Some(output_tensor) = outputs.tensor("output_name")? {
    let data = output_tensor.try_extract_tensor::<f32>()?;
    // 处理数据...
}

登录后复制

常见问题与注意事项

内存管理：确保传入的数据生命周期足够长，尤其是在异步或多线程场景下。
数据类型匹配：输入的数据类型必须与模型期望的一致（通常是 f32 或 i64）。
维度匹配：输入张量的 shape 必须完全匹配模型定义。
性能优化：
- 可以启用 GPU 支持（如果 ONNX Runtime 编译时启用了 CUDA 或 DirectML）；
- 使用固定大小的输入批处理提高吞吐量；
- 避免频繁创建 Session，尽量复用。