Lib.rs

› 科学 › 机器学习
#onnx #model #inference #di

bin+lib stonnx

一个用于在 ONNX 模型上运行推理的 Rust 库

由 Atari2 拥有。

4 个版本

0.1.3 2023 年 12 月 21 日
0.1.2 2023 年 12 月 20 日
0.1.1 2023 年 12 月 20 日
0.1.0 2023 年 12 月 20 日

#249 in 机器学习

每月 24 次下载

MIT 许可证

315KB
7.5K SLoC

Stonnx

主要贡献者

名称来源于 Steelix（金属宝可梦，选择它是因为该项目是用 Rust 编写的）和 ONNX 的结合。此外，它也是一种双关语，与 stonks 有关。

项目描述

该项目涉及使用 Rust 语言实现一个 ONNX 解释器。需要满足以下要求：

  • 创建一个解析器，从 ONNX 文件中提取创建网络所需的信息；
  • 实现 ONNX 的子集操作符；
  • 使用并行执行网络；
  • （可选）与其他语言的绑定。

使用方法

  • 使用位于 scripts 文件夹中的脚本 download_models.ps1（Windows）或 download_models.sh（Linux/macOS）下载 ONNX 模型，或者从 这里 手动下载并解压到存储库根目录。最初，模型是通过 git-lfs 添加到存储库的，但这是因为我们从 politecnico 收到的存储库空间和带宽配额太低，所以出现了问题；

  • 执行命令 cargo build --release 编译项目；

    • 默认情况下，项目使用 rayon 进行并行编译（要使用我们自己实现的并行编译，请使用命令 cargo build --features custom-threadpool --release）；

  • 执行命令 cargo run --release -- --model <modelname> 或者 cargo run --release --features custom-threadpool -- --model <modelname>

    • modelname 是包含模型的文件夹路径。如果路径是相对路径，它将被假定为相对于 $pwd/models 的相对路径，其中 $pwd 代表可执行文件所在的文件夹。在包含模型的文件夹中，应该有一个 inputs.json 文件，其结构如下

      {
    "inputs": [
        "<percorso verso input del modello>",
        ...
    ],
    "outputs": [
        "<percorso verso output attesi del modello>",
        ...
    ],
    "modelpath": "<percorso al file .onnx contenente il modello>"
}

      此文件中的所有路径都可以是相对路径或绝对路径，如果路径是相对路径，则相对于 $pwd/models/$modelname

    • 在克隆仓库后，将会有一些不在 models 文件夹中的模型，这些模型将在第一次构建项目时自动从互联网上下载。

  • 如果想要查看操作器的执行情况，可以在之前的命令中添加 --verbose 选项（默认情况下详细程度为0）。

    • 示例： cargo run --release -- --model <modelname> --verbose 1
    • verbose = 0：不显示任何内容
    • verbose = 1：显示有关操作执行的信息
    • verbose = 2：将每个操作器的输出写入一个 .npy 文件
    • verbose = 4：将每个操作器的中间输出也写入一个 .npy 文件

  • 执行 googlenet 模型的示例如下： cargo run --release --features custom-threadpool -- --verbose 0 --model googlenet-12

  • 此外，还可以添加 --gengraph 命令以生成可用于绘制模型图的文件。程序可以生成一个专有的 json 格式（可以被这个 工具 读取）或者通用的 dot 格式（用于 graphviz）。可以通过 --graphtype 选项来控制图形的格式，可以是 json 或 dot（默认： json）。生成的文件将放置在执行模型所在的同一文件夹中。

  • 执行以下命令 cargo doc --open 以查看项目的文档。

支持的模型

测试的模型参考了archive部分中的ONNX官方仓库，因为截至2023年12月初，已经更新并添加了新的模型，但这个程序的开发开始于模型库更新之前。以下是被测试的模型：

注意：在选择模型时，选择了ONNX仓库archive部分中的最新版本。

使用的重要Crate

使用了以下Crate

  • ndarray：用于管理多维数组；
  • anyhow：用于处理错误；
  • clap：用于处理命令行参数；
  • bytemuck：用于类型转换；
  • petgraph：用于创建网络图；
  • serde：用于数据结构的序列化和反序列化；
  • protobuf：用于处理protobuf文件；

