WebAssembly (WASM) データ フロー グラフで ONNX 推論を実行する

この記事では、小規模な Open Neural Network Exchange (ONNX) モデルを WebAssembly モジュール内に埋め込んで実行し、Azure IoT Operationsデータ フロー グラフの一部としてインバンド推論を実行する方法について説明します。 このアプローチは、外部予測サービスを呼び出さずにストリーミング データに対して直接、低待機時間のエンリッチメントと分類を行う場合に使用します。

[前提条件]

開始する前に、次の項目があることを確認します。

• データ フロー グラフ機能を備えたAzure IoT Operationsデプロイ。
• Azure Container Registryなどのコンテナー レジストリと、構成されたコンテナー レジストリ エンドポイント。 詳細については、 レジストリ エンドポイントの構成に関するページを参照してください。
• WebAssembly モジュール開発用に設定された開発環境。 オプションと詳細な環境セットアップ手順については、「 WebAssembly モジュールの開発」を参照してください。

インバンドONNX推論の利点

Azure IoT Operationsデータ フロー グラフを使用すると、外部予測サービスを呼び出す代わりに、小さな ONNX モデル推論をパイプラインに直接埋め込むことができます。 この方法には、次の利点があります。

• 待機時間が短い: データが到着するのと同じオペレーター パスで、リアルタイムのエンリッチメントまたは分類を実行します。 各メッセージにはローカル CPU 推論のみが必要です。ネットワーク ラウンド トリップを回避します。
• ストリーム処理にインラインで組み込み: 特徴がグラフ内にすでに集約されているマルチソースのストリーム処理の中で推論を実行し、イベントタイムセマンティクスに整合させることで、推論でも他の演算子と同じタイムスタンプを使用できます。
• 簡単な更新: WASM と埋め込みモデルを使用して新しいモジュールを出荷し、グラフ定義参照を更新します。 別のモデル レジストリや外部エンドポイントの変更は必要ありません。
• 水平スケーリング: 推論は、データ フロー グラフのスケーリングに合わせてスケーリングされます。 ランタイムによってスループット用のワーカーが追加されると、各ワーカーは埋め込みモデルを読み込み、負荷分散に参加します。

推論は、WebAssembly システム インターフェイス (WASI) wasi-nn インターフェイスを介して CPU 上で実行されます。 サポートされているモデル形式とハードウェアの制約については、 WASM データ フロー グラフでの ONNX 推論の制限事項を参照してください。

インバンド ONNX 推論はどのような場合に使用すべきですか？

次の機能が必要な場合は、データ フロー グラフでインバンド ONNX 推論を使用します。

• 取り込み時にメッセージをインラインで強化または分類するための低遅延。
• MobileNet クラスのビジョン モデルなど、小規模で効率的なモデル。
• イベント時処理に合わせ、他の演算子と同じタイムスタンプを使用する推論。
• 新しいモジュール バージョンを配布することで、単純なモデルを更新します。

次の機能が必要な場合は、インバンド ONNX 推論を使用しないでください。

• 大規模なトランスフォーマー モデル、GPU または TPU アクセラレーション、または高度な A/B ロールアウト。
• 複数のテンソル入力、キー値キャッシュ、またはサポートされていない ONNX 演算子を必要とするモデル。

データ フロー グラフでの ONNX 推論のアーキテクチャ パターン

データ フロー グラフでの ONNX 推論の一般的なパターンには、次のステージが含まれます。

1. データの前処理: モデルの予想される形式に合わせて生の入力データを変換します。 イメージ モデルの場合、通常、このプロセスには次の処理が含まれます。
   • 画像バイトをデコードする。
   • ターゲット ディメンションへのサイズ変更 (例: 224×224)。
   • 色空間の変換 (RGB から BGR など)。
   • ピクセル値を想定される範囲 (0 ~ 1 または -1 から 1) に正規化します。
   • 適切なテンソル レイアウトでのデータの配置: NCHW (バッチ、チャネル、高さ、幅) または NHWC (バッチ、高さ、幅、チャネル)。
2. 実行推論: wasi-nn インターフェイスを使用して前処理されたデータをテンソルに変換し、CPU バックエンドで埋め込まれた ONNX モデルを読み込み、実行コンテキストに入力テンソルを設定し、モデルのフォワード パスを呼び出して、生の予測を含む出力テンソルを取得します。
3. 後処理の出力: 生モデルの出力を意味のある結果に変換します。 一般的な操作:
   • softmax を適用して分類確率を生成します。
   • top-K 予測を選択します。
   • 信頼度しきい値を適用して、信頼度の低い結果をフィルター処理します。
   • 予測インデックスを人間が判読できるラベルにマップします。
   • 下流処理で使用するために結果をフォーマットします。

Rust WASM 演算子の IoT サンプルこのパターンに従う 2 つのサンプルを見つけることができます。

• Data 変換 "format" サンプル: 画像をデコードして RGB24 224×224 にサイズ変更します。
• Image/Video processing "snapshot" サンプル: MobileNet v2 ONNX モデルを埋め込み、CPU 推論を実行し、softmax を計算します。

グラフ定義を構成する

データ フロー グラフで ONNX 推論を有効にするには、グラフ構造とモジュール パラメーターの両方を構成します。 グラフ定義ではパイプライン フローを指定しますが、モジュール構成では、前処理と推論の動作をランタイムでカスタマイズできます。

