Rust AI 生态扫描：从 tract 到 candle，本地推理库该怎么选？

Rust AI 生态扫描：从 tract 到 candle，本地推理库该怎么选？

一、当推理必须跑在本地：Rust AI 库的现实需求

AI 推理不一定非得依赖云端 API。在边缘设备、离线环境、数据隐私敏感的场景下，把模型跑在本地是刚需。Python 生态有 PyTorch、ONNX Runtime，但部署时 Python 运行时的体积和启动开销是个问题。Rust 在这个方向上的优势很明确：编译为单一二进制、无运行时依赖、内存安全、启动快。

但 Rust 的 AI 生态还远没有 Python 成熟。选择一个推理库，不只是看它支持哪些模型格式，还要看它的推理性能、跨平台能力、API 稳定性，以及社区活跃度。选错了，可能模型加载不了，可能在 ARM 板子上跑不起来，可能 API 在下个版本全变了。

实际项目中，我遇到过这样的场景：需要在树莓派上跑一个图像分类模型，要求冷启动时间低于 500ms，内存占用不超过 200MB。Python 方案直接出局，Rust 方案需要在 tract 和 candle 之间做选择。这个选择不是看哪个 star 多，而是看哪个在目标平台上真的能跑。

二、三大 Rust AI 推理库的架构差异

目前 Rust 生态中，本地推理库主要有三个选择：tract、candle 和 burn。它们的架构设计思路差异很大。

graph LR
    subgraph tract
        A1[ONNX/TensorFlow 前端] --> B1[计算图优化]
        B1 --> C1[多后端推理: NumPy/CPU/GPU]
    end
    subgraph candle
        A2[纯 Rust 算子实现] --> B2[自动微分]
        B2 --> C2[CUDA/Metal/CPU 后端]
    end
    subgraph burn
        A3[Burn IR 中间表示] --> B3[后端抽象层]
        B3 --> C3[WGPU/CUDA/NDArray 后端]
    end

tract 的核心思路是"加载即推理"。它直接解析 ONNX 和 TensorFlow 的模型文件，不需要你用 Rust 重写模型逻辑。它的计算图优化器会在推理前做算子融合、常量折叠等优化。适合"我有一个现成的模型，想用 Rust 跑起来"的场景。

candle 来自 HuggingFace，设计目标是提供纯 Rust 的张量运算和自动微分。它不是模型加载器，而是一套张量计算框架。你需要用 candle 的 API 重新实现模型的推理逻辑，或者使用它提供的预实现模型（如 LLaMA、Whisper 等）。优势是和 HuggingFace 生态天然对接。

burn 的野心更大，它试图做一套统一的前端-后端抽象。前端定义模型，后端可以切换到 WGPU（跨平台 GPU）、CUDA 或纯 CPU。它的设计更接近一个深度学习框架，而不只是推理引擎。

三、用 tract 加载 ONNX 模型：一个完整的推理流程

以下代码展示了用 tract 加载一个 ONNX 图像分类模型并执行推理的完整流程：

use tract_onnx::prelude::*;
use anyhow::{Context, Result};

/// 图像分类推理器
struct ImageClassifier {
    /// tract 优化后的推理计划
    plan: SimplePlan<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
    /// 模型输入尺寸
    input_size: (usize, usize),
}

impl ImageClassifier {
    /// 从 ONNX 文件创建推理器
    fn from_onnx(path: &str, input_size: (usize, usize)) -> Result<Self> {
        // 加载 ONNX 模型并优化计算图
        let model = tract_onnx::onnx()
            // 设置输入形状，帮助优化器做静态推断
            .with_input_fact(
                0,
                f32::fact([1, 3, input_size.0, input_size.1]).into(),
            )
            .into_optimized()?
            .into_runnable()?;

        Ok(Self { plan: model, input_size })
    }

    /// 执行推理，返回 Top-K 分类结果
    fn predict(&self, image_data: &[f32]) -> Result<Vec<(usize, f32)>> {
        // 构造输入张量，形状为 [1, 3, H, W]
        let input: Tensor = tract_ndarray::Array4::from_shape_vec(
            (1, 3, self.input_size.0, self.input_size.1),
            image_data.to_vec(),
        )?.into();

        // 执行推理
        let result = self.plan.run(tvec!(input.into()))?;
        let output = result[0].to_array_view::<f32>()?;

        // 提取 Top-5 分类结果
        let mut scores: Vec<(usize, f32)> = output
            .iter()
            .enumerate()
            .map(|(i, &s)| (i, s))
            .collect();
        scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        scores.truncate(5);

        Ok(scores)
    }
}

fn main() -> Result<()> {
    let classifier = ImageClassifier::from_onnx(
        "models/resnet50.onnx",
        (224, 224),
    ).context("模型加载失败，请检查 ONNX 文件路径")?;

    // 模拟输入数据（实际应从图像预处理获取）
    let fake_input = vec![0.0f32; 3 * 224 * 224];
    let top5 = classifier.predict(&fake_input)?;

    for (class_id, score) in top5 {
        println!("类别 {}，置信度：{:.4}", class_id, score);
    }

    Ok(())
}

这段代码的关键点在于 with_input_fact 的调用。tract 是静态图推理引擎，它需要在加载时就知道输入的形状，这样才能做算子融合和内存规划。如果你的模型有动态形状输入，tract 也能处理，但需要用 Symbol 来标记动态维度，配置会更复杂。

四、三个库的适用边界与选型取舍

tract 的优势与局限：tract 最大的优势是"开箱即用"——直接加载 ONNX 文件就能推理，不需要重写模型逻辑。但它的局限也很明显：只支持推理，不支持训练；对 ONNX 算子的覆盖不完整，遇到不支持的算子会报错；GPU 后端支持有限，目前主要是 CUDA，Metal 支持还在早期。

candle 的优势与局限：candle 的优势是纯 Rust 实现，和 HuggingFace 生态深度绑定，预置了多种主流模型的推理实现。但你需要用 candle 的 API 重新组织推理逻辑，学习成本比 tract 高。candle 的 CUDA 后端相对成熟，Metal 后端也在推进中。

burn 的优势与局限：burn 的后端抽象设计最灵活，WGPU 后端意味着它可以在任何支持 WebGPU 的平台上跑（包括浏览器）。但它的 API 稳定性目前还不够，文档也不完善，适合愿意跟进上游变化的早期采用者。

选型建议：如果你的需求是"加载 ONNX 模型，在 CPU 上跑推理"，tract 是最务实的选择。如果你需要 GPU 推理且模型是 HuggingFace 格式，candle 更合适。如果你需要跨平台 GPU 推理（尤其是 WebGPU），可以关注 burn 的进展。

一个常被忽略的考量是模型格式的兼容性。ONNX 生态最成熟，但不是所有模型都能顺利导出为 ONNX。Safetensors 是 HuggingFace 推荐的格式，candle 原生支持，tract 不支持。如果你的模型来源是 HuggingFace Hub，candle 的对接成本最低。

五、总结

Rust AI 推理库的选择取决于三个维度：模型格式、目标平台、是否需要 GPU。tract 适合 ONNX 模型的 CPU 推理，开箱即用但算子覆盖有限。candle 适合 HuggingFace 生态，GPU 支持较好但需要用其 API 重写推理逻辑。burn 的后端抽象最灵活，但 API 稳定性待观察。在边缘部署场景下，Rust 的编译产物体积和启动速度优势明显，但生态成熟度仍需时间。选型时建议先用 tract 验证可行性，遇到性能瓶颈再考虑 candle 或 burn。