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 crate作为依赖项添加（如果您有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]}

然后运行python3并使用上面的Python代码！

有关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）

贡献：实现新的算子

即使没有在深度学习、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开始递增1，输入在前，输出在后。如果绑定的数量超过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函数进行测试，并将其放置在测试文件夹中。测试可以如下所示

#[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/](https://tera.netlify.app/docs/)

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

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

已知限制

• 操作 Clip、Resize、Reshape、Split、Pad 和 ReduceSum 接受（通常是可选的）次要输入来设置各种参数（例如，axis）。这些输入仅支持作为初始化器张量提供时（即不依赖于输入并且不是其他操作的输出），因为WONNX预先将所有操作编译为着色器，并且必须提前知道这些参数。

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

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

形状推断

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

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

要使用命令行界面执行形状推断，请运行类似以下命令的命令（此处 batch_size 和 sequence_length 是动态维度参数；- 标志启用形状推断）

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