wasi-nn 機能を有効にする

ONNX 推論に対して WebAssembly ニューラル ネットワーク (wasi-nn) インターフェイスを有効にするには、グラフ定義に wasi-nn 機能を追加します。

moduleRequirements:
  apiVersion: "1.1.0"
  runtimeVersion: "1.1.0"
  features:
    - name: "wasi-nn"

推論パイプラインの操作を定義する

ONNX 推論パイプラインを形成する操作を構成します。 この例では、一般的な画像分類ワークフローを示します。

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 モジュールの動作をカスタマイズするためのランタイム パラメーターを定義します。 これらのパラメーターは、初期化時に WASM モジュールに渡されます。

詳細については、 モジュール構成パラメーターを参照してください。

モデルをパッケージ化する

ONNX モデルを WASM コンポーネントに直接埋め込むと、アトミックデプロイとバージョンの一貫性が確保されます。 この方法により、配布が簡略化され、外部モデル ファイルまたはレジストリへのランタイムの依存関係が削除されます。

ONNX モデルの準備要件

モデルを埋め込む前に、WASM デプロイの要件を満たしていることを確認します。

• 実際の WASM 読み込み時間とメモリ制約のために、モデルを 50 MB 未満にしてください。
• モデルが共通の形式 (float32 または uint8) で 1 つのテンソル入力を受け入れることを確認します。
• WASM ONNX ランタイム バックエンドが、モデルで使用するすべてのオペレーターをサポートしていることを確認します。
• ONNX 最適化ツールを使用して、モデルのサイズを小さくし、推論速度を向上させます。

WASM モジュールに ONNX モデルを埋め込む手順

次の手順に従って、ONNX モデルと関連リソースを WASM モジュールに埋め込みます。

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. エラーを適切に処理する: モデルの読み込みエラー、サポートされていない演算子、ランタイム推論エラーに対して適切なエラー処理を実装します。

完全な実装パターンについては、"snapshot" サンプルを参照してください。

推奨される WASM プロジェクト構造

懸念事項を明確に分離して 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

スナップショット サンプルからのファイル レイアウトの例

「スナップショット」サンプルの次のファイル レイアウトを参照として使用します。

• Labels ディレクトリ - さまざまなラベル マッピング ファイルが含まれています
• Models ディレクトリ - ONNX モデル ファイルとメタデータが含まれています

ONNX モデルを埋め込むための最小 Rust の例

次の Rust の例は、WASM モジュールに ONNX モデルを埋め込むために必要な最小コードを示しています。 パスは、マクロを含むソース ファイルに対する相対パスです。

// 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 グラフと実行コンテキストを再利用する

すべてのメッセージの ONNX グラフと実行コンテキストを再作成しないようにするには、それらを 1 回初期化して再利用します。 パブリック スナップショットサンプルでは、静的なLazyLockを使用して、ワーカーごとに 1 回、グラフと実行コンテキストを初期化します。

use crate::wasi::nn::{
    graph::{load, ExecutionTarget, Graph, GraphEncoding, GraphExecutionContext},
    tensor::{Tensor, TensorData, TensorDimensions, TensorType},
};

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).compute()?;
    }
}

ONNX モジュールをローカルでデバッグしてテストする

Azure IoT Operationsにデプロイする前に、ONNX 推論モジュールをローカルでテストして、機能とパフォーマンスを検証します。 詳細については、以下をご覧ください。

• VS Code で WASM モジュールをデバッグする
• WASM モジュールをテストする

ONNX モジュールを構成してデプロイする

Azure IoT Operations データ フローまたはコネクタで ONNX 推論モジュールを使用する準備ができたら、次の手順に従います。

• データ フロー グラフとコネクタ用に WebAssembly (WASM) グラフ定義を構成する
• WebAssembly (WASM) モジュールとグラフ定義をデプロイする

例: MobileNet の画像分類

IoT パブリック サンプルには、画像分類用のグラフにワイヤードされた 2 つのサンプルが用意されています。

• "format" サンプルは、画像のデコードとサイズ変更の機能を提供します。
• "snapshot" サンプルは、ONNX 推論とソフトマックス処理を提供します。

これらのモジュールを使用するサンプルの詳細と実行については、「 例 2: 複雑なグラフをデプロイする」を参照してください。

WASM データ フロー グラフでの ONNX 推論の制限事項

WASM データ フロー グラフでの推論には、次の制限があります。

• モデル形式: ONNX のみ。 データ フロー グラフは、TFLite などの他の形式をサポートしていません。
• ハードウェア: CPU のみ。 GPU と TPU アクセラレーションはサポートされていません。
• モデル サイズ: 大規模なモデルとメモリ集中型推論はサポートされていません。 MobileNet クラスのアーキテクチャなどの小さなモデルを使用します。
• モデル入力: 単一テンソル入力モデルのみがサポートされます。 複数入力モデル、キー値キャッシュ、高度なシーケンスまたは生成シナリオはサポートされていません。
• オペレーターのサポート: WASM ランタイムの ONNX バックエンドは、モデルで使用するすべてのオペレーターをサポートする必要があります。 演算子がサポートされていない場合、読み込み時または実行時に推論が失敗します。

関連するコンテンツ

• WebAssembly モジュールの開発
• データ フロー グラフとコネクタ用に WebAssembly (WASM) グラフ定義を構成する
• WebAssembly (WASM) モジュールとグラフ定義をデプロイする