[1907.01989] On-Device Neural Net Inference with Mobile GPUs

Google Research より、TFLite GPU のアーキテクチャ設計についての論文。

• 1907.01989 On-Device Neural Net Inference with Mobile GPUs
• tensorflow/tensorflow/lite/delegates/gpu at master · tensorflow/tensorflow · GitHub

ざっとまとめると、

• モバイル CPU で NN の推論は難しい
  • 計算量、熱、バッテリー
  • 特に既存のロジックが CPU で動いてるのに、そこに計算負荷を追加すると…
  • 実際、Fig8 で iPhone XS の MobileNetV1 の CPU 推論、3 分ちょいで熱でパフォーマンス落ちてる
• モバイル GPU を NN の推論に使う
  • NPU などのアクセラレータ、一部のハイエンド機しかないよね
  • 多くのデバイスをターゲットとするために GPU
  • モデルにもよるが、TFLite の CPU→GPU で x2-9 は速くなる
• TFLite GPU のデザイン
  • vendor specific な実装を避け、OpenGL や Metal をバックエンドとした
    • OpenCL はサポートしてない端末多いよね
  • delegate
    • GPU オフロード可能なサブグラフをまとめて、delegate node
    • それ以外は CPU
  • 最適化
    • コマンド数やメモリ IO を少なくするため、よくある最適化はしている
    • コードだとこの辺のモジュール
    • tensorflow/fuse_inline.cc at 46252f3 · tensorflow/tensorflow · GitHub
    • operator fusion
    • 意味のない命令は消す（x1 のリサイズとか）
    • パラメータをシェーダにインライン化
    • ドライバでの最適化も期待し、TFLite としてはあくまでコードを出力
    • 一部のシェーダは人手で最適化
  • データレイアウト
    • モバイル GPU は、文字通りグラフィクス用途に適した 4 要素のベクトル（x,y,z,w）の計算や読み書きに特化
    • これに合わせて NN も計算しないと遅いので、HWC のテンソルも C を 4 チャネル単位で分解して、PHWC4 と呼ぶレイアウトにしている
    • （TFLite じゃないフレームワークでも、この”4”で pack/unpack はよく見る気がする
    • CHW->PHWC4 への変換の実装
      • tensorflow/tensorflow/lite/delegates/gpu/gl/converters at 46252f3 · tensorflow/tensorflow · GitHub
  • WorkGroup 数の決め方
    • GPU によって、WorkGroup 数のチューニングのシビアさが変わる
      • Mali はチューニング頑張っても 5%程度しか向上しないが、Adreno は 30％向上する場合も
    • 総当たりで調べるのは、デバイスの状態（熱など）にもよるので、必ずしも最適解は得にくい
    • そこで推論時間の関数を勾配法で最適化して、デバイスや属性ごとにいくつかのセットをつくってるみたい
      • 最適化部分のコードはわからないけど、結果は以下のようにディスパッチ
        • tensorflow/conv.cc at 46252f3 · tensorflow/tensorflow · GitHub
        • tensorflow/default_calculator.cc at 46252f3 · tensorflow/tensorflow · GitHub
  • メモリ管理
    • メモリのフットプリントを小さくするため、中間テンソルをいつ確保するか
    • Greedy と MCFP の 2 つのアルゴリズムを実装
      • （詳細はよんでないけど、
        • 中間テンソルは使い回す前提で、最大サイズでアロケートしておき、使いまわせるタイミングで使い回す（雑な理解
      • モデルにより良し悪しが変わるのでデフォルト Greedy
      • ただナイーブな実装よりは MB 単位で数倍は軽くなる
• 実験
  • iOS は、CoreML より TFLite GPU が総じて速い
    • （これなんでだろう？
  • Android は、vendor specific なフレームワーク（SNPE(Qualcomm)）より 2 倍程度遅いケースもある
    • それより多くの端末で動くほうが嬉しいよねという感じ
    • （MACE（Xiaomi）、こんど使ってみるかーと思ってたけど、これなら TFLite のほうが良さそう…

TFLite の実装はちゃんと見てないけど、この論文で説明がある構成になってるんだろう。 MediaPipe など見ていると、スマホ AR などで TFLite を使って様々なアプリケーションを高速に世に出していきたいんだろうなーと思えていて、それを実現する技術として、この選択肢をとったかと思うと面白い。

Table5、Android 端末で TFLite GPU と MACE（Xiaomi）、SNPE（Qualcomm）を比較してるんだけど、 この速度差ならどっちのフレームワークをつかう？ってのは人によって判断が変わりそうで面白い。 数倍の遅延も許容できない、かつ、サポートしたいプラットフォームがすくなければ、vendor specific なフレームワークでもいいかもしれない。

しかし、新しいアーキテクチャがでたとか、BlazeFace のように 7x7 がいいよみたいなときに、実際に早く速く動かせるものはどれ？と考えると、 Research→Deployment までのエコシステムの充実度、ベンダーに依存しない拡張の早さが見込めるなどで TFLite を選択するモチベーションは大きいんじゃないかな。

実装はちら見した程度だけど、カーネルの実装の中にも GpuType::ADRENO とかでてきて、やっぱり避けられないよなーという感じ。

合わせて読みたい論文としては以下。

Machine Learning at Facebook: Understanding Inference at the Edge - Facebook Research