Rust深度学习推理框架tract-hir的使用

tract-hir是一个高效的Rust深度学习推理框架，专注于解析和执行神经网络模型的HIR(High-Level Intermediate Representation)中间表示。它提供了对神经网络模型的高效处理能力。

安装

在您的项目目录中运行以下Cargo命令：

cargo add tract-hir

或者在您的Cargo.toml中添加以下行：

tract-hir = "0.21.13"

基本使用示例

以下是一个使用tract-hir加载和运行ONNX模型的完整示例：

use tract_hir::prelude::*;

async fn main() -> TractResult<()> {
    // 1. 加载ONNX模型
    let model = tract_hir::onnx()
        .model_for_path("path/to/model.onnx")?
        .into_optimized()?
        .into_runnable()?;
    
    // 2. 准备输入数据
    let input_shape = model.input_shape(0)?;
    let input: Tensor = tract_hir::tensor::Tensor::zero::<f32>(&input_shape)?;
    
    // 3. 运行模型
    let outputs = model.run(tvec![input.into()])?;
    
    // 4. 处理输出
    for output in outputs {
        println!("Output tensor: {:?}", output);
    }
    
    Ok(())
}

更完整的示例

下面是一个更完整的示例，展示了如何从零构建一个简单的神经网络模型：

use tract_hir::prelude::*;

fn main() -> TractResult<()> {
    // 1. 创建一个简单的计算图
    let mut model = InferenceModel::default();
    
    // 2. 添加输入节点
    let input = model.add_source("input", f32::fact(dims!(1, 3, 224, 224)))?;
    
    // 3. 添加卷积层
    let conv = model.wire_node(
        "conv1",
        tract_hir::ops::cnn::Conv::default()
            .hwio()
            .strides(vec![1, 1])
            .padding(tract_hir::ops::cnn::PaddingSpec::Valid),
        &[input],
    )?;
    
    // 4. 添加激活函数(ReLU)
    let relu = model.wire_node("relu1", tract_hir::ops::nn::Relu, &conv)?;
    
    // 5. 添加池化层
    let pool = model.wire_node(
        "pool1",
        tract_hir::ops::cnn::Pool::default()
            .pool_type(tract_hir::ops::cnn::PoolType::Max)
            .kernel_shape(vec![2, 2])
            .strides(vec![2, 2]),
        &[relu[0]],
    )?;
    
    // 6. 添加全连接层
    let matmul = model.wire_node(
        "fc1",
        tract_hir::ops::math::MatMul::default(),
        &[pool[0]],
    )?;
    
    // 7. 添加输出节点
    model.set_output_outlets(&matmul)?;
    
    // 8. 优化模型
    let model = model.into_optimized()?.into_runnable()?;
    
    // 9. 准备输入数据
    let input = Tensor::zero::<f32>(&[1, 3, 224, 224])?;
    
    // 10. 运行模型
    let outputs = model.run(tvec![input.into()])?;
    
    println!("Model output: {:?}", outputs);
    
    Ok(())
}

完整示例代码

以下是一个使用tract-hir进行图像分类的完整示例：

use tract_hir::prelude::*;

async fn classify_image() -> TractResult<()> {
    // 1. 加载预训练的ONNX模型
    let model = tract_hir::onnx()
        .model_for_path("mobilenetv2.onnx")?  // 替换为你的模型路径
        .with_input_fact(0, f32::fact(dims!(1, 3, 224, 224)).into())?
        .into_optimized()?
        .into_runnable()?;

    // 2. 加载并预处理图像
    let image = image::open("test.jpg")?  // 使用image crate加载图像
        .resize_to_fill(224, 224, image::imageops::FilterType::Triangle)
        .to_rgb8();
    
    // 3. 转换图像为模型输入格式 (CHW, 归一化等)
    let tensor = tract_ndarray::Array4::from_shape_fn((1, 3, 224, 224), |(_, c, y, x)| {
        let mean = [0.485, 0.456, 0.406][c];
        let std = [0.229, 0.224, 0.225][c];
        (image[(x as _, y as _)][c] as f32 / 255.0 - mean) / std
    }).into();

    // 4. 运行推理
    let outputs = model.run(tvec![tensor])?;

    // 5. 处理输出结果
    let best = outputs[0]
        .to_array_view::<f32>()?
        .iter()
        .enumerate()
        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap());
    
    println!("预测结果: 类别={}, 置信度={}", best.unwrap().0, best.unwrap().1);
    
    Ok(())
}

关键特性

1. 高效的HIR中间表示：tract-hir使用高级中间表示来优化和执行神经网络模型
2. 多框架支持：可以导入ONNX、TensorFlow等格式的模型
3. 硬件加速：支持多种后端，包括CPU和GPU
4. 灵活的模型构建：可以从零开始构建自定义神经网络

tract-hir是MIT或Apache-2.0许可下的开源项目，由Mathieu Poumeyrol维护。