在 WebAssembly （WASM） 数据流图中运行 ONNX 推理

本文介绍如何在 WebAssembly 模块中嵌入和运行小型开放式神经网络交换（ONNX）模型，以作为 Azure IoT作数据流图的一部分执行带内推理。 在流式处理数据上直接使用此方法进行低延迟扩充和分类，而无需调用外部预测服务。

为何使用带内 ONNX 推理

使用 Azure IoT 操作数据流图，您可以直接在管道中嵌入小型 ONNX 模型推理，而不是调用外部预测服务。 此方法提供了多种实际优势：

• 低延迟：在数据到达的相同运算符路径中执行实时扩充或分类。 每条消息只涉及本地 CPU 推断，避免网络往返。
• 紧凑设计：面向 MobileNet 级模型等紧凑模型。 此功能不适用于大型转换器模型、GPU/TPU 加速或频繁的 A/B 模型推出。
• 针对特定用例进行优化：
  • 与多源流处理一致，其中特征已在图中并置
  • 与事件时间语义保持一致，因此推理使用与其他运算符相同的时间戳
  • 保持足够小以嵌入到模块中，且不会超出实际的 WASM 大小和内存限制
• 简单更新：使用 WASM 和嵌入式模型交付新模块，然后更新图形定义参考。 无需单独更改模型注册表或外部终结点。
• 硬件约束：ONNX 后端通过 WebAssembly 系统接口（WASI） wasi-nn在 CPU 上运行。 无 GPU/TPU 目标；仅运行支持的 ONNX 算子。
• 水平缩放：推理缩放为数据流图缩放。 当运行时为吞吐量添加更多辅助角色时，每个辅助角色都会加载嵌入式模型并参与负载均衡。

何时使用带内 ONNX 推理

如果您有以下需求，请使用带内推理：

• 低延迟需求，要求在引入时对消息进行内联扩充或分类
• 小型高效模型，如 MobileNet 类视觉模型或类似的紧凑型模型
• 推理需要与事件时间处理保持一致，且时间戳与其他算子相同
• 通过交付具有更新的模型的新模块版本来简单更新

具有以下需求时，不要使用带内推理：

• 大型 Transformer 模型、GPU/TPU 加速或复杂的 A/B 部署
• 模型需要多个张量输入、键值缓存或不受支持的 ONNX 算子

先决条件

在开始之前，请确保具备：

• 具有数据流图功能的 Azure IoT 运维部署。
• 访问容器注册表（如 Azure 容器注册表）。
• 为 WebAssembly 模块开发设置的开发环境。

有关详细的设置说明，请参阅 开发 WebAssembly 模块。

体系结构模式

数据流图中 ONNX 推理的常见模式包括：

1. 预处理数据：转换原始输入数据以匹配模型的预期格式。 对于图像模型，此过程通常涉及：
   • 解码图像字节。
   • 调整为目标尺寸（例如 224×224）。
   • 将颜色空间（例如 RGB 转换为 BGR）。
   • 将像素值规范化为预期范围（0-1 或 -1 到 1）。
   • 在正确的张量布局中组织数据：NCHW（批量、通道、高度、宽度）或 NHWC（批量、高度、宽度、通道）。
2. 运行推理：使用接口将预处理的数据转换为张量，使用 wasi-nn CPU 后端加载嵌入式 ONNX 模型，在执行上下文上设置输入张量，调用模型的转发传递，并检索包含原始预测的输出张量。
3. 后处理输出：将原始模型输出转换为有意义的结果。 常见操作
   • 应用 softmax 以生成分类概率。
   • 选择前 K 个预测结果。
   • 应用置信度阈值来筛选低置信度结果。
   • 将预测索引映射到人工可读标签。
   • 格式化结果以供下游使用。

在 Rust WASM 操作员的 IoT 示例 中，可以找到两个遵循此模式的示例：

• 数据转换“格式”示例：将图像解码并调整为 RGB24 224×224。
• 图像/视频处理“快照”示例：嵌入 MobileNet v2 ONNX 模型、运行 CPU 推理和计算 softmax。

配置图形定义

