WONNX:基于WebGPU的高性能ONNX推理运行时

wonnx

WONNX:基于WebGPU的高性能ONNX推理运行时

WONNX是一个令人兴奋的开源项目,为深度学习模型的部署和推理提供了一个高性能、跨平台的解决方案。这个项目由100%的Rust代码编写而成,利用WebGPU进行GPU加速计算,可以在原生环境和Web环境中运行ONNX格式的深度学习模型。

主要特性

WONNX具有以下几个突出的特性:

基于WebGPU的GPU加速: 利用WebGPU API实现GPU加速计算,提供卓越的推理性能。
跨平台支持: 得益于WebGPU的跨平台特性,WONNX可以在Windows、Linux、macOS等多个平台上运行。
Web兼容: 通过WebAssembly,WONNX可以直接在浏览器中运行,为Web应用提供高性能的深度学习能力。
ONNX格式支持: 兼容广泛使用的ONNX模型格式,可以直接加载和运行各种深度学习框架导出的模型。
Rust语言优势: 利用Rust语言的安全性和性能优势,提供稳定可靠的运行时环境。

支持的平台

WONNX基于wgpu库实现,因此支持多种图形API:

Windows: 支持Vulkan和DirectX 12
Linux & Android: 支持Vulkan和OpenGL ES 3
macOS & iOS: 支持Metal

这种广泛的平台支持使得WONNX成为一个真正跨平台的深度学习推理解决方案。

性能表现

根据项目提供的基准测试数据,WONNX在多个模型上展现出了优秀的性能:

模型	WONNX	Onnxruntime Web CPU	Onnxruntime Web GPU (WebGL)
MNIST	~1ms	~1ms	~1ms
Squeezenet	~26ms	~40ms	~22ms
Tiny YOLO	~200ms	~380ms	~120ms

这些数据表明,WONNX在某些模型上甚至可以超越Onnxruntime Web GPU的性能,展现出其优秀的计算效率。

使用方法

WONNX提供了多种使用方式,以适应不同的应用场景:

命令行工具: 提供了nnx命令行工具,可以方便地进行模型信息查看、推理等操作:

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 API: 可以直接在Rust项目中使用WONNX:

use wonnx::Session;

let session = Session::from_path("path/to/model.onnx")?;
let result = session.run(inputs)?;

Python绑定: 提供Python包,可以在Python环境中使用:

from wonnx import Session

session = Session.from_path("path/to/model.onnx")
result = session.run(inputs)

Web集成: 通过WebAssembly,可以在浏览器中使用WONNX:

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

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);