Google Edge TPUとTensorFlow Lite(基礎篇)

こんにちは。次世代システム研究室のY.R.です。外国人です。近来Google Edge TPUとTensorFlow Liteを勉強しているところですが、こちらで習ったとこを皆さんと共有致します。宜しくお願い致します。

はじめに

エッジコンピューティング

エッジコンピューティングとは、利用者や端末と物理的に近い場所に処理装置（エッジプラットフォーム）を分散配置して、ネットワークの端点でデータ処理を行う技術の総称。多くのデバイスが接続されるIoT時代となり提唱されるようになった[1]。Google GCP[2]とAmazon AWS[3]も相関の開発を促進している。

図[1]

更に、この図で示された通り将来の数年間でスマートデバイスはすごく増加すると見通す。エッジAIは注目される領域になると信じる。これから、重要なハードウェアとソフトウェアをそれぞれに紹介する。

Cloud TPUとは

テンソル・プロセッシング・ユニット（Tensor processing unit、TPU）はGoogleが開発した機械学習に特化した特定用途向け集積回路(ASIC)。グラフィック・プロセッシング・ユニット(GPU)と比較して、ワットあたりのIOPSをより高くするために、意図的に計算精度を犠牲に（8ビットの精度）した設計となっており、ラスタライズ/テクスチャマッピングのためのハードウェアを欠いている 。チップはGoogleのテンソルフローフレームワーク専用に設計されているがGoogleはまだ他のタイプの機械学習にCPUとGPUを使用している 。他のAIアクセラレータの設計も他のベンダーからも登場しており、組み込みやロボット市場をターゲットとしている。[4]

Google TPUで以下のような特徴を持っている

• 巨大なオンチップRAMを備えている。メモリのアクセス時間を減らす。
• 8ビット 定量化(8 bit quantitation)。32ビットのfloating-point計算を入れ替わって8ビットのinteger 計算を採用される。ある程度の精度を犠牲したが、効率をよく改善した。
• Systolic arrayというアーキテクチャを使用する。Systolic array自体新しい技術でない。「積和演算」などの特定の操作用に配線されている。ディープラーニングによく適用できる。
• 良い冷却システムがある。リソースコストを節約できる。

現在普通のユーザはGCPを通してTPU [5]を利用できる。Cloud TPU Pod[6]というプラットホームを利用してもスーパーコンピュータのように動ける。

Nvidia Tesla Volta 100のGPUと比べてTPU v2の性能はそのほど高いと言えない。しかし、コストを減らすことは明らかになる。

Edge TPUとは

2019年でエッジAI に向かって、Google Coralチーム[7]はEdge TPUをリリースした。

今も様々なデバイスをリリースされた。研究のため、筆者はUSB Acceleratorを購入した。

Edge TPUの使い方は複雑ではない。普通のPCからAIモデルをEdge TPUのデバイスに導入した後で、Edge TPUと別のデバイスを一緒に動けることは可能になる。System-on-Module (SoM)でEdge TPUのデバイス自体もAI専用の端末のように動ける。

詳細的に説明したら、Edge TPUのデバイスとTensorFlow lite緊密的に繋がる。

Pythonは常に運用されるから、Python SDKも特に支えられる。

Edge TPUの性能について、上の図で示された通りEdge TPU自体の性能は強いし、AI加速装置として、Edge TPUは非常に既存のCPU性能を向上できる。

TensorFlow Lite

TensorFlow Lite は2017年11月14日に Google より公開された。TensorFlow Liteのおかげで、モバイル、組み込み、IoTデバイスでTensorFlowモデルを変換して実行できるようになる[8]。

全体的にTensorFlow Liteは上記の4つのコンポーネントで構成される。「Converter」の役割はトレーニングのTensorFlowモデルを一つの効率的なフォームに変換することになる。ユーザは「Converter」を常に使用する。「Interpreter」を通して携帯電話、Linuxシステムの組み込みハードウェア、マイクロコントローラー（マイクロコントローラー）など、さまざまなハードウェアで（最適化された）モデルは実行できる。「Interface to Hardware Acceleration」を通して、導入してモデルは異なるハードウェアに向かって最適なパフォーマンスで実行できる。