若要在数据流图中启用 ONNX 推理，需要同时配置图形结构和模块参数。 图形定义指定管道流，而模块配置允许运行时自定义预处理和推理行为。

启用 WASI-NN 支持

若要启用 WebAssembly 神经网络接口支持，请将 wasi-nn 该功能添加到图形定义中：

moduleRequirements:
  apiVersion: "0.2.0"
  hostlibVersion: "0.2.0"
  features:
    - name: "wasi-nn"

定义操作和数据流

配置构成推理流程的操作。 此示例显示了典型的图像分类工作流：

operations:
  - operationType: "source"
    name: "camera-input"
  - operationType: "map"
    name: "module-format/map"
    module: "format:1.0.0"
  - operationType: "map"
    name: "module-snapshot/map"
    module: "snapshot:1.0.0"
  - operationType: "sink"
    name: "results-output"

connections:
  - from: { name: "camera-input" }
    to: { name: "module-format/map" }
  - from: { name: "module-format/map" }
    to: { name: "module-snapshot/map" }
  - from: { name: "module-snapshot/map" }
    to: { name: "results-output" }

此配置创建管道，其中：

• camera-input 从源接收原始图像数据
• module-format/map 预处理图像（解码、调整大小、格式转换）
• module-snapshot/map 运行 ONNX 推理和后处理
• results-output 向接收器传送分类结果

配置模块参数

定义运行时参数以自定义模块行为，而无需重新生成。 这些参数在初始化时传递给 WASM 模块：

moduleConfigurations:
  - name: module-format/map
    parameters:
      width:
        name: width
        description: "Target width for image resize (default: 224)"
        required: false
      height:
        name: height
        description: "Target height for image resize (default: 224)"
        required: false
      pixelFormat:
        name: pixel_format
        description: "Output pixel format (rgb24, bgr24, grayscale)"
        required: false

  - name: module-snapshot/map
    parameters:
      executionTarget:
        name: execution_target
        description: "Inference execution target (cpu, auto)"
        required: false
      labelMap:
        name: label_map
        description: "Label mapping strategy (embedded, imagenet, custom)"
        required: false
      scoreThreshold:
        name: score_threshold
        description: "Minimum confidence score to include in results (0.0-1.0)"
        required: false
      topK:
        name: top_k
        description: "Maximum number of predictions to return (default: 5)"
        required: false

操作员 init 可以通过模块配置接口读取这些值。 有关详细信息，请参阅 模块配置参数。

打包模型

将 ONNX 模型直接嵌入 WASM 组件可确保原子部署和版本一致性。 此方法简化了分发，并删除了外部模型文件或注册表上的运行时依赖项。

模型准备指南

在嵌入模型之前，请确保它满足 WASM 部署的要求：

• 对于实际 WASM 加载时间和内存约束，请将模型保持在 50 MB 以内。
• 验证模型是否接受采用通用格式（float32 或 uint8）的单张张量输入。
• 验证 WASM ONNX 运行时后端是否支持模型使用的每个操作员。
• 使用 ONNX 优化工具减小模型大小并提高推理速度。

嵌入工作流

按照以下步骤嵌入模型和关联的资源：

1. 组织模型资产：将 .onnx 模型文件和可选 labels.txt 文件放在源树中。 使用专用目录结构，例如 src/fixture/models/ 和 src/fixture/labels/ 用于明确的组织。
2. 在编译时嵌入：使用特定于语言的机制在二进制文件中包括模型字节。 在 Rust 中，使用 include_bytes! 处理二进制数据，使用 include_str! 处理文本文件。
3. 初始化 WASI-NN 图：在操作员的 init 函数中，根据嵌入的字节创建图形 wasi-nn ，并指定 ONNX 编码和 CPU 执行目标。
4. 实现推理循环：对于每个传入消息，预处理输入以匹配模型要求，设置输入张量，执行推理，检索输出并应用后处理。
5. 正常处理错误：为模型加载失败、不受支持的运算符和运行时推理错误实现正确的错误处理。

有关完整的实现模式，请参阅 “快照”示例。

建议的项目结构

