用 Rust 跑推理：从 ONNX Runtime 到 Candle 的框架选型与实战

用 Rust 跑推理：从 ONNX Runtime 到 Candle 的框架选型与实战

一、Python 之外的推理之路——Rust AI 推理的工程痛点

大部分 AI 推理都跑在 Python 生态里，PyTorch 和 TensorFlow 已经把门槛压得很低。但当你要把模型部署到边缘设备、嵌入到 CLI 工具、或者追求毫秒级冷启动时，Python 的 GIL、依赖膨胀和启动延迟就成了硬伤。

我在尝试用 Rust 构建一个本地 AI 工具时，踩到的第一个坑就是：Rust 的 AI 推理生态远没有 Python 成熟。没有一行 pip install 就能搞定的方案，每个框架都有自己的坑。ONNX Runtime 的 Rust binding 文档稀少，Candle 还在快速迭代中 API 频繁变动，tract 的模型支持有限。选错框架，后续就是无尽的适配工作。

这篇文章梳理当前 Rust AI 推理框架的选型思路，从底层机制到实际代码，帮助同样在 Rust 上做 AI 推理的开发者少走弯路。

二、Rust AI 推理框架的架构分层与运行机制

先看整体架构，理解不同框架在推理链路中的位置：

graph LR
    A[训练框架<br>PyTorch/TensorFlow] -->|导出| B[模型格式<br>ONNX/Safetensors/GGML]
    B -->|加载| C{Rust 推理框架}
    C --> D[ONNX Runtime Rust<br>调用 C++ 库]
    C --> E[Candle<br>纯 Rust 实现]
    C --> F[tract<br>纯 Rust 实现]
    C --> G[candle-ggml/gguf<br>GGUF 格式推理]

    D --> H[CPU / GPU / NPU]
    E --> H
    F --> H
    G --> H

    style D fill:#f9f,stroke:#333
    style E fill:#bbf,stroke:#333
    style F fill:#bfb,stroke:#333
    style G fill:#ffb,stroke:#333

2.1 ONNX Runtime Rust：工业级但绑定层薄

ONNX Runtime 是微软维护的跨平台推理引擎，Rust 绑定通过 FFI 调用 C++ 核心。它的优势是模型支持广、后端多（CPU/CUDA/TensorRT/CoreML），但 Rust binding 的 API 覆盖率远不如 C++/Python。

底层机制：Rust 侧的 ort crate 封装了 C API，每次推理调用都经过 Rust -> C FFI -> C++ 核心 -> 硬件后端的链路。FFI 边界有开销，但相比推理计算本身可以忽略。

2.2 Candle：HuggingFace 的纯 Rust 推理

Candle 是 HuggingFace 用纯 Rust 写的推理框架，核心目标是零 C 依赖、小二进制体积、支持 GPU（通过 CUDA）。它的设计哲学是"最小化抽象"，直接操作张量，没有 ONNX 那样的计算图优化层。

底层机制：Candle 使用自定义的张量存储后端，CPU 上用 BLAS 加速，GPU 上通过 cudarc crate 直接调用 CUDA API。没有计算图编译步骤，推理就是逐算子执行。

2.3 tract：纯 Rust 的多格式推理

tract 支持 ONNX 和 TensorFlow 模型格式，纯 Rust 实现，目标是可嵌入、可交叉编译。它的优化器可以对计算图做常量折叠、算子融合等变换。

2.4 GGUF 推理：本地大模型的事实标准

GGUF 是 llama.cpp 使用的模型格式，Rust 侧通过 candle-gguf 或 llama-cpp-rs 加载。这是目前 Rust 跑本地 LLM 最成熟的路径。

三、框架选型与生产级代码实践

3.1 ONNX Runtime：适合已有 ONNX 模型的生产部署

