一个简单例子——Run AI In Rust系列01

原文链接

一个简单例子——Run AI In Rust系列01

一、导言

开一个系列，讲述如何在rust中运行AI模型。计划主题如下：

1、如何在rust中运行pytorch的模型
2、如何使用candle运行大语言模型LLM
3、如何使用axum为模型写一个API服务
4、如何在candle使用gguf格式模型来实现低资源部署
5、如何在rust中运行sdxl

先立flag，再来实现。

二、极简入门

在python模型开发部分，我们用jit导出一个jit的模型。

import torch

# 使用一个简单的测试模型
class Test_Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(in_features=5, out_features=2)
        ...
    def forward(self,inputs:torch.Tensor)->torch.Tensor:
        out = self.linear(inputs)
        outputs = torch.nn.functional.softmax(out,dim=-1)
        return outputs
    ...

test_model = Test_Model()
# 参看一下输出,便于之后验证rust模型推断的正确行
test_model.eval()
inputs = torch.tensor([1,2,3,4,5],dtype=torch.float32)
with torch.no_grad():
    pt_outputs = test_model(inputs)
print(pt_outputs)

#导出jit格式模型,这样我们能脱离pytorch框架来运行模型
trace_model = torch.jit.trace(test_model, inputs)
trace_model.save("./data/model.jit")

rust部分

去官网下载对应环境的libtorch解压，然后要注意添加如下环境变量

#  added by torch
export LIBTORCH=Yourpath/libtorch   
export DYLD_LIBRARY_PATH=Yourpath/libtorch/lib

并在运行时设定环境变量: DYLD_LIBRARY_PATH=Yourpath/libtorch/lib，然后新建一个rust项目，在Cargo.toml中添加libtorch的依赖项

[dependencies]
# libtorch
tch = "0.13.0"

main.rs文件写入如下代码：

use tch::jit;
use tch::Tensor;

fn main() {
// 下面来个模型导出并在rust中使用的例子
// 载入模型
let model = jit::CModule::load("./data/model.jit").unwrap();
let inputs = Tensor::from_slice(&[1f32, 2f32, 3f32, 4f32, 5f32]);
let outputs = model.forward_ts(&[inputs]).unwrap();
// 将结果转化为rust向量格式
let a = Vec::<f32>::try_from(outputs).unwrap();
println!("{a:?}");

}

是不是超级简单。当然如果你没rust开发经验，你会看到一些觉得奇怪的语法，诸如::、<>、.、&。其实容易猜出来是什么意思。

::是路径
<>里面是类型
&表示引用
.表示方法

和python不一样路径和方法都用的点.rust里面做了明确区分，然后其他写法和python没有什么区别。