wonnx - 基于Rust的多平台GPU加速ONNX推理工具

项目介绍：WONNX

WONNX 是一个专门为 Web 应用提供的 GPU 加速 ONNX 神经网络推理运行时。其核心代码是用 Rust 语言百分百编写的。

支持的平台

WONNX 利用 wgpu 支持多个平台的 GPU 加速：

Windows: 支持 Vulkan 和 DX12，如果是 Windows 10，还支持 DX11。
Linux 和 Android: 支持 Vulkan 和 OpenGL ES3。
macOS 和 iOS: 支持 Metal。

快速开始

WONNX 提供多种使用路径，以下是各平台的使用方法：

从命令行

首先确保您的系统支持 Vulkan、Metal 或 DX12。然后可以下载一个二进制版本，或者安装 Rust 并运行以下命令来安装 CLI 工具：

cargo install --git https://github.com/webonnx/wonnx.git wonnx-cli

CLI 工具名为 nnx，它为操作 ONNX 模型提供了便利的接口。例如，查看模型信息或进行推理：

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 wonnx

查看 examples 文件夹中的示例代码来学习如何使用 WONNX，或者浏览 WONNX 的 API 文档。

从 Python 编程语言

运行以下命令安装 WONNX：

pip install wonnx

然后在 Python 代码中使用：

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]}

在浏览器中使用 WebGPU + WebAssembly

首先安装 WONNX 的 NPM 包：

npm install @webonnx/wonnx-wasm

然后，在客户端代码中导入和初始化：

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

// 首先检查 WebGPU 的可用性：if(navigator.gpu) { .. }
await init();
const session = await Session.fromBytes(modelBytes);
const input = new Input();
input.insert("x", [13.0, -37.0]);
const result = await session.run(input);
session.free();
input.free();

开发指南

若要参与 WONNX 的开发，首先需要安装:

Rust 编译器
Vulkan、Metal 或 DX12 作为 GPU API
克隆项目代码库

克隆项目的命令：

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

然后可以通过 cargo 运行示例：

cargo run --example squeeze --release

支持的模型

WONNX 已测试过以下模型：

Squeezenet
MNIST
BERT

GPU 选择

除了在 WebAssembly 中运行时外，您可以设置以下环境变量以影响 GPU 的选择：

WGPU_ADAPTER_NAME: 选择您想使用的适配器名的一部分（如: 1080 匹配 NVIDIA GeForce 1080ti）。
WGPU_BACKEND: 指定您想使用的后端列表（如: vulkan, metal, dx12, dx11, 或 gl）。
WGPU_POWER_PREFERENCE: 在未指定具体适配器名称时选择电源偏好（如: high 或 low）。

贡献指南

项目欢迎所有对深度学习、WGSL 或 Rust 感兴趣的贡献者，即使之前没有相关经验。通过该项目，大家可以共同学习这些技术。有关实现新操作符的详细步骤可以参考项目代码库。

总之，WONNX 致力于为 Web 应用提供高效的深度学习模型推理解决方案，凭借 Rust 的安全性和 WebGPU 的高性能，开发者能够轻松高效地处理复杂的神经网络推理任务。