エッジデバイス推論向けの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
   • 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
   • TensorFlow Model Optimization Toolkit - Pruning API
5. Binarization(Extremely low-bit quantization)
   • 現時点では利用できません
6. Neural Architecture Search(NAS)
   • AutoML Vision Edge modelを利用することで取り込むことは可能です。

プロセッサごとの動き

モバイル(CPU+GPU)

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

GPU(NVIDIA)

• 準備中のLow-order precision floating point(TFLiteConverterによるFP16変換)とは別途、TensorRTによるTensorFlowモデルのインポートが利用可能です。

  • TensorFlowではありませんが、TensorRTの利用効果については当ブログの過去記事でも紹介していますので、あわせて参照ください。 tech-blog.optim.co.jp

• GTC 2017の発表内容によれば、TensorRTによる高速化の1手法としてINT8量子化を採用しているとのことです。

  • TensorRT: Int8 calibration with hand-tuned scale factors

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変換を行っており、手法としての差異はありません。

論文

• 量子化の効用については、例えばCloud TPU(v1)のアーキテクチャ論文にて以下のような記述があります。
  • In-Datacenter Performance Analysis of a Tensor Processing Unit
    • 推論時にかかるFP16(IEEE754)の乗算をINT8の乗算にすると、エネルギー消費とプロセッサ面積は1/6に節約。加算であればエネルギーは1/13、プロセッサ面積は1/38に節約
  • 参考記事: GoogleのTPUって結局どんなもの？ 日本法人が分かりやすく説明
• また、Binarizationの1つであるXNOR-Netの論文にも以下のような記載があります。
  • ImageNet Classification Using Binary Convolutional Neural Networks
    • 推論時のCNN演算weightを2値化(強い量子化)すると、1/32のメモリ使用量で58倍速
• 精度劣化を最小限に抑えつつ解析的にactivationの量子化を行う方法については、以下論文にて示されています。
  • PACT: Parameterized Clipping Activation for Quantized Neural Networks
  • Post-training 4-bit quantization of convolution networks for rapid-deployment
  • ResNet系だとINT4量子化まではFP16レベルの精度を維持できることが読み取れます。
  • ただし、INT8未満のbit数にて量子化を行っても速度パフォーマンス改善につながらないケースもあります。(プロセッサが低い桁数の整数表現に対応していない場合など)
• weightやactivationに対する量子化ではなく、gradientやerrorを対象にした量子化手法もあります。
  • DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients
  • Training and Inference with Integers in Deep Neural Networks
• Compressionについては下記に記載されています。
  • Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding
• NASについては、例えばMobileNetV3のモデル作成にモバイル向けNAS(MnasNet)が利用されています。
  • Searching for MobileNetV3
  • MnasNet: Platform-Aware Neural Architecture Search for Mobile

おわりに

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

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

www.optim.co.jp

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

www.optim.co.jp