组织 WASM 模块项目，明确分离关注点：

src/
├── lib.rs                 # Main module implementation
├── model/
│   ├── mod.rs            # Model management module
│   └── inference.rs      # Inference logic
└── fixture/
    ├── models/
    │   ├── mobilenet.onnx      # Primary model
    │   └── mobilenet_opt.onnx  # Optimized variant
    └── labels/
        ├── imagenet.txt        # ImageNet class labels
        └── custom.txt          # Custom label mappings

示例文件引用

使用“快照”示例中的以下文件布局作为参考：

• 标签目录 - 包含各种标签映射文件
• Models 目录 - 包含 ONNX 模型文件和元数据

最小嵌入示例

以下示例演示了最少的 Rust 嵌入。 路径相对于包含宏的源文件：

// src/lib.rs (example)
// Embed ONNX model and label map into the component
static MODEL: &[u8] = include_bytes!("fixture/models/mobilenet.onnx");
static LABEL_MAP: &[u8] = include_bytes!("fixture/labels/synset.txt");

fn init_model() -> Result<(), anyhow::Error> {
  // Create wasi-nn graph from embedded ONNX bytes using the CPU backend
  // Pseudocode – refer to the snapshot sample for the full implementation
  // use wasi_nn::{graph::{load, GraphEncoding, ExecutionTarget}, Graph};
  // let graph = load(&[MODEL.to_vec()], GraphEncoding::Onnx, ExecutionTarget::Cpu)?;
  // let exec_ctx = Graph::init_execution_context(&graph)?;
  Ok(())
}

性能优化

若要避免为每个消息重新创建 ONNX 图形和执行上下文，请初始化一次并重复使用它。 公共示例使用静态LazyLock：

use std::sync::LazyLock;
use wasi_nn::{graph::{load, Graph, GraphEncoding, ExecutionTarget}, GraphExecutionContext};

static mut CONTEXT: LazyLock<GraphExecutionContext> = LazyLock::new(|| {
  let graph = load(&[MODEL.to_vec()], GraphEncoding::Onnx, ExecutionTarget::Cpu).unwrap();
  Graph::init_execution_context(&graph).unwrap()
});

fn run_inference(/* input tensors, etc. */) {
  unsafe {
    // (*CONTEXT).set_input(...)?; (*CONTEXT).compute()?; (*CONTEXT).get_output(...)?;
  }
}

部署您的模块

重用简化的示例生成器或在本地构建：

• 若要使用简化的 Docker 生成器，请参阅 Rust Docker 生成器示例
• 有关常规 WebAssembly 部署和注册表步骤，请参阅 将 WebAssembly 与数据流图配合使用

遵循以下部署过程：

1. 在发布模式下生成 WASM 模块并生成 <module-name>-<version>.wasm 文件。
2. 使用 OCI 注册表作为存储（ORAS），将模块及（可选的）图形定义推送到注册表。
3. 在 Azure IoT操作中创建或复用注册表终结点。
4. 创建引用图定义项目的数据流图资源。

示例：MobileNet 图像分类

IoT 公共示例提供两个用于图像分类的样本，这些样本被连接到图中：“格式”示例 提供图像解码和调整大小功能，“快照”示例 提供 ONNX 推理和 softmax 处理。

若要部署此示例，请从公共注册表拉取项目，将其推送到注册表，并部署数据流图，如 示例 2：部署复杂图形。 复杂图形使用这些模块来处理图像快照并发出分类结果。

局限性

WASM 数据流图中的推理具有以下限制：

• 仅限 ONNX。 数据流图不支持其他格式（如 TFLite）。
• 仅限 CPU。 无 GPU/TPU 加速。
• 建议使用小型模型。 不支持大型模型和内存密集型推理。
• 支持单张量输入模型。 不支持多输入模型、键值缓存和高级序列或生成方案。
• 确保 WASM 运行时中的 ONNX 后端支持模型的运算符。 如果不支持运算符，推理会在加载或执行时失败。

相关内容

• 开发 WebAssembly 模块
• 配置 WebAssembly 图定义
• 将 WebAssembly 与数据流图配合使用