WebAssembly AI推理

概述

WebAssembly通过wasi-nn接口支持AI推理，使小型ML模型可以在边缘的WASM组件中运行。¹

wasi-nn核心概念

wasi-nn是WASI的机器学习接口规范，允许WebAssembly程序访问host提供的ML功能：

抽象	描述
Backend	推理后端（OpenVINO/ONNX Runtime/TFLite）
Graph	模型图（加载为不透明字节）
Tensor	输入/输出张量
ExecutionTarget	CPU/GPU执行目标

设计原则

wasi-nn采用”图加载器”API设计：

┌─────────────┐     不透明字节      ┌─────────────┐
│   Wasm Guest │ ────────────────▶  │   Host      │
│              │                    │ OpenVINO/   │
│              │ ◀───────────────   │ ONNX Runtime│
│              │   输出张量         │              │
└─────────────┘                    └─────────────┘

特点：

• 模型作为字节传递，无需转换
• host处理硬件加速（GPU/SIMD）
• 图加载器与具体框架解耦

接口定义

WIT格式

package wasi:nn;
 
interface nn {
    // 图编码类型
    enum graph-encoding {
        openvino,
        onnx,
        tflite,
        pytorch,
    }
 
    // 执行目标
    enum execution-target {
        cpu,
        gpu,
        tpu,
    }
 
    // 图操作
    load: func(data: list<list<u8>>, encoding: graph-encoding) -> result<graph, error>;
    init-execution-graph: func(graph: graph) -> result<execution-graph, error>;
    
    // 执行操作
    compute: func(execution-graph: execution-graph, inputs: list<u32>, outputs: list<u32>) -> result<_, error>;
    
    // 张量操作
    get-output: func(execution-graph: execution-graph, index: u32) -> result<list<u8>, error>;
}

模型部署实战

模型准备

# 1. 导出为ONNX格式
python convert_to_onnx.py --model mobilenet_v2 --input-size 224x224
 
# 2. 优化模型
python -m onnxoptimizer --input mobilenetv2.onnx --output mobilenetv2_optimized.onnx
 
# 3. 验证模型大小（建议<50MB）
ls -lh mobilenetv2_optimized.onnx

模型准备指南

要求	建议值	说明
模型大小	<50MB	WASM加载时间与内存限制
输入格式	单张量float32/uint8	简化预处理
运算符	主流运算符	检查运行时支持
输出	单张量	多输出需额外处理

Rust实现

use wasi_nn::*;
 
pub struct ImageClassifier {
    graph: Graph,
    execution_ctx: ExecutionGraph,
    input_tensor: Tensor,
    output_tensor: Tensor,
}
 
impl ImageClassifier {
    pub fn new(model_bytes: &[u8]) -> Result<Self, Error> {
        // 1. 加载模型
        let graph = unsafe {
            load(
                &[model_bytes.to_vec()],
                GRAPH_ENCODING_ONNX,
                EXECUTION_TARGET_CPU,
            )
        }.map_err(|e| Error::ModelLoadFailed)?;
 
        // 2. 初始化执行图
        let execution_ctx = unsafe {
            init_execution_graph(graph)
        }.map_err(|e| Error::InitFailed)?;
 
        Ok(Self {
            graph,
            execution_ctx,
            input_tensor: vec![0f32; 3 * 224 * 224],
            output_tensor: vec![0f32; 1000],
        })
    }
 
    pub fn classify(&mut self, image_data: &[u8]) -> Result<Vec<f32>, Error> {
        // 1. 预处理图像
        self.preprocess_image(image_data);
 
        // 2. 设置输入张量
        let input_indices = &[0u32];
        let output_indices = &[0u32];
 
        // 3. 执行推理
        unsafe {
            compute(
                self.execution_ctx,
                input_indices,
                output_indices,
            )
        }.map_err(|e| Error::ComputeFailed)?;
 
        // 4. 获取输出
        let output = unsafe {
            get_output(self.execution_ctx, 0)
        }.map_err(|e| Error::GetOutputFailed)?;
 
        Ok(self.parse_classification_output(&output))
    }
}

Azure IoT Operations集成

Azure IoT Operations支持在数据流图中嵌入ONNX推理：

# graph-definition.yaml
moduleRequirements:
  apiVersion: "1.1.0"
  runtimeVersion: "1.1.0"
  features:
    - name: "wasi-nn"
 
modules:
  - name: image-classifier
    type: wasm
    path: ./image_classifier.wasm
    config:
      model: ./models/mobilenet_v2.onnx
      labels: ./models/labels.txt

