RTen

RTen（Rust 张量引擎†）是用于从 ONNX 格式转换的机器学习模型的运行时，您可以从 PyTorch 和其他框架中导出这些模型。

该项目还提供各种领域常见的前处理和后处理任务的支持库。这使得 RTen 成为运行模型的更完整的 Rust 应用工具包。

† 该名称也是对 PyTorch 的 ATen 库的引用。

目标

• 提供一个小巧高效的神经网络运行时，使您能够轻松地将使用 PyTorch 等框架创建的模型在 Rust 应用程序中运行。
• 易于编译和运行在各种平台上，包括 WebAssembly
• 全 Rust。该项目及其所有依赖项均使用 Rust 编写。

局限性

该项目有一些局限性需要注意。计划在未来解决这些问题。

• 目前仅支持 CPU 推理。目前不支持在 GPU 或其他加速器上运行模型。
• 目前不支持所有 ONNX 运算符。有关目前支持的运算符，请参阅 src/schema.fbs 中的 OperatorType。对于已实现的运算符，某些属性或输入形状可能不受支持。
• 支持的数据类型有限：仅支持 float32 和 int32 张量。int64 和布尔张量将转换为 int32。
• RTen 的优化不如 ONNX Runtime 或 TensorFlow Lite 等成熟的运行时。性能差异取决于使用的运算符、模型结构、CPU 架构和平台。

入门指南

开始的最佳方式是克隆此仓库并尝试在本地运行一些示例。转换脚本使用流行的 Python 机器学习库，因此您需要安装 Python >= 3.10。

示例位于 rten-examples/ 目录中。请参阅 README，其中包含所有示例的描述以及运行它们的步骤。作为一个快速入门，以下是在运行图像分类示例时的步骤：

git clone https://github.com/robertknight/rten.git
cd rten

# Install model conversion tool
pip install -e rten-convert

# Install dependencies for Python scripts
pip install -r tools/requirements.txt

# Export an ONNX model. We're using resnet-50, a classic image classification model.
python -m tools.export-timm-model timm/resnet50.a1_in1k

# Convert model to this library's format
rten-convert resnet50.a1_in1k.onnx resnet50.rten

# Run image classification example. Replace `image.png` with your own image.
cargo run -p rten-examples --release --bin imagenet mobilenet resnet50.rten image.png

ONNX模型的转换

RTen不能直接加载ONNX模型。ONNX模型必须通过转换工具运行，该工具生成一个基于FlatBuffers格式的优化模型（.rten），该格式是引擎可以加载的。这从概念上类似于TensorFlow Lite和ONNX Runtime使用的.tflite和.ort格式。

转换工具需要Python >= 3.10。要转换现有的ONNX模型，请运行：

pip install rten-convert
rten-convert your-model.onnx

有关用法和版本兼容性的更多信息，请参阅 rten-convert README。

在JavaScript中的使用

要在JavaScript应用程序中使用此库，有两种方法

1. 在JavaScript中准备模型输入，并使用rten库内置的WebAssembly API运行模型，然后返回一个张量，该张量需要在JS中进行后处理。这种方法对于预处理简单的任务可能最容易。

   图像分类示例使用这种方法。

2. 创建一个使用rten并在Rust端进行输入预处理和输出后处理的Rust库，并公开一个特定领域的WebAssembly API。如果需要进行复杂和/或计算密集型的预/后处理，这种方法更为合适。

在运行示例之前，您需要遵循下面"Building the WebAssembly library"（构建WebAssembly库）下的步骤。

使用RTen内置的WebAssembly API在JavaScript项目中运行模型的通用步骤是：

1. 开发一个模型或找到一个你想要运行的预训练模型。可以从ONNX模型库或Hugging Face获取ONNX格式的预训练模型。
2. 如果模型尚未是ONNX格式，则将其转换为ONNX。PyTorch用户可以使用torch.onnx进行此操作。
3. 使用本存储库中的rten-convert包将模型转换为RTen使用的优化格式。请参阅上面关于转换模型的章节。
4. 在你的JavaScript代码中，获取WebAssembly二进制文件，并使用init函数初始化RTen。
5. 获取准备好的.rten模型，并使用它从这个库中实例化一个Model类。
6. 每次你想运行模型时，准备一个或多个包含模型预期格式的输入数据的Float32Array，并调用Model.run。这将返回一个TensorList，它提供了访问输出形状和数据的方式。

在构建库之后，Model和TensorList类的API文档可在dist/rten.d.ts中找到。

构建WebAssembly库

先决条件

为了构建RTen WebAssembly版本，你需要

• 最新的稳定版Rust
• make
• （可选）可以从Binaryen获取的wasm-opt工具用于优化.wasm二进制文件，以提高性能
• （可选）运行演示的最新版本Node

构建 rten

git clone https://github.com/robertknight/rten.git
cd rten
make wasm

由 make wasm 创建的构建需要支持 WebAssembly SIMD，自 Chrome 91、Firefox 89 和 Safari 16.4 以来可用。您可以使用 make wasm-nosimd 来构建没有 WebAssembly SIMD 支持的库，或者使用 make wasm-all 同时构建。非 SIMD 构建速度明显较慢。

在运行时，您可以通过调用此包导出的 binaryName() 函数来找出支持的构建。