今月は仕事の空き時間はずっとTensorFlow Lite C-API上でのGPU delegateの高速化に取り組んでいました。高速化といっても、GPU delegate内部に手を付けるのではなく、CPU → GPU間をデータ転送するときの転送待ち時間をなくす作業をしていました。
GPUでの計算はCPUとは比較にならないほど早いですが、結局なにかロジックを変更するのはCPU。GPUですべてのプログラムを動かすことはできません。
そして、CPUからGPUへのデータのアップロード、ダウンロードは、CG周りをやる方はよく悩んでいると思いますが、かなり…遅いのです。
私はUnity上でTensorFlow Liteを使っていて、Unity上ではほとんどの画像をGPU上で扱っています。しかし、TensorFlow LiteのAPIにはCPUメモリ上でデータをやり取りするするものしか用意されていませんでした。
その問題はこちらのissueでやり取りしています。↓ github.com
例えば、セマンティックセグメンテーションのタスクをUnityのTextureで実行して、Unity上で結果を表示しようとすると以下のような経路を辿らなくてはいけません。
Unity GPU → Unity CPU → TFL CPU → TFL GPU Delegate → TFL CPU → Unity CPU → Unity GPU
Unityで扱っているGPU上のデータをTensorFlowのCPU側のAPIを抜かして、TensorFlowのGPU Delegateへ直接送ることができればかなり高速化できるはずです。
実際にGoogleの公開しているMediaPipeというライブラリではそのような方法で高速化をしているようです。（逆にGitHub上でコードを探した限り、MediaPipeでしかやられてないんじゃないかとも思ってます。）
色々とやっていた試行錯誤は↓こちら。Zennのスクラップへまとめています。 zenn.dev
そして、最終的には以下のように、UnityのGPUとTensorFlowを直接つなげることができました。
Unity GPU → TFL GPU Delegate → Unity GPU
iOSやmacOSで使われているMetal Delegate ではUnityのComputeBufferはMTLBufferへ繋がります。
Androidの環境はちょっと複雑で、
GL DelegateとGPU Delegate v2という2種類が用意されています。
GL DelegateはOpenGLES上で計算するもの。 GPU Delegateは、内部的には、OpenCL, OpenGLESの２パターンがOpenCLに対応してないAndroidの場合に自動的にOpenGLES側へフォールバックするようです。この辺ちゃんとドキュメントにかかれてなくて、ソースコード読んでやっと理解しました。またOpenCLで動いている場合でもOpenGLの Shader Storage Buffer Object (SSBO)へとつなげることができるようです。
世界中でどれだけ需要があるかわかりませんが、8割方の案件がARで機械学習をつかってる私が使うためにTensorFlow本体の方へも反映すべく、PRやり取り中です。マージされるといいな〜。
