Lib.rs

Wonnx 是一个100%用 Rust 编写的 GPU 加速 ONNX 推理运行时，适用于 Web。

支持的平台（由 wgpu 启用）

✅ = 优先级支持 — 🆗 = 尽力支持 — 🚧 = 不支持，但支持中

入门

从命令行

确保您的系统支持 Vulkan、Metal 或 DX12 以访问 GPU。然后，下载二进制发布版或安装 Rust 并运行 cargo install --git https://github.com/webonnx/wonnx.git wonnx-cli 来安装 CLI。

CLI 工具 (nnx) 提供了一个方便的界面来调整模型（有关更多信息，请参阅README）

nnx info ./data/models/opt-squeeze.onnx
nnx infer ./data/models/opt-squeeze.onnx -i data=./data/images/pelican.jpeg --labels ./data/models/squeeze-labels.txt --top 3

从 Rust

将 wonnx 包作为依赖项添加（如果您有 cargo-add，请运行 cargo add wonnx）。然后，查看示例以获取用法示例，或浏览 API 文档。

从 Python

pip install wonnx

然后，使用

from wonnx import Session
session = Session.from_path(
    "../data/models/single_relu.onnx"
)
inputs = {"x": [-1.0, 2.0]}
assert session.run(inputs) == {"y": [0.0, 2.0]}

然后运行带有上述 Python 代码的 python3！

有关 Python 包的更多详细信息，包括构建说明，请参阅wonnx-py。

在浏览器中使用WebGPU + WebAssembly

npm install @webonnx/wonnx-wasm

然后，在客户端

import init, { Session, Input } from "@webonnx/wonnx-wasm";

// Check for WebGPU availability first: if(navigator.gpu) { .. }
await init();
const session = await Session.fromBytes(modelBytes /* Uint8Array containing the ONNX file */);
const input = new Input();
input.insert("x", [13.0, -37.0]);
const result = await session.run(input); // This will be an object where the keys are the names of the model outputs and the values are arrays of numbers.
session.free();
input.free();

该包 @webonnx/wonnx-wasm 提供了访问WONNX的接口，它作为WebAssembly模块包含在内，并使用浏览器的WebGPU实现。有关更完整的用法示例（包括打包器），请参阅 wonnx-wasm-example。

有关包括构建说明在内的JS/WASM包的更多详细信息，请参阅 wonnx-wasm。

开发

要为wonnx本身工作，请遵循以下步骤

• 安装Rust
• 为GPU API安装Vulkan、Metal或DX12。
• 使用git克隆此仓库。

git clone https://github.com/webonnx/wonnx.git

然后，一切准备就绪！您可以通过cargo运行其中包含的示例之一

cargo run --example squeeze --release

运行其他模型

• 要运行ONNX模型，首先使用 nnx prepare 简化它（在仓库内部使用时替换为 cargo run -- prepare）

nnx prepare -i ./some-model.onnx ./some-model-prepared.onnx

要指定动态维度参数，例如添加 --set batch_size=1。

您还可以使用外部工具，例如 onnx-simplifier，使用以下命令

# pip install -U pip && pip install onnx-simplifier
python -m onnxsim mnist-8.onnx opt-mnist.onnx

• 然后，您可以使用CLI（请参阅 README 或按照 示例文件夹中的示例 编程方式运行它。要运行示例

cargo run --example mnist --release

测试过的模型

• Squeezenet
• MNIST
• BERT

GPU选择

除非在WebAssembly中运行，否则您可能需要设置以下环境变量以影响WGPU的GPU选择

• WGPU_ADAPTER_NAME 使用您要使用的适配器名称的子串（例如，1080 将匹配 NVIDIA GeForce 1080ti）。
• WGPU_BACKEND 使用您要使用的后端的逗号分隔列表（vulkan、metal、dx12、dx11 或 gl）。
• WGPU_POWER_PREFERENCE 使用当未指定特定适配器名称时的电源偏好（high 或 low）

贡献：实现新运算符

即使没有DL、WGSL或Rust的丰富经验，也欢迎贡献。我希望这个项目可以成为我们学习这些技术的沙盒，而不仅仅是项目的初始范围。

要实现运算符，您只需做以下事情

1. 在 compiler.rs 中添加一个新的匹配模式
2. 使用 get_attribute 函数检索其属性值

    let alpha = get_attribute("alpha", Some(1.0), node);
    // or without default value
    let alpha = get_attribute::<f32>("alpha", None, node);

3. 使用 context 在WGSL着色器中添加您想要使用的任何变量。
4. 在 templates 文件夹中编写一个新的WGSL模板。

可用类型在 structs.wgsl 中，但您也可以在模板中生成新的类型。

5. 请尊重每个条目的绑定布局，从0开始递增，输入在前，输出在后。如果绑定数超过4，则递增绑定组。您可以在 sequencer.rs 中更改输入。
6. 编写逻辑。

在上下文中存在默认变量

• {{ i_lens[0] }}：输入0的长度。这同样适用于输出：{{ o_lens[0] }}和其他输入{{ i_lens[1] }}
• {{ i_shape[0] }}：输入0的维度数组。要获取数组的第一个维度，只需使用：{{ i_shape[0][0] }}
• {{ i_chunks[0] }}：输入0每个维度的块大小。默认情况下，每个变量表示为长值数组，要获取特定值，您必须通过块来移动。这些块在此变量中表示。要获取第一个维度的块大小，请使用：{{ i_chunks[0][0] }}。
• {{ op_type }}：操作类型，如激活等使用相同的模板。

