Unity BaracudaでTensorFlowおよびONNXモデルを使用する方法

2020年3月9日 23:09

以下の記事を参考に書いてます。

・How to use (some) TensorFlow and ONNX computer vision models in Unity

1. はじめに

しばらく前に、「TensorFlow Sharp Plugin」によるを使用したUnityでのTensorFlowモデルの使い方を紹介しました。「画像分類」は十分に機能しましたが、「物体検出」のパフォーマンスは低いという結果になりました。それでも、Unityで機械学習を必要とする人にとっては、良い出発点になると思いました。

しかし残念なことに、Unityは「TensorFlow」のサポートを終了し、「Barracuda」というコードネームの独自の推論エンジンの開発に移行しました。「TensorFlow Sharp Plugin」も引き続き使用できますが、TensorFlow 1.7.1向けに実装されいるため、それ以降のバージョンで訓練されたモデルはまったく動かない可能性があります。

「Baracuda」はTensorFlowのモデルを「Baracuda」のサポート形式に簡単に変換できます。しかしまだプレビュー中であり、注意事項があります。

2. Baracudaで特定のアーキテクチャを試す方法

TensorFlowのAndroidのサンプルを元に、Unityで同様のサンプルを作ることにしました。mobilenet_v1アーキテクチャも試してみました。サンプルにはありませんが、必要なのはinput/output名とstd/mean値を置き換えることだけです。

「Baracuda」では、もう少し複雑です。Unityで特定のアーキテクチャを試すには、3つの方法があります。

(1) 既に持っているONNXモデルを使用。
(2) TensorFlowからONNXコンバータを使用してTensorFlowモデルを変換。
(3) Unityが提供するTensorFlowからBarracudaスクリプトを使用してBarracuda形式に変換。

「inceptionモデル」や「ssd-mobilenetモデル」はいずれもうまくいきませんでした。変換に問題があったか、Unityがサポートされていないテンソルを使っていたためロードできないか、推論中にクラッシュするか、無意味な結果を返します。事前に訓練された「inception ONNX」は、動いたように見えましたが、途中で奇妙なエラーや結果が出力されました。

しかし、いくつかのモデルはうまくいきました。

3. 画像分類

「MobileNet」は、その名前からわかるように、まさにモバイルのために作成されたた、モバイル推論に優れたアーキテクチャです。小さく、高速であり、サイズ/レイテンシと精度のトレードオフを提供するさまざまなバージョンがあります。

私は最新のmobilenet_v3を試しませんでしたが、v1とv2はONNXとしても、tf-barracuda変換でもうまく機能しました。.onnxモデルがある場合、そのまま設定できます。.pbモデル（protobuf形式のTensorFlowモデル）がある場合、「tensorflow-to-barracuda converter」を使用して簡単に変換できます。

$ python tensorflow_to_barracuda.py ../mobilenet_v2_1.4_224_frozen.pb ../mobilenet_v2.nn

Converting ../mobilenet_v2_1.4_224_frozen.pb to ../mobilenet_v2.nn
Sorting model, may take a while...... Done!
IN: 'input': [-1, -1, -1, -1] => 'MobilenetV2/Conv/BatchNorm/FusedBatchNorm'
OUT: 'MobilenetV2/Predictions/Reshape_1'
DONE: wrote ../mobilenet_v2.nn file.

コンバータは、inputs/outputs自体を把握します。これらは、後でコードで変更する必要がある場合に備えておくと便利です。ここでは、入力名「input」と出力名「MobilenetV2/Predictions/Reshape_1」があります。また、検査のためにこのモデルを選択すると、Unity Editorでそれらを見ることができます。

注意すべき点が1つあります。「mobilenet_v2」では、コンバータとUnityインスペクタに誤った入力次元が表示されます。代わりに[1, 224, 224, 3]にする必要がありますが、実際にはこれは重要ではないようです。

次に、ドキュメントに記載されているように、「Barracuda」を使用してこのモデルをロードして実行できます。

public NNModel modelFile;
...
var model = ModelLoader.Load(this.modelFile);
var worker = WorkerFactory.CreateWorker(WorkerFactory.Type.ComputePrecompiled, model);
...
var inputs = new Dictionary<string, Tensor>();
inputs.Add(INPUT_NAME, tensor);
worker.Execute(inputs);
var output = worker.PeekOutput(OUTPUT_NAME);

推論作業を行うために入力画像で行わなければならない重要なことは、「正規化」です。これは、ピクセル値を[0; 255]の範囲から[-1; 1]にシフトおよび縮小することを意味します。

public static Tensor TransformInput(Color32[] pic, int width, int height)
{
    float[] floatValues = new float[width * height * 3];

    for (int i = 0; i < pic.Length; ++i)
    {
        var color = pic[i];

        floatValues[i * 3 + 0] = (color.r - IMAGE_MEAN) / IMAGE_STD;
        floatValues[i * 3 + 1] = (color.g - IMAGE_MEAN) / IMAGE_STD;
        floatValues[i * 3 + 2] = (color.b - IMAGE_MEAN) / IMAGE_STD;
    }

    return new Tensor(1, height, width, 3, floatValues);
}

「Barracuda」には実際にはTexture2Dからテンソルを作成するメソッドがありますが、スケーリングとバイアスのパラメータは受け入れません。画像で推論を実行する前に必要なステップであることが多いため、これは奇妙です。

独自のモデルを試す際には注意が必要です。一部のモデルでは、バイアスレイヤーをモデル自体の一部として実際にスケーリングしている可能性があるため、使用する前にUnityで必ず確認してください。

4. 物体検出