项目各个部分的描述

  • src/main.rs：项目的主要文件，包含main函数和命令行参数处理；

  • src/operators：包含ONNX运算符的实现文件；

    • 其中实现的运算符包括（但不仅限于）：Add、AveragePool、BatchNormalization、Concat、Conv、Dropout、Flatten、Gemm、GlobalAveragePool、MaxPool、MatMul、Mul、Relu、Reshape、Softmax、Sum、Transpose；

  • src/onnxparser：包含ONNX解析器实现文件；

    • 该目录中的文件在构建时由protobuf编译器生成（请参阅build.rs），使用protobuf_codegen库。

  • src/executor：包含网络执行的实现，包括

    • 创建图和执行运算符的逻辑；
    • 创建线程池（自定义或使用rayon库）和管理线程间通信的逻辑；
    • 比较预期输出和实际输出的逻辑；

  • src/parallel：包含线程池并行实现的文件；

    • 并行实现有两种方式
      • 使用rayon库（默认）；
      • 使用自定义线程池（可通过--features custom-threadpool标志激活）；

  • src/protograph：包含用于创建包含网络图的 .json 文件的实现文件；

  • src/protos：包含用于创建包含 protobuf 文件结构的 .rs 文件的 onnx.proto 文件；

  • src/common/mod.rs：包含用于处理 ONNX 文件的文件结构；

    • 处理 verbose；
    • 处理各种模型的输入和输出文件的路径；
    • 处理包含网络图的 .json 文件；
    • 处理操作符的 opset_version；
    • 处理执行操作符得到的结果；
    • 处理数据类型；
    • 处理用于表示张量的特征（ArrayElement）；
    • 处理读取二进制格式数据以创建张量的操作；

  • src/utils：处理用于创建和管理张量的有用操作；

    • 如果 verbose 设置为大于 0 的值，则在执行网络时创建包含张量的 .npy 文件；
    • 也有函数可以创建根据形状和数据类型或形状、字节和数据类型创建的张量；
    • 也有函数可以根据模型描述创建输入和输出张量；

程序架构

程序有四个主要步骤

  • 解析和读取 .onnx 文件，其外部输入，以及使用这些输入初始化张量；
  • 创建两个由 HashMap 表示的图，一个将每个操作符与其输入连接起来，另一个将每个输入（张量）与其使用的操作符连接起来，实际上是第一个的逆。因此，我们可以拥有每个操作符的依赖关系图，这些图将在执行模型时使用。实际上，当操作符的所有依赖项都得到满足（例如，其输入由前面的操作符生成）时，它可以放入线程池的队列中；
  • 执行推理：模型并行执行，从没有依赖项或依赖项已完全满足的操作符开始，这些操作符将放入线程池队列，每当一个操作符完成时，其结果将通知主线程，主线程将更新依赖关系图，并启动现在因结果而满足其依赖项的操作符，这样，直到依赖关系图为空，我们就会得到最终结果；
  • 线程池实现：表示线程池的主要结构由一个 queue（用于任务管理的同步队列）、workers（线程集合）和一个 queuestate（跟踪队列中操作数数量的计数器）组成；
    • 当创建线程时，它开始循环等待执行的操作，当操作因其依赖项得到解决而变得可用时，添加该操作的线程通过一个 Condvar 通知工作者队列，因此，一个空闲的线程获取操作并执行它，当操作完成时，线程通知主线程任务已完成，主线程更新依赖关系图，并将现在可以执行的操作添加到队列中，以此类推，直到没有更多任务要执行；
    • 代码使用 Mutex 和 Condvar 来同步访问队列以及线程之间的通信；
    • 关于操作符内部的并行处理，这仅在少数操作符中实现（在认为更方便的地方，以避免在处理非常简单且快速的计算时加重程序执行负担，以及处理各种锁和线程之间的上下文切换），特别是，在这些操作符中可以找到 Conv、Exp、Sqrt、Pow、MaxPool、AveragePool 等，但不止这些。特别是，并行处理被引入到执行非常重的循环的地方，例如在操作符内部执行卷积计算，而不是使用我们自己编写的 ThreadPool，主要利用了 rayon 的并行迭代器功能；
  • 以下屏幕截图显示了并行线程的示例用法（使用 VTune 获取），以及在该时间段内执行的部分图形（使用 Netron 获取）

Googlenet

Inception

  • 输出比较：程序还会读取“参考”输出，这些输出应该是通过执行模型获得的，并将其与实际获得的输出进行比较，检查两个结果的单个值之间的差异是否最大为 10e-4，并打印一些比较统计信息。

将 Stonnx 作为库使用

程序被编译为动态库 (.dll / .so / .dylib)，名称为 stonnx_api，可以通过公开的绑定正常使用。

以下语言提供了绑定

  • Python
  • C
  • C++
  • C#

目前，绑定非常有限，通过提供一些用于创建网络和执行网络的函数来实现。

基准测试

见 BENCHMARKS

依赖关系

~15–28MB
~411K SLoC