使用tch-rs在Rust中加载和运行PyTorch模型

贾耀斐

于 2025-06-08 09:20:50 发布

阅读量280

点赞数 3

本文链接：https://blog.csdn.net/gitblog_01183/article/details/148508730

版权

使用tch-rs在Rust中加载和运行PyTorch模型

tch-rs Rust bindings for the C++ api of PyTorch. 项目地址: https://gitcode.com/gh_mirrors/tc/tch-rs

前言

在深度学习领域，PyTorch因其易用性和灵活性而广受欢迎。然而在生产环境中，我们往往需要将训练好的模型部署到高性能的Rust应用中。tch-rs项目为这一需求提供了完美的解决方案，它允许我们在Rust中直接加载和运行PyTorch模型。

Torch Script简介

Torch Script是PyTorch提供的一种模型序列化格式，它能够将Python中定义的模型转换为可移植的、高性能的表示形式。这种格式不仅包含模型架构，还包含了训练好的权重参数，使得模型可以脱离Python环境运行。

模型转换步骤

Python端模型导出

首先我们需要在Python环境中将训练好的模型导出为Torch Script格式。以ResNet-18为例：

import torch
import torchvision

# 加载预训练模型
model = torchvision.models.resnet18(pretrained=True)
model.eval()  # 设置为评估模式

# 创建示例输入
example = torch.rand(1, 3, 224, 224)

# 使用追踪方法生成Torch Script
traced_script_module = torch.jit.trace(model, example)

# 保存模型
traced_script_module.save("model.pt")

关键点说明：

model.eval()确保模型处于评估模式，这对批归一化等层的表现至关重要
追踪(tracing)方法会记录模型在给定输入上的操作序列
生成的model.pt文件包含了完整的模型定义和参数

Rust端模型加载与推理

环境准备

确保你的Rust项目中已经添加了tch-rs依赖。tch-rs提供了与PyTorch C++ API的绑定，使得在Rust中使用PyTorch模型成为可能。

核心代码解析

pub fn main() -> anyhow::Result<()> {
    // 解析命令行参数
    let args: Vec<_> = std::env::args().collect();
    let (model_file, image_file) = match args.as_slice() {
        [_, m, i] => (m.to_owned(), i.to_owned()),
        _ => bail!("usage: main model.pt image.jpg"),
    };
    
    // 加载并预处理图像
    let image = imagenet::load_image_and_resize(image_file)?;
    
    // 加载Torch Script模型
    let model = tch::CModule::load(model_file)?;
    
    // 执行推理
    let output = model.forward_ts(&[image.unsqueeze(0)])?.softmax(-1);
    
    // 输出Top-5预测结果
    for (probability, class) in imagenet::top(&output, 5).iter() {
        println!("{:50} {:5.2}%", class, 100.0 * probability)
    }
    
    Ok(())
}

关键组件详解

图像预处理：
- imagenet::load_image_and_resize函数负责加载图像并调整到224x224分辨率
- 同时会应用ImageNet的标准归一化处理
模型加载：
- tch::CModule::load方法从文件加载Torch Script模型
- 加载后的模型可以直接用于推理
推理执行：
- forward_ts方法执行前向传播
- softmax将输出转换为概率分布
- 也可以使用更简洁的apply方法：image.unsqueeze(0).apply(&model).softmax(-1)
结果解析：
- imagenet::top函数提取概率最高的5个类别及其置信度

实际应用示例

假设我们有一张老虎的图片，运行程序后会输出类似以下结果：

tiger, Panthera tigris                             96.33%
tiger cat                                           3.56%
zebra                                               0.09%
jaguar, panther, Panthera onca, Felis onca          0.01%
tabby, tabby cat                                    0.01%