7. 使用utils函数进行测试，并将其放置在tests文件夹中。测试可能如下所示

#[test]
fn test_matmul_square_matrix() {
    // USER INPUT

    let n = 16;
    let mut input_data = HashMap::new();

    let data_a = ndarray::Array2::eye(n);
    let mut data_b = ndarray::Array2::<f32>::zeros((n, n));
    data_b[[0, 0]] = 0.2;
    data_b[[0, 1]] = 0.5;

    let sum = data_a.dot(&data_b);

    input_data.insert("A".to_string(), data_a.as_slice().unwrap());
    input_data.insert("B".to_string(), data_b.as_slice().unwrap());

    let n = n as i64;
    let model = model(graph(
        vec![tensor("A", &[n, n]), tensor("B", &[n, n])],
        vec![tensor("C", &[n, n])],
        vec![],
        vec![],
        vec![node(vec!["A", "B"], vec!["C"], "MatMul", "MatMul", vec![])],
    ));

    let session =
        pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");

    let result = pollster::block_on(session.run(input_data)).unwrap();

    // Note: it is better to use a method that compares floats with a tolerance to account for differences
    // between implementations; see `wonnx/tests/common/mod.rs` for an example.
    assert_eq!((&result["C"]).try_into().unwrap(),sum.as_slice().unwrap());
}

查看tera文档以获取其他模板操作：https://tera.netlify.app/docs/

8. 如果您在任何时候想要对多个节点进行优化，可以在sequencer.rs中进行。

支持的运算符（参考ONNX IR）

已知限制

• The Clip, Resize, Reshape, Split, Pad and ReduceSum ops接受（通常是可选的）二级输入来设置各种参数（例如，axis）。这些输入只有在作为初始化器张量提供时才受支持（即不依赖于输入并且不是其他操作的输出），因为WONNX预先编译所有操作到着色器中（并且必须事先知道这些参数）。

• 内部不支持64位整数（原因是它们在当前的WGSL版本中不受支持）；64位标量输入和初始化器转换为32位值（可能溢出）。

• 对于MatMul和Gemm，矩阵维度必须能被2整除，或者输出矩阵的大小必须是(1, N)。矩阵乘法只支持浮点数，不支持整数（这是WebGPU/WGSL的限制）。

形状推理

WONNX需要知道每个操作输入和输出张量的形状，以便生成执行它的着色器代码。然而，ONNX模型并不总是包含中间值的信息。形状推理是从输入和输出的形状以及每个操作的特性推断中间值形状的过程。

WONNX支持有限的形状推理形式（即确定模型图中各个节点形状的过程）。形状推理可以通过程序方式和CLI来实现。在进行形状推理之前，所有动态维度参数都需要替换为静态值。形状推理仅从输入形状推断特定支持的操作的输出形状（请参阅上面的表格）。如果节点的任何输入形状未知，推理将无法成功。已经为输出定义了完全形状的节点保持不变（并且输出用于使用这些输出作为输入的节点上的形状推理）。

要使用CLI执行形状推理，运行类似于以下命令的命令（这里batch_size和sequence_length是动态维度参数；-i标志启用形状推理）

nnx prepare model.onnx model-prepared.onnx --set batch_size=1 --set sequence_length=255 -i

要程序化执行形状推理，请使用来自wonnx_preprocessing::shape_inference模块的apply_dynamic_dimensions和infer_shapes。

常量折叠

一些模型包含输出可以静态确定的子图，因为它们不依赖于推理期间提供的特定输入。WONNX可以将此类常量中间值替换为静态值（“常量折叠”）。以下情况下支持

• Constant操作类型的节点输出（这些用初始化器替换）
• Shape操作类型的节点输出，其中输入的形状已知（在推理前或推理过程中）
• 所有输入都是常量（可能经过折叠）的节点输出，并且操作由WONNX支持。

常量折叠作为形状推理的一部分执行，除非禁用（从CLI传递--no-fold-constants以禁用）。这是为了支持使用Shape/Squeeze/Unsqueeze等操作动态计算形状的模型，这些操作依赖于动态设置的维度参数（例如批量大小）。

许可证

根据以下任一许可证获得许可

• Apache License, Version 2.0 (LICENSE-APACHE或http://www.apache.org/licenses/LICENSE-2.0）
• MIT许可证（LICENSE-MIT或http://opensource.org/licenses/MIT）由您选择。

除以下文件外

• 数据/模型:

  • mobilenetv2-7.onnx：[源](https://github.com/onnx/models/blob/main/vision/classification/mobilenet/model/mobilenetv2-7.onnx)，仅Apache-2.0许可证。
  • squeezenet-labels.txt：[源](https://github.com/onnx/models/blob/main/vision/classification/synset.txt)，仅Apache-2.0许可证。

• 数据/图像:

  • pelican.jpeg：[源](https://en.wikipedia.org/wiki/Pelican#/media/File:Pelikan_Walvis_Bay.jpg)，版权所有 Rui Ornelas，[CC-BY 2.0](https://creativecommons.org/licenses/by/2.0/)。
  • bald_eagle.jpeg：[源](https://en.wikipedia.org/wiki/Bald_eagle#/media/File:Bald-Eagle-9114-cropped.jpg)，版权所有 David R. Tribble，[CC BY-SA 4.0](https://creativecommons.org/licenses/by-sa/4.0/)。

贡献

除非您明确声明，否则您有意提交以供包括在您的工作中的任何贡献都应如上所述双许可，不附加任何额外条款或条件。