Onnx RuntimeをUnityで動かす

Onnx Runtimeをネイティブプラグインとして、Unity上で動かす実験とサンプルを公開しています。

開発の動機

4年前に、TensorFlow LiteをUnityで動かす実験を初めて、はじめは全くの趣味で始めたものが、今では海外からいただく相談の半分以上が機械学習関連になっています。

四年前に始めた実験↓ asus4.hatenablog.com

ところが、実際にシェアを見ると、研究関連ではPytorchのシェアが圧倒的。Unityの公式推論ライブラリBarracudaやTensorFlow Liteで動かすために一旦Onnxに変換するなどの事例なども増え始め、速度的にはTFLiteは非常に満足していますが、サクッとモデルを試してみたいという時に、変換するのが億劫になってきていました。公式ツールで変換しようにもOnnxやPytorchのNCHWからTFLiteのNHWCに変換するときに大量のTransposeが挟まり速度が逆に遅くなることがあるのが不満でもありました。（この辺の高速化は PINTOさんのonnx2tfなどのツールでも対応されています）

Unity SentisのOnnx対応

Unityの公式ML推論ライブラリBarracudaもOnnxフォーマットを読み込みます。今年リニューアルしてSentisという名前になりました。 unity.com

実際試したことがある方はわかるかも知れませんが、SentisではOnnxフォーマットを読み込みますが、実際の実行エンジンはUnityが独自開発しているため、対応オペレーターに結構差があります。Onnx Runtimeで動いてもSentisで動かないことがままあります。（体感的には読み込み成功の打率は半分以下な気がします。）

Sentisの前身、Barracudaでの説明ではありますが、KeijiroさんによるCEDEC公演でも、後半、結構トリッキーなことをして、Onnxモデルの非対応オペラーターをBrracuda上で対応する構造に書き換えるということをしています。

youtu.be

Sentis自体はマルチプラットフォーム対応を謳い、今後対応オペレーターの互換性も増えていくと思いますが、Onnx自体の進化も早く、今後完全なOnnx互換となることは難しいように思います。

もちろんNintendo Switch, PlaystationからWeb Playerまでを対応しなくては行けないUnityが、独自推論エンジンSentis開発を進めることは正しく思いますが、私のようにiOS,Android,PC,macOSくらいで動けば良いユーザーからすると、Onnxの互換性が高くなると嬉しいなと思っていました。

Onnx Runtimeの対応プラットフォーム

一方、Onnx Runtimeの方の進化も早く、最新のHardware Acceleration対応を見ると、

CPU上ではXNNPack対応
iOS, macOSではCoreML
AndroidではNNApi
Windows/LinuxではDirectML,CUDA,TensorRT

と、はじめからPC, Mobile, さらにRaspberry PiのようなEdge Deviceまでを考慮に入れたプラットフォーム対応に見えます。
Googleの開発するTensorFlow Liteが、PC上でのHardware Accelerationが未だに公式にサポートされてないことを考えると、Onnx Runtimeのマルチプラットフォーム対応は、もしUnityでそのまま動けば大きな強みになりそうです。

またMicrosoftが開発していることもありC#によるC FFIライブラリのWrapperが、はじめからほぼ全て整備されていることも魅力でした。
TensorFlow Liteのときは半分以上自分でC#をFFI Wrapperを作っていたので、関心しました。

実際のサンプル

という経緯から、OnnxがUnity上で動くか試し始めたのですが、
Onnx RuntimeのC#設計が良いのかTFLiteでの経験が生きているのか、Unityで動かすのは、すんなりいきました。
現在、macOS, iOS, Androidでの動作を確認しています。

AppleのMobile OneというImage Classificationが100fpsで動く例↓

Hi ML with Unity friends, Have you ever struggled to make the ONNX model compatible with Sentis (formerly Barracuda)? I tested native ONNX runtime on Unity. It will be compatible with most models, at least on CPU. It's just PoC, but will it be an alternative ML solution? pic.twitter.com/wjIG0l5gtg
— Koki Ibukuro (@asus4) December 20, 2023

Yoloxが60fps以上で動く例↓

Testing multiplatform Onnx runtime with Unity. Yolox-nano Object Detection runs at over 60FPS. #madewithunity pic.twitter.com/ph6XFJVrK5
— Koki Ibukuro (@asus4) December 25, 2023

またライブラリに含めるサンプルも4年前からアップデートして、新し目のモデルを選んだので、4年間でのモデルの精度、高速化技術の向上に驚きました。

Onnx Runtime for Unityの使い方

シンプルな事例、Image ClassificationモデルMobile Oneで説明します。

Netronでモデルの入出力をみてみるとMobileOneは224x224のRGB画像を受取り、ImageNetの1000種類の画像分類それぞれの確率を返すシンプルなモデルです。

画像を入力として受け取るクラスでやることは大体同じなので、ImageInference.csというPreprocessingをやってくれる抽象クラスを用意しました。

Preprocess

通常はImageInference内でやってくれるので、気にする必要はありませんが、概要だけ書きます。

UnityのWebカメラはPC上では問題ないが、スマートフォンでは回転したままなので、デバイスの向きに応じて回転を補正する。TextureSource内で自動でやってくれます。
モデルの入力に合わせて、画像をリサイズする。３種類のリサイズ方法を用意しているので、用途に合わせて使い分けてください。
Aspect Mode
Onnxモデルの入力に合わせたTensorを作る。
- mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]の正規化をする。
- カメラ画像のテクスチャはNHWC(N=1)のメモリ配列なので、Onnxのメモリ配列NCHWに並べ替える。

Onnx実行

SessionOptions.Run()を実行するだけです。いくつかInput/Output Tensorの取得方法があるのですが、生のbyte配列を受け渡すより、OrtValueを使う方法がおすすめされているようです。こちらも通常はImageInference内でやってくれます。

// 事前にモデルから読み込んでOrtValueを作っておく
string[] inputNames;
OrtValue[] inputs;
string[] outputNames;
OrtValue[] outputs;

public void Run()
{
   // inputsで画像をTensorへいれる。
   PreProcess();

   // 実行
   session.Run(null, inputNames, inputs, outputNames, outputs);

   // outputsから値を取り出す。
   PostProcess();
}

Postprocess

MobieOne.cs内でPostProcessメソッドをoverrideしています。

protected override void PostProcess()
{
    // Output Tensorから値を読み込み
    var output = outputs[0].GetTensorDataAsSpan<float>();
    for (int i = 0; i < output.Length; i++)
    {
        labels[i].score = output[i];
    }
    // スコア順に並べる
    TopKLabels = labels.OrderByDescending(x => x.score).Take(topK);
}

以上です。C#のAPIがよく出来ているので、カスタマイズも色々出来そうですが、ひとまずは一番シンプルな方法を使いました。