边缘计算架构

典型部署模式

┌─────────────────────────────────────────────────────────────────┐
│                     Edge Node                                   │
├─────────────────────────────────────────────────────────────────┤
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐              │
│  │ Data Flow   │  │   WASM      │  │   ONNX      │              │
│  │ Graph       │──▶│  Component  │──▶│  Runtime    │              │
│  └─────────────┘  └─────────────┘  └─────────────┘              │
│         │                │                │                      │
│         ▼                ▼                ▼                      │
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐              │
│  │ Preprocess  │  │ Inference   │  │ Postprocess │              │
│  │ (Rust)      │  │ (wasi-nn)  │  │ (Rust)      │              │
│  └─────────────┘  └─────────────┘  └─────────────┘              │
└─────────────────────────────────────────────────────────────────┘
           │                │                │
           ▼                ▼                ▼
      Sensor Data     ML Model        Classification Results

冷启动优势

指标	容器方案	WASM方案
冷启动	100-500ms	<1ms
内存占用	200-500MB	10-50MB
模型推理延迟	10-50ms	15-60ms

运行时支持

Wasmtime wasi-nn

use wasmtime::{Engine, Config};
use wasmtime_wasi_nn::NNTable;
 
let mut config = Config::default();
config.wasi_nn_plugin_nn_modules(true);
 
let engine = Engine::new(&config)?;
let mut linker = Linker::new(&engine);
 
// 添加wasi-nn支持
wasmtime_wasi_nn::wit::add_to_linker(&mut linker, |ctx| ctx)?;

支持的Backends

Backend	编码格式	加速	生产就绪
OpenVINO	ONNX	CPU SIMD	✅
ONNX Runtime	ONNX	CPU/GPU	✅ (需配置)
TFLite	TFLite	CPU	⏳

版本状态

# Cargo.toml
wasmtime-wasi-nn = "41"  # 最新稳定版
 
[features]
default = ["openvino", "onnx"]  # 支持多种backend
onnx-cuda = ["onnx", "ort/cuda"]  # CUDA加速

生产最佳实践

1. 模型嵌入

将模型字节直接嵌入WASM组件：

// 在build时嵌入模型
include_bytes!("../models/mobilenet_v2.onnx");
 
pub fn load_model() -> &'static [u8] {
    include_bytes!("../models/mobilenet_v2_optimized.onnx")
}

优势：

• 原子部署（模型+逻辑一起版本化）
• 无外部依赖
• 简化分发

2. 错误处理

pub enum InferenceError {
    ModelLoadFailed(String),
    UnsupportedOperator(String),
    TensorShapeMismatch { expected: Vec<u32>, got: Vec<u32> },
    BackendError(String),
}
 
impl std::fmt::Display for InferenceError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            InferenceError::ModelLoadFailed(msg) => write!(f, "Model load failed: {}", msg),
            InferenceError::UnsupportedOperator(op) => write!(f, "Unsupported operator: {}", op),
            InferenceError::TensorShapeMismatch { expected, got } => {
                write!(f, "Tensor shape mismatch: expected {:?}, got {:?}", expected, got)
            }
            InferenceError::BackendError(msg) => write!(f, "Backend error: {}", msg),
        }
    }
}

3. 性能优化

// 预热执行（避免首次调用延迟）
pub fn warm_up(&mut self) {
    // 用虚拟数据预热
    let dummy_input = vec![0f32; 3 * 224 * 224];
    let _ = self.classify(&dummy_input);
}
 
// 批量处理
pub fn batch_classify(&mut self, images: Vec<&[u8]>) -> Vec<Vec<f32>> {
    images
        .iter()
        .map(|img| self.classify(img).unwrap())
        .collect()
}

已知限制

当前限制

限制	说明	缓解方案
仅CPU	无GPU/TPU加速	选择CPU优化模型
仅ONNX	不支持TFLite/PyTorch直接格式	转换模型
单输入	多输入模型不直接支持	预处理合并
模型大小	大模型内存受限	模型量化/剪枝

适用场景

✅ 适合WASM推理：

• 边缘数据预处理+分类
• 轻量级模型（<50MB）
• 延迟不敏感应用
• 需低冷启动时间

❌ 不适合WASM推理：

• 大型模型（LLM、扩散模型）
• GPU密集型推理
• 实时视频处理
• 需要最新硬件加速

参考资料

Footnotes

1. wasi-nn - WebAssembly/WASI ↩