はなちるのマイノート

Unityをメインとした技術ブログ。自分らしくまったりやっていきたいと思いますー!

【Unity】Tensorflow Liteを利用してUnityで機械学習の推論を行う(tf-lite-unity-sampleを利用)

はじめに

Unity SYNCを見ていたところ、以下の講演にて面白い事を話されていました。

youtu.be

初手Tensorflow Lite(以下 TFlite)でCPU推論がお勧め。理由は以下。

  • ONNXRuntimeやBarracudaと比較して多くのケースで推論が最速

私もONNXRuntimeBarracudaを利用したことがあるのですが、Tensorflow Liteがそんなに早いとは知りませんでした。

というかどの推論ライブラリを利用するか選定していたときに完全に見落としていて、やってしまったなという感じですね。

導入も簡単そうなので、実際にTensorflow Liteをお試しで触ってみようと思います。

環境

MacOS
Unity 2021.3.0f1
tf-lite-unity-sample 2.10.0

やり方

動画でも紹介されていましたが、Koki Ibukuroさんが作られたtf-lite-unity-sampleを利用することで簡単にUnityに導入することができます。
github.com


↓ 作者さんのブログ
asus4.hatenablog.com
asus4.hatenablog.com

こちらを利用していきたいと思います。

TensorFlow Lite librariesをインストールする

サンプルコードを試したいならリポジトリごと取ってくれば何か特別な操作も必要なく動作確認できます。
またTensorFlow Lite librariesのみを利用したい場合はUPMを使ってできるようにしてくれているみたいです。(神機能)

manifest.jsonに以下のコードを加えます。

{
  "scopedRegistries": [
    {
      "name": "package.openupm.com",
      "url": "https://package.openupm.com",
      "scopes": [
        "com.cysharp.unitask"
      ]
    },
    {
      "name": "npm",
      "url": "https://registry.npmjs.com",
      "scopes": [
        "com.github.asus4"
      ]
    }
  ],
  "dependencies": {
    // Core TensorFlow Lite libraries
    "com.github.asus4.tflite": "2.10.0",
    // Utilities for TFLite
    "com.github.asus4.tflite.common": "2.10.0",
    // Utilities for MediaPipe
    "com.github.asus4.mediapipe": "2.10.0",
    ...// other dependencies
  }
}


入れてみたサンプル↓

MNISTを動作するようにしてみる

モデル(tflite)をサンプルから利用させていただき、本当に動作するのか確認してみます。
tf-lite-unity-sample/mnist.tflite at master · asus4/tf-lite-unity-sample · GitHub

mnist.tfliteStreamingAssetsに入れて以下のコードを書きました。

using System.IO;
using TensorFlowLite;
using UnityEngine;

public class Test : MonoBehaviour
{
    // 入力:28x28のグレスケール
    private readonly float[,] _inputs = new float[28, 28];

    // 出力:0~9の数字の確率
    private readonly float[] _outputs = new float[10];
    
    // 推論
    private Interpreter _interpreter;
    
    private void Start()
    {
        // モデルを読み込む
        var modelPath = Application.streamingAssetsPath + "/mnist.tflite";
        var model = File.ReadAllBytes(modelPath);
        
        // Optionを設定する
        var options = new InterpreterOptions()
        {
            threads = 2,
            useNNAPI = false,
        };
        
        // Interpreterを生成する
       _interpreter = new Interpreter(model, options);
        
        // 入力バッファを確保する
        _interpreter.ResizeInputTensor(0, new int[] { 1, 28, 28, 1 });
        _interpreter.AllocateTensors();
    }

    private void Update()
    {
        // 入力データをセット
        _interpreter.SetInputTensorData(0, _inputs);
        
        // 推論開始
        _interpreter.Invoke();

        // 出力データを取得
        _interpreter.GetOutputTensorData(0, _outputs);
        
        // 一番確率が高いものを出力する
        var number = 0;
        var max = 0f;
        for (var i = 0; i < _outputs.Length; i++)
        {
            if (!(_outputs[i] > max)) continue;
            number = i;
            max = _outputs[i];
        }
        Debug.Log(number);
    }

    private void OnDestroy()
    {
        _interpreter?.Dispose();
    }
}

かなりベタ書きですが、私の環境ではエラーなく動作しました。

_inputsの中身をちゃんと変えてあげるようにすれば、サンプルのようなものも得られるはずです。
github.com

またサンプルではTextureからグレースケールの配列に変換するのにCompute Shaderを利用しているようですね。賢いと思います。

さいごに

manifest.jsonに書き込むだけでセットアップが完了するまで整えてくれたリポジトリ作者さんには本当に感謝です。

tfliteへのモデルの変換が少し大変そう(まだ試せていない)ですが、速度が早いだけでなくInterpreterOptionsも優秀で素晴らしいと思います。