Edge TPUに向けて、トレーニングのTensorFlowモデル導入フローは以下のようになる。

先ず、既存のTensorFlowモデルを「Frozen graph」という操作で導出する。簡単的に「Frozen graph」でモデルのパラメータをフリーズされたpbファイルが作られる。そしてTensotFlow Liteのモデル転換によって、TensorFlow Liteモデルのtfliteファイルに生成する。最後、モデルファイルをEdge TPUのデバイスにロードして運行できます。携帯電話など設備へ導入とロードのフローも似ていると考える。

モデル転換はよく利用する方法ですが、様々なトレーニングのモデルもリリースしたから、現在私たちは直接的にモデルファイルをEdge TPUの設備に導入して実行できるようになる。上記の図で示されるようにこれは一番使いやすい方法になる。

Google Edge TPUとTensorFlow Liteの検証

これから簡単的な検証を皆さんに共有致します。

環境

ノートPC OS: mac 10.14

Coral USB accelerator

問題点

最初、検証はCoral AIの正式な教程[9]に参考してやってみましたが、エラーに出会った。

RuntimeError: Internal: Unsupported data type: 0Node number 1 (EdgeTpuDelegateForCustomOp) failed to prepare.

対策

相関の議論[10]を参考して、最新のEdge TPU runtime（このブログを書く時点でedgetpu_runtime_20201204.zip）とTensorFlow lite libraryの互換性において問題点がると考える。議論[11]で一つ古いバージョンのEdge TPU runtime （edgetpu_runtime_20200331.zip）を発見しました。古いバージョンのEdge TPU runtimeを利用して動けるようになる！

結果

CPUバージョンとEdge TPUバージョンのMobileNet画像分類モデルの比較は以下のようになる。

CPU バージョン

python3 classify_image.py \
> --model models/mobilenet_v2_1.0_224_inat_bird_quant.tflite \
> --labels models/inat_bird_labels.txt \
> --input images/parrot.jpg

----INFERENCE TIME----
Note: The first inference on Edge TPU is slow because it includes loading the model into Edge TPU memory.
40.3ms
37.2ms
36.5ms
36.4ms
36.5ms

Edge TPUバージョン

python3 classify_image.py \
> --model models/mobilenet_v2_1.0_224_inat_bird_quant_edgetpu.tflite \
> --labels models/inat_bird_labels.txt \
> --input images/parrot.jpg

----INFERENCE TIME----
Note: The first inference on Edge TPU is slow because it includes loading the model into Edge TPU memory.
15.5ms
3.4ms
3.4ms
3.1ms
3.1ms
-------RESULTS--------
Ara macao (Scarlet Macaw): 0.77734

判断結果は同じくレベルになったが、Edge TPUは10倍以上速くなる。

更に、正式な教程で紹介された通り、最大動作周波数（maximum operating frequency）を起動したら、Coral USB acceleratorは顯著に熱くなった。

感想と次の探索

TPUに感想

CPUとGPUと違ってTPUはAI特にディープラーニングのために設計された。巨大オンチップ RAM を備えているし、Systolic arrayのようなアーキテクチャも使用された。

しかし、GoogleによりTensorFlowしか支えられないし、Google GCPしかCloud TPUを使用できない。それに対して、ディープラーニングでPyTorchは人気があるし、PyTorch/XLA が出ている[12]。将来、筆者もPyTorch/XLAを 試してみたい。

次の探索

今のところ、TensorFlow LiteとEdge TPUの探索はまだ早い段階になります。これから、両方に深く研究したいと思う。

モデルについて、Mask R-CNN[13]系列の物体検知モデルとMobileBERT[14]という自然言語処理用の機械学習モデルに高い興味が持っている。

TensorFlow Liteのカーネルに深く理解して、性能をよく発揮させることに頑張る。

 

終わり

次世システム研究室では、ビッグデータ解析プラットホームの設計・開発を行うアーキテクトとデータサイエンティストを募集しています。興味を持って頂ける方がいらっしゃいましたら、ぜひ 募集職種一覧からご応募をお願いします。 一緒に勉強しながら楽しく働きたい方のご応募をお待ちしております。