エッジデバイス推論向けのTensorFlow最適化を掘り下げてみる #GoogleMLSummit

まえがき

R&Dチームの徳田(@dakuton)です。 昨日、Google主催の機械学習エンジニア向けイベント「Google Developers ML Summit Tokyo」に参加してきました。 このイベントでは、エッジデバイス推論(On-Device Machine Learning)に関するセッションもいくつかありました。

  • 該当セッション
    • On-Device Machine Learning with TensorFlow Lite
    • TensorFlow Model Optimization:Quantization and Sparsity

今回は、上記セッションで新たにわかったこと+自分たちで調べていた内容を整理してご紹介します。同じような取り組みをされている方の参考になれば幸いです。

エッジ推論のメリットと課題

  • クラウド実行と比較して、下記のようなメリットが挙げられます。
    1. 推論リクエスト通信によって発生するlatencyを抑える
    2. ネットワークから切断された環境でも推論可能
    3. プライバシー配慮データが必要な際、アップロードなしで推論可能
  • 一方で、4つのリソース制約に対処する必要があります。(精度との両立が求められる)
    1. プロセッサ(CPUやGPUなど)が限られている
    2. メモリが限られている
    3. バッテリー消費
    4. モデルサイズ

主なエッジ推論最適化

セッションでは1〜4について説明されていました。

  1. Post-training quantization
    • 学習済みモデルのもつweightやactivationを少ないbit数の整数表現に変換
    • Hybrid quantization: weightのみ量子化
    • Integer quantization: activationとweightの両方を量子化
  2. Quantization-aware training(During training)
    • モデルの生成時にあらかじめ量子化することを前提に学習
  3. Low-order precision floating point
    • 整数ではなく、FP32->FP16など少ない実数表現に変換
  4. Compression
    • Pruning(ネットワークのweightが0となる数を増やして計算短縮)やファイル圧縮などを組み合わせて、実行速度とモデルサイズ削減のバランスをとる
    • ML Kitにて提供予定と発表されたLearn2Compressもこちらに該当します。
  5. Binarization(Extremely low-bit quantization)
    • 2値化(bit数が整数よりさらに小さい量子化)表現に置き換えて、論理演算可能にすることで高速化
  6. Neural Architecture Search(NAS)
    • 強化学習や遺伝的アルゴリズムを用いたモデルネットワーク・パラメータの最適化

TensorFlow 2.0でのAPI対応状況

準備中のものやTensorFlow以外を利用するものもあるため、現在のサポート状況を整理しておきます。

  1. Post-training quantization
  2. Quantization-aware training(During training)
    • 現時点では利用できません(Converter Python API guide)が、tf.keras向けに準備中とのアナウンスがありました。
  3. Low-order precision floating point
    • 現時点では利用できませんが、TFLiteConverterによるFP16変換を準備中とのアナウンスがありました。
      • TargetSpecにてsupported_type=[tf.float16]を指定してPost-training実施
  4. Compression
  5. Binarization(Extremely low-bit quantization)
    • 現時点では利用できません
  6. Neural Architecture Search(NAS)

プロセッサごとの動き

モバイル(CPU+GPU)

  • CPU実行環境での量子化モデル推論高速化として、gemmlowpが該当します。
    • 例えば、ARM NEON対応のCPUであれば、INT8に精度を落とすことで並列計算が可能なSIMD命令セットに置き換えて実行されます。
  • また、GPU利用を指定した場合でも部分的にCPUにて動作し、CPUのもつオペレータによって上記と同等の効果が得られることがあります。
    • TFLite on GPUでサポート記載外のオペレータが必要な場合、CPU側のオペレータにフォールバック実行されます。
  • 下記に端末ごとのベンチマーク結果がまとまっていて参考になります。

GPU(NVIDIA)

Edge TPU

  • Edge TPUでの高速実行のためにINT8への量子化が必要です。
  • Edge TPUの専用オペレータを利用するため、2のconvert過程でINT8量子化を行います。(Post-trainingではなく、Quantization-aware trainingが必要)
    1. TensorFlowによる学習
    2. TFLiteConverterを用いたTensorFlow Liteモデルへの変換
    3. Edge TPU Compilerによる変換
  • 上記変換を行わなかった場合やEdge TPUで非サポートのオペレータの場合、CPU側のオペレータにフォールバック実行されます。
    • なおこちらは実動作ベース(論文やドキュメント等で明確な記載がない)での確認であり、今後動作が変化する可能性があります。
    • 参考: TensorFlow models on the Edge TPU

(追記)ブラウザ

  • TensorFlow.jsを用いたブラウザ推論に対応する場合はtensorflowjs_converterによるモデル変換を行います。このとき量子化オプションとしてquantization_bytesを指定します。
    • 利用するconverterは別ですがTensorFlow LiteのPost-training quantizationと同様のweight変換を行っており、手法としての差異はありません。

論文

おわりに

TensorFlow 2.0正式リリース向けて準備中のものもありますが、公式にてサポートされた最適化手法も増えてきたという点は非常にありがたいです。 ただし、期待する高速化・最適化が実際に行われているか?については各デバイスごとに動かしてみてベンチマークをとらなければわからない部分も多いため、調べてみて現時点でも不明な部分についてはいくつか試してみたいと思います。

オプティムの今年のインターン(R&Dコース)では、本記事にてご紹介したようなエッジデバイスを用いた推論もテーマとして選択できます。 興味がある方はぜひご応募ください。

www.optim.co.jp

インターン以外でも募集中です。

www.optim.co.jp