TensorFlow LiteをUnityから動かすその2

この記事は Unity #2 Advent Calendar 2019 - Qiita の1日目です。他の日程に面白そうなタイトルが沢山並んでいます。

TL;DR;

TensorFlow LiteのUnity Pluginを使い MNIST, SSD, DeepLab, PoseNetなどの基本的なものを動作させるサンプルを作りました。
github.com

TensorFlow Liteとは

TensorFlowをスマートフォンやRaspberry Piなどのデバイスで動かすことを目的としています。学習はTensorFlow本体のpythonなどでモデルを作り、それを*.tfliteというファイルに書き出すことで、色々なデバイスで推論部分のみを小さなプログラムサイズで動かせるようになっています。コア部分はC++でUnity Pluginも非常に簡単に作ることができます。

先日この記事を書きました。 asus4.hatenablog.com

その後、GPUで動かすことによる高速化。複数プラットフォーム対応、入力と出力の型が見えないなどの足りない機能を追加して、使えそうな感じになってきました。TensorFlow の本家リポジトリへも PRを送ってやり取り頑張っています。
Pull Requests · tensorflow/tensorflow · GitHub

TensorFlow Lite Unity Pluginの使い方。

MNISTと呼ばれる0から9の手書き数字識別するプログラムの使い方をサンプルコードで説明します。簡略化しているので実際に動作するコードはGitHubを見てください。

    // 入力の画像 28 x 28のグレースケール
    float[,] inputs = new float[28, 28];

   // 出力の10個の数字の確率
    float[] outputs = new float[10];

    void Start()
    {
        // モデルのbyte配列を読み込んで`Interpreter`を作る
        interpreter = new Interpreter(File.ReadAllBytes(path));

        // 入力のバッファを確保
        interpreter.ResizeInputTensor(0, new int[] { 1, 28, 28, 1 });
        interpreter.AllocateTensors();
    }
    
   void Update()
   {
        // 入力データをセット
        interpreter.SetInputTensorData(0, inputs);

        // なんか計算してくれる！
        interpreter.Invoke();
        
        // 出力のデータを取得
        interpreter.GetOutputTensorData(0, outputs);
    }

    void OnDestroy()
    {
        // Interpreterを開放
        interpreter?.Dispose();
    }

とてもシンプルにできています。サンプルでもコードの多くは、画像をinputsに変換する部分だったりします。
また次のコードの用に、コンストラクタのオプションで、マルチスレッドや、GPUを使った高速化などを設定できます。

    void Start()
    {
        // オプション
        var options = new Interpreter.Options()
        {
            // CPU時にマルチスレッドで動作 
            threads = 2,
            // iPhone, macでメタルGPU上で動かす
            gpuDelegate = new MetalDelegate(new MetalDelegate.Options()
            {
                allowPrecisionLoss = false,
                waitType = MetalDelegate.WaitType.Passive,
            })
        };
        
        // option付きのコンストラクタ
        interpreter = new Interpreter(File.ReadAllBytes(path), options);
    }

Android/iOSで書かれた公式のExamplesのいくつかをUnityへ移植したので、それぞれのサンプルの説明をしていきます。

MNIST

Tensorflow Lite CPU on Unity is 100x faster #madewithunity pic.twitter.com/CRw9f495a4
— Koki Ibukuro (@asus4) November 15, 2019

MLのHallo World的存在MNISTです。28 x 28 ピクセルの数字を 0 - 9までのどれが書かれたの確率を教えてくれます。

f:id:asus4:20191128200929p:plain — MNIST ネットワーク図

ネットワーク図を見るとぎょっとしますが、レイヤーの機能は理解ぜずとも、inputs/outputsの部分に注目すれば大丈夫です。

入力

index	type	dimensions	説明
0	float	[28,28,1]	28x28のグレースケール画像です。0.0 ~ 1.0へ正規化します。

出力

index	type	dimensions	説明
0	float	[10]	10個の0.0 ~ 1.0までの確率です

SSD - Single Shot MultiBox Detector

画像上の物体の名前と矩形範囲を推定する。サンプルに使ったモデルでは90種類の物体を認識します。