use ort::{Environment, SessionBuilder, Value};
use ndarray::Array2;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 初始化 ONNX Runtime 环境
    // 使用默认执行提供者，生产环境可切换到 CUDA/TensorRT
    let environment = Environment::builder().build()?;
    let session = SessionBuilder::new(&environment)?
        .with_model_from_file("model.onnx")?;

    // 构造输入张量——这里以 1x784 的 MNIST 输入为例
    let input = Array2::<f32>::from_shape_vec((1, 784), vec![0.0; 784])?;
    let input_value = Value::from_array(session.allocator(), &input)?;

    // 执行推理
    let outputs = session.run(vec![input_value])?;

    // 解析输出
    let output: Array2<f32> = outputs[0].try_into_array()?;
    println!("推理结果形状: {:?}", output.shape());

    Ok(())
}

3.2 Candle：适合需要纯 Rust、小体积的场景

use candle_core::{Device, Tensor, DType};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    // 构造输入张量
    let input = Tensor::randn(0f32, 1f32, (1, 784), &device)?;

    // 加载 Safetensors 模型权重（简化示例）
    // 实际使用时需要根据模型架构手动构建计算图
    let var_store = candle_nn::VarBuilder::from_pth("model.pt", DType::F32, &device)?;

    // 构建模型——以简单的全连接网络为例
    let model = candle_nn::Linear::new(
        var_store.pp("layer1").get((128, 784), "weight")?,
        var_store.pp("layer1").get(128, "bias")?,
    );

    // 前向推理
    let output = model.forward(&input)?;
    println!("输出形状: {:?}", output.shape());

    Ok(())
}

3.3 GGUF 本地 LLM 推理：最实用的路径

use candle_core::Device;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    // 加载 GGUF 格式的量化模型
    // 量化模型体积小，适合本地部署
    let model_path = "qwen2-1.5b-instruct-q4_k_m.gguf";

    // 使用 candle-gguf 加载
    let model = candle_transformers::models::quantized_llama::ModelWeights::from_gguf(
        model_path,
        &mut candle_transformers::models::quantized_llama::Config::default(),
        &device,
    )?;

    // 简化的推理循环——实际需要 tokenizer 和采样器配合
    println!("模型加载成功，参数量级: 1.5B");

    Ok(())
}

3.4 框架选型决策表

四、Rust AI 推理的现实代价与边界

4.1 生态成熟度是最大瓶颈

Python 的 AI 生态经过多年积累，几乎任何模型都有现成的推理脚本。Rust 生态目前只能覆盖主流模型架构，遇到非标准算子或自定义层，需要自己实现。Candle 的算子库还在扩展中，tract 对动态形状的支持有限。

4.2 调试体验差距明显

Python 推理出错时，有丰富的可视化工具和社区问答。Rust 推理出错时，错误信息往往停留在张量形状不匹配的层面，定位问题需要深入到算子实现。我在用 Candle 跑一个 BERT 模型时，shape mismatch 的报错花了两个小时才定位到是 tokenizer 的 padding 策略不对。

4.3 性能并非总是优于 Python

纯 Rust 实现的推理框架在 CPU 上不一定比 Python + ONNX Runtime 快。因为 ONNX Runtime 的 C++ 核心有大量手工优化的算子实现（AVX/NEON 指令集），而 Rust 侧的算子优化程度参差不齐。Candle 的优势更多在于部署体积小、无运行时依赖，而非绝对推理速度。

4.4 量化与精度损失

GGUF 格式的量化模型（Q4_K_M、Q5_K_S 等）在压缩模型体积的同时会引入精度损失。对于分类、检测等任务，Q4 量化的精度下降通常可接受；但对于生成式任务，量化可能导致输出质量明显下降。选型时必须在体积、速度和精度之间做权衡。

五、总结

Rust AI 推理框架的选型没有银弹。如果已有 ONNX 模型且追求稳定性，ONNX Runtime 是最稳妥的选择；如果需要纯 Rust、小体积部署，Candle 是当前最活跃的项目；如果要跑本地 LLM，GGUF 格式配合 candle-gguf 或 llama-cpp-rs 是最成熟的路径。

落地路线建议：先用 Python 验证模型和推理流程，确认模型格式和输入输出规范后，再迁移到 Rust。不要一开始就在 Rust 里调模型——调试效率的差距会让你怀疑人生。Rust 的价值在于部署阶段，而非实验阶段。