現在最も頻繁に使用されている2つの物体検出アーキテクチャがあるようです。「SSD-MobileNet」と「YOLO」です。

残念ながら、SSDは「Barracuda」によってまだサポートされていません（この問題で述べられているように）。「YOLO v2」で解決する必要がありましたが、もともと「YOLO」は「DarkNet」に実装されており、TensorflowまたはONNXモデルを取得するには、まず「DarkNet」の重みを必要な形式に変換する必要があります。

幸いなことに、すでに変換されたONNXモデルが存在しますが、完全なネットワークはモバイル推論には大きすぎるように思えたため、ここで入手できるTiny-YOLO v2モデル（opsetバージョン7または8）を選択しました。しかし、すでにTensorflowモデルを持っている場合、「tensorflow-to-barracuda converter」は同様に機能します。実際、試用できるリポジトリ内のWorksフォルダにもあります。

面白いことに、ONNXモデルには既に画像ピクセルを正規化するためのレイヤーがありますが、このモデルは正規化を必要とせず[0; 255]の範囲のピクセルで正常に動作するため、実際には何もしないようです。

YOLOの最大の問題は、その出力が「SSD-Mobilenet」よりもはるかに多くの解釈を必要とすることです。ONNXリポジトリからの「Tiny-YOLO」出力の説明は次のとおりです。

「出力は（125x13x13）テンソルで、13x13は画像が分割されるグリッドセルの数です。各グリッドセルは125チャンネルに対応し、グリッドセルによって予測される5つのバウンディングボックスと25個のデータ要素で構成されます各境界ボックス（5x25 = 125）を記述します。」

幸いなことに、Microsoftは.NETでのオブジェクト検出にONNXモデルを使用するための優れたチュートリアルを用意しています。

5. ブロック

「TensorFlow Sharp」のサンプルは、メインスレッドで1秒に1回推論を実行するだけで、カメラの再生をブロックしているため、並列処理の点でかなりいまいちでした。別のスレッドでモデルを実行するなど、より合理的なアプローチを示す他のサンプルがありました（MatthewHallbergに感謝）。

ただし、「Barracuda」を別のスレッドで実行しても機能しなかったため、見苦しいクラッシュが発生しました。ドキュメントで判断すると、「Barracuda」はデフォルトで非同期であり、GPUで推論を自動的にスケジュールする（可能な場合）ので、Execute()を呼び出して、後で結果を照会するだけです。実際には、まだ警告があります。

そのため、「Barracuda workerクラス」には、推論できるExecute()、ExecuteAsync()、ExecuteAndWaitForCompletion()の3つのメソッドがあります。最後の1つはブロックします。プロセス中にアプリをフリーズさせない限り、使用しないでください。Execute()は非同期に機能するため、次のようなことができるはずです。

worker.Execute(inputs);
yield return new WaitForSeconds(0.5f);
var output = worker.PeekOutput(OUTPUT_NAME);

または完全に別のメソッドで出力をクエリします。ただし、Execute()を呼び出しただけで、カメラフィードにわずかなジッターが生じる場合でも、わずかな遅延があることに気付きました。これは新しいデバイスでは目立たない可能性があるため、購入する前に試してください。

ExecuteAsync()は、推論を非同期的に実行するための非常に優れたオプションのようです。StartCoroutine(worker.ExecuteAsync(inputs))で実行できる列挙子を返します。ただし、内部的にこのメソッドは各レイヤーの後にnullを返します。これは、フレームごとに1つのレイヤーを実行することを意味します。それができるよりも（私が見つけたように、画像分類のためのmobilenetモデルの場合です）。YOLOモデルは、他のメソッドよりもExecuteAsync()でうまく動作するように見えますが、テストできるデバイスの量はかなり限られています。

モデルを実行するためのさまざまな方法で遊んで、別の可能性を見つけました。ExecuteAsync()はIEnumeratorなので、手動で反復して、フレームごとに必要な数のレイヤーを実行できます。

var enumerator = this.worker.ExecuteAsync(inputs);

while (enumerator.MoveNext())
{
    i++;
    if (i >= 20)
    {
        i = 0;
        yield return null;
    }
};

これは少しハックであり、完全にプラットフォームに依存していません。自分で判断してください。

6. 結論

繰り返しますが、「Baracuda」はまだプレビュー段階のため、多くのことが頻繁にそして根本的に変わります。Barracuda 0.4.0でテストされている私のサンプルは、0.5.0では間違った結果が返されます（サンプルを試す場合は必ず0.4.0をインストールしてください。新しいバージョンについては後で調べます）。

しかし、TensorFlowがクロスプラットフォームで動作し、ONNXをサポートし、Unityに埋め込むことができる推論エンジンは素晴らしいです。

「Baracuda」のサンプルは、「TensorFlow Sharp」のサンプルよりもパフォーマンスが優れているかどうかというと、判断は困難です。

特に物体検出は異なるアーキテクチャが使用されいるため困難です。「Tiny-YOLO」は「SSD-Mobilenet」よりもラベルがはるかに少なく、TFSharpのサンプルはOpenGLESで高速ですが、「Barracuda」はVulcanでより高性能です。非同期戦略も異なります。

Galaxy S8でMobileNet v1と画像分類を比較すると、推論時間は、OpenGLES 2/3でのTFSharpで約400ミリ秒、VulcanでのBarracudaで約110ミリ秒となりました。しかし、より重要なことは、「Baracuda」は今後数年間開発、サポート、改善される可能性が高いということです。したがって、将来さらに大きなパフォーマンスの向上が期待できます。

私のgithubで完全なコードをチェックしてください。

この記事が気に入ったらサポートをしてみませんか？