SSD runs fast enough on Unity iPhoneXS #madewithunity pic.twitter.com/xPWjuAgYF0
— Koki Ibukuro (@asus4) November 28, 2019

f:id:asus4:20191128201008p:plain — SSD ネットワーク図

入力

index	type	dimensions	説明
0	sbyte	[1,300,300,3]	300x300のRGB画像です。それぞれのピクセルは 0 - 255

出力

認識されたオブジェクト上位10種の情報を返します。

index	type	dimensions	説明
0	float	[10, 4]	10個のバウンディングボックスがtop, left, bottom, rightの順で入っています
1	float	[10]	認識した物体の IDが入っています。floatだけど、intにキャストして使えば大丈夫
2	float	[10]	それぞれの確率が入っています。
3	float	[1]	全体で認識した物体の数

DeepLab

SSDと似て、画像の中の物体を推定しますが、矩形ではなく、ピクセル単位で推定します。

Added a DeepLab segmentation sample. #TensorFlowLite #madewithunity https://t.co/KzfXcQrFTH pic.twitter.com/FSulqALIqm
— Koki Ibukuro (@asus4) November 28, 2019

f:id:asus4:20191128201042p:plain — DeepLab ネットワーク図

入力

index	type	dimensions	説明
0	float	[1,257,257,3]	257x257のRGB画像です。それぞれのピクセルは 0.0 ~ 1.0 に正規化しています

出力

すべてのピクセルを21種類のラベルに分類します。

index	type	dimensions	説明
0	float	[1,257,257,21]	257x257のピクセル全てにどのラベルの可能性が高いかのスコアがはいっています

PoseNet

画面上の人間の骨格を、2D座標で取得するプログラムです。

f:id:asus4:20191128201139p:plain — PoseNet ネットワーク図

直列のとてもシンプルなネットワーク図に見えます。その分ポストプロセスでやることがあります。

入力

index	type	dimensions	説明
0	float	[1,257,257,3]	257x257のRGB画像です。それぞれのピクセルは 0.0 ~ 1.0 に正規化しています

出力

index	type	dimensions	説明
0	float	[9,9,17]	9x9のブロックに分割したエリアそれぞれに17個のパーツIDの可能性が入っています。
1	float	[9,9,34]	9x9のブロックのそれぞれのパーツIDの本当の場所へのオフセット x,y値
2	float	[9,9,32]	今回1人の認識では使わない
3	float	[9,9,32]	同上

使うモデルによってブロックの分割数、精度が変わります。

出力したデータの整形が一番面倒なモデルです。このブログが一番わかりやすかったです。 medium.com

正確さを犠牲にシンプルに言えば、

9 x 9に分割したブロックそれぞれに、体のパーツ17種の確率が入っている。
一番確率が高いパーツにオフセット値を足して、正しいXY座標を得る
17x17分割のモデルを使うことで精度を上げることが可能。遅くなる

ようです。

Unityでの実装で注意すること。

Height, Widthの順番: UnityではすべてのAPIがwidth, heightで表されますが、TF Liteでは height, widthです。
モデルの種類によっては、GPUオプションが使えないものもある
- ドキュメントにあるops (演算子?)をサポートしてます

TODOs

Android版は私の端末Sony Xperia XZ2で動作中にフリーズするので、原因を調査
- 熱によるカメラのフリーズ問題のよう
Windows, Linuxのプラグインを作る

まとめ

TensorFlow LiteはUnityでも使えそう
公式のモデルをUnityで動くサンプルを作った

TensorFlow Liteの公式サンプルでは、Android/iOSネイティブは全く別の実装をしています。Unity PluginではAndroid/iOS/Unity Editorを同じプログラムで作れるので便利です。今回使ったのは決して最新ではなく、ある程度こなれたモデルですが、Unityで使えるようになることで、ゲーム組み合わせたりするとまた可能性が広がりそうだと思いました。
Texture2DのピクセルをC#上で直接触っているところなどがあり、まだ高速化できる余地があるので、もう少し触っていきます。実案件へは未投入ですが、仕事で使いたいって相談もお待ちしています。

Asus4 onD