Edge TPU の性能を引き出すためには？

はじめに

オプティムの R&D チームで Deep な画像解析をやっている奥村です。

Edge TPU は NVIDIA GPU と同じような感覚で使うことはできません。NVIDIA GPU よりもメモリの制約が強く、Edge TPU の性能を引き出したり、複数のモデルを１つの Edge TPU で同時に実行するにはいくつかのコツが必要になります。Edge TPU Compiler | Coral をベースに、意訳・追記したものをメモしました。

Edge TPU はモデルのパラメータデータをキャッシュするための 8MB 程度の SRAM を持っており、ここにモデルが乗り切らない場合、都度外部メモリから Edge TPU にデータを転送する必要があるため、性能低下を引き起します。また、複数モデルを 1 つの Edge TPU で実行する場合、同時コンパイル (Co-Compile) しないと同様に性能低下を引き起こす可能性が高いです。

以下は Edge TPU の良い母艦である Jetson Nano に２つの USB アクセラレータ版 Edge TPU を接続した様子です。

Edge TPU に関しては、過去記事も参考になると思います。

• Edge TPU の July 2019 Updates を追う - OPTiM TECH BLOG
• 2019インターン紹介: Edge TPU向けPoseNetのリアルタイム姿勢推定 - OPTiM TECH BLOG

意訳

Edge TPU コンパイラ (Edge TPU Compiler)

Edge TPU コンパイラ (edgetpu_compiler) は、TensorFlow Lite モデル (.tflite ファイル) を Edge TPU に互換性のあるファイルにコンパイルするコマンドラインツールです。ここではコンパイラの使い方とどのように動作するかを少し説明します。

コンパイラを使う前に、Edge TPU に互換性のあるモデルについて確認しておく必要があります。互換性の詳細については TensorFlow models on the Edge TPU を参照してください。

システム要件 (System requirements)

Edge TPU コンパイラは、モダンな Debian ベースの Linux システムで実行可能です。特に以下です。

• Debian 6.0 以降、またはその派生 (Ubuntu 10.0 以降 など)
• x86-64 または ARMv8 命令セットを持つ ARM64 アーキテクチャのシステム

Coral Dev Board での Mendel もサポートされています。

ダウンロード

システム要件にあるような Linux システム上で、以下のコマンドでインストールできます。

curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list
sudo apt-get update
sudo apt-get install edgetpu

使い方

edgetpu_compiler [options] model...

コンパイラは、ひとつ以上の TensorFlow Lite モデルへのファイルパスを受け取り (model 引数)、いくらかのオプションを加えます。複数のモデルを（空白文字で区切って）指定すると、パラメータデータキャッシュ用に Edge TPU の RAM を共有できるように同時コンパイルされます。詳細は parameter data caching を参照してください。

コンパイルされた各モデルのファイル名は、foo.tflite を指定すると foo_edgetpu.tflite のようになり、カレントディレクトリに保存されます。--out_dir オプションでディレクトリを指定すると、そこに保存されるようになります。

利用可能なオプションは以下のとおりです。

• -o, --out_dir dir
  • コンパイルされたモデルやログを dir で指定されたディレクトリに出力します。既定はカレントディレクトリです。
• -m, --min_runtime_version val
  • 互換性のあるモデルが求める Edge TPU の最低ランタイムバージョン。例えば、モデルを実行する予定のデバイスが Edge TPU ランタイムバージョン 10 の場合 (かつランタイムバージョンを更新できない場合)、モデルの互換性を確保するために 10 と設定する必要があります。(モデルは常に新しい Edge TPU ランタイムと上位互換性があります。バージョン 10 ランタイム用にコンパイルされたモデルは、バージョン 12 と互換性があります。)
  • 既定値は使用しているコンパイラのバージョンに依存します。--help での出力で確認してください。詳細は compiler and runtime versions を確認してください。
• -s, --show_operations
  • Edge TPU にマッピングされたオペレータを示すログを出力します。foo.tflite のログは foo.log として同じディレクトリに出力されます。
• -v, --version
  • コンパイラバージョンを出力して終了します。
• -h, --help
  • コマンドラインのヘルプを出力して終了します。

パラメータデータキャッシング (parameter data caching)

Edge TPU は、モデルのパラメータデータをキャッシュすることができる 8MB 程度の SRAM を持っています。しかし、モデルの推論の実行のために、最初に少量の RAM が予約されるため、パラメータデータはその後に残る空き領域を使用します。当然、Edge TPU の RAM のパラメータを抑えると、外部メモリからパラメータデータをフェッチする場合に比べて推論速度が速くなります。

Edge TPU の「キャッシュ」は、実際には従来のキャッシュではありません。コンパイラが割り当てたスクラッチパッド（メモ用の）メモリです。Edge TPU コンパイラは、モデル内部に小さな実行可能ファイルを追加します。これが、推論を実行する前に、モデルのパラメータのうち特定量を Edge TPU RAM に書き込みます。

モデルを個々にコンパイルすると、コンパイラは各モデルにユニークな「caching token」(64 ビットの数値) を割り当てます。あるモデルを実行すると、Edge TPU ランタイムはその caching token が現在キャッシュされているデータと token と比較します。token が一致すると、ランタイムはそのキャッシュされたデータを使います。token が一致しないと、ランタイムはキャッシュを消して新しいモデルのデータを変わりに書き込みます。つまり、モデルが個々にコンパイルされると、一度にひとつのモデルだけがデータをキャッシュできます。この過程を以下に図示します。

※ Figure 1. Flowchart showing how the Edge TPU runtime manages model cache in the Edge TPU RAM を元に作成

システムは、必要になったときだけキャッシュをクリアしてモデルデータを書き込むということに注意してください。そうなると、推論に遅延が生じます。よって、モデルを最初に実行するときは常に遅めです。その後の推論は、既に書き込まれたキャッシュを使うためより高速に実行されます。しかし、アプリケーションが複数のモデルをしょっちゅう切り替える場合、キャッシュのスワッピングが発生して、アプリケーション全体の性能に大きなオーバーヘッドが加わります。そこで、同時コンパイル機能が登場します。

複数モデルの同時コンパイル (Co-compiling multiple models)

同じ Edge TPU 上で複数のモデルを連続して実行するときに性能を向上させるために、コンパイラは同時コンパイルをサポートしています。モデルを同時コンパイルすると、複数のモデルが Edge TPU RAM を共有してパラメータデータを一緒にキャッシュできるため、異なるモデルを実行する度にキャッシュをクリアする必要がなくなります。

複数のモデルをコンパイラに指定すると、コンパイルされた各モデルには同じ caching token が割り当てられます。そのため、２つ目のモデルを最初に実行するとき、最初のモデルのキャッシュをクリアせずに２つ目のモデルのデータを書き込むことができます。上記の図でいうと、２つ目の分岐が「はい」となります。

しかし、各モデルに割り当てられた RAM の量はコンパイル時に固定され、コンパイラコマンドに表示される順に基づいて優先順位付けされるという点に注意してください。例えば、以下のように２つのモデルを同時コンパイルする場合を考えます。

edgetpu_compiler model_A.tflite model_B.tflite

この場合、キャッシュの空き領域はまずモデル A のデータに（可能な限り）割り当てられます。空きがあれば、モデル B のデータにキャッシュが与えられます。モデルのデータのいくらかが Edge TPU RAM に収まらない場合、キャッシュの代わりに実行時に外部メモリからフェッチする必要があります。

複数のモデルを同時コンパイルすると、いくつかのモデルはキャッシュを得られない可能性があり、外部メモリからすべてのデータを読み込む必要があります。この場合、キャッシュを使う場合よりも当然遅くなります。とはいえ、複数のモデルを素早く連続して実行する場合、異なるモデルを実行するたびにキャッシュをスワップするよりは、同時コンパイルした方がまだ高速に実行できるでしょう。

注意：パラメータデータは、一度にひとつのレイヤーをキャッシュするために割り当てられます。特定のレイヤーのすべてのパラメータデータがキャッシュに収まるよう書き込まれなかったり、そのレイヤーのデータが大きすぎると、そのレイヤーのためのすべてのデータを外部メモリからフェッチする必要があります。

性能検討 (Performance considerations)

各モデルに割り当てられたキャッシュは従来のキャッシュではなく、コンパイラが割り当てたスクラッチパッドメモリであることを覚えておくことは重要です。

Edge TPU コンパイラは、Edge TPU RAM のサイズを知っており、各モデルの実行とパラメータデータにどれくらいメモリが必要になるかを知っています。そのため、コンパイラはコンパイル時に各モデルのパラメータデータのために固定されたキャッシュ空間を割り当てることができます。edgetpu_compier コマンドは、この情報を与えられた各モデルについて出力します。例えば以下のとおりです。

On-chip memory available for caching model parameters: 6.91MiB
On-chip memory used for caching model parameters: 4.21MiB
Off-chip memory used for streaming uncached model parameters: 0.00B

この場合、モデルのパラメータデータはすべて Edge TPU RAM に収まっており「Off-chip memory used」が 0 となっています。

しかし、２つのモデルを同時コンパイルし、最初のモデルが 6.91 MiB の RAM のうち 4.21 MiB を使用し、空き容量が 2.7 MiB になったとします。すべてのパラメータデータのための十分な空きがない場合、残りは外部メモリからフェッチする必要があります。このような場合、コンパイラは２つ目のモデルのために以下のような情報を出力します。

On-chip memory available for caching model parameters: 2.7MiB
On-chip memory used for caching model parameters: 2.49MiB
Off-chip memory used for streaming uncached model parameters: 4.25MiB

２つ目のモデルの「Off-chip memory used」が 4.25 MiB であることに注意してください。このシナリオを図示します。

※ Figure 2. Two co-compiled models that cannot both fit all parameter data on the Edge TPU RAM を元に作成

あなたのアプリケーションがモデル B だけを実行する場合でも、Edge TPU RAM にはそのデータの一部しか格納されません。それは、モデル A という他のモデルと同時コンパイルして、利用可能な容量がそのように決定されたからです。

静的設計の主な利点は、モデルを同時コンパイルしたときに性能が決定的であることと、RAM の頻繁な書き換えに時間がかからないことです。そしてもちろん、モデルが Edge TPU RAM にすべて収まれば、外部メモリからの読み取りや Edge TPU RAM の書き換えなしに、最大の性能を達成することができます。

同時コンパイルを使うかどうかを決めるときには、すべてのモデルをコンパイルして、Edge TPU RAM にすべてのパラメータデータが収められるかどうかを（コンパイラの出力を見て）確認するべきです。収まりきらない場合、各モデルの使用頻度を検討してください。最もよく使われるであろうモデルをコンパイラに最初に指定すれば、そのモデルのすべてのパラメータデータがキャッシュされるようになります。すべてが収まりきらずかつ切り替えが稀な場合、Edge TPU RAM を定期的に書き換えるよりも外部メモリからの読み取りの方がコストがかかるため、同時コンパイルは有益ではありません。ベストを決めるには、さまざまなコンパイルオプションでテストする必要があります。

Tip: 複数のモデルのデータがすべてキャッシュに収まらない場合、edgetpu_compiler に指定するモデルの順序を変えてみてください。前述したとおり、データは一度に一つのレイヤーが割り当てられます。より小さいレイヤーをより多くキャッシュできるような順序が見つかるかもしれません。

警告: 複数の Edge TPU で同時コンパイルを使うには注意が必要です。モデルを同時コンパイルしても、実際に別々の Edge TPU でモデルが実行される場合、モデルが不必要に外部メモリにパラメータデータを格納する可能性があります。そのため、同時コンパイルされたモデルは同じ Edge TPU で実際に実行されることを確認する必要があります。

コンパイラとランタイムのバージョン (Compiler and runtime versions)

Edge TPU 用にコンパイルされたモデルは、対応するバージョンの Edge TPU ランタイムを使って実行する必要があります。最近コンパイルされたモデルを古いランタイムで実行しようとすると、以下のようなエラーが発生します。

Failed precondition: Package requires runtime version (12), which is newer than this runtime version (10).

これを解決するには、デバイスのランタイムのバージョンを更新してください。

• Dev Board 向けの手順
• USB アクセラレータ向けの手順

デバイスのランタイムを更新できない場合、代わりにモデルを再コンパイルし、古いランタイムバージョンと互換性を持たせるようにしてください。--min_runtime_version フラグを付けて edgetpu_compiler を実行すると、そうできます。

edgetpu_compiler --min_runtime_version 10 your_model.tflite

以下の表は、Edge TPU コンパイラのバージョンと、コンパイラが既定で要求する Edge TPU ランタイムのバージョンの対応を示しています。前述したとおり、新しいコンパイラを使用すれば、常に古いランタイムと互換性のあるモデルを作成できます。

コンパイラバージョン	コンパイラが既定で要求するランタイムバージョン
2.0	12
1.0	10

コンパイラのバージョンは以下のようにして確認できます。

edgetpu_compiler --version

ランタイムのバージョンは以下のようにして確認できます。

python3 -c "import edgetpu.basic.edgetpu_utils; print(edgetpu.basic.edgetpu_utils.GetRuntimeVersion())"

ヘルプ (Help)

モデルをコンパイルしたときに「Mapped to the CPU」というログが多く表示された場合、モデル要件を慎重に確認し、Edge TPU にマップされる操作を増やすための変更を試みてください。

コンパイラがモデルのコンパイルに完全に失敗し、一般的なエラーメッセージが表示される場合には連絡して問題を報告してください。その際、問題をデバッグできるようにコンパイルしようとしている TensorFlow Lite モデルを共有してください。(完全に訓練されたモデルは不要です。ランダムに初期化されたパラメータがあれば問題ありません。)

考察

Models | Coral にある EfficientNet を Jetson Nano に USB アクセラレータ版の Edge TPU を接続して速度を測定してみると・・・

• EfficientNet-EdgeTpu (L) (13.3 MiB)
  • 初回推論: 50 ミリ秒前後
  • 次回以降: 40 ミリ秒前後
• EfficientNet-EdgeTpu (M) (9 MiB)
  • 初回推論: 35 ミリ秒前後
  • 次回以降: 17 ミリ秒前後
• EfficientNet-EdgeTpu (S) (7 MiB)
  • 初回推論: 30 ミリ秒前後
  • 次回以降: 10 ミリ秒前後

どれも Edge TPU に用意されているパラメータデータ用の SRAM 的にはつらいものがありますが、初回推論の性能でみても EfficientNet-EdgeTpu (L) で 20 fps 出るのはなかなか、ではないでしょうか。幸い Jetson Nano は USB 3.0 なので、Edge TPU への転送速度が速いのかもしれません。USB 2.0 だともう少し性能が落ちると思われます。

MobileNet-v2 SSD (COCO) のモデルは、容量が 7MiB ほどあるので、他のモデルと一緒に１つの Edge TPU で実行するのは非常に厳しいです。

PoseNet のモデルは、高解像度向けでも容量が 2.5 MiB ほどなので、同時コンパイルすれば他のモデルと同時に高速推論できる可能性があります。PoseNet は解像度で大・中・小の三種類のモデルがあり、Jetson Nano + USB アクセラレータ版ではそれぞれ 25 fps、60 fps、100 fps 程度の性能が出ます。とはいえ、PoseNet の Edge TPU 用のモデルは作成方法が不明なので、現状では同時コンパイルができません。同時コンパイルには、Edge TPU 用にコンパイルする前の TensorFlow Lite モデルが必要だからです。ちなみに、Edge TPU 用にコンパイルされたモデルを edgetpu_compiler でコンパイルしようとすると、以下のようなエラーが出ます。

$ edgetpu_compiler posenet_mobilenet_v1_075_481_641_quant_decoder_edgetpu.tflite 
Edge TPU Compiler version 2.0.258810407
INFO: Initialized TensorFlow Lite runtime.
Invalid model: posenet_mobilenet_v1_075_481_641_quant_decoder_edgetpu.tflite
Model already compiled for Edge TPU

まとめ

Edge TPU はものすごいコストパフォーマンスを持ったデバイスだと思いますが、その性能を引き出すには Edge TPU の特性をよく理解する必要があります。特にモデルのサイズに起因することが多いということをなんとなく感じることができたのではないでしょうか。

Edge TPU では、NVIDIA GPU でよく実行されるような 100 MiB 超えのモデルの推論は厳しいです。FP32 なモデルなら INT8 に量子化して容量が 1/4 程度になるとはいえ、それでも 16 MiB 程度です。なので、最初から Edge TPU で動作させることを前提にモデルを用意していく必要があります。当然といえば当然ですが・・・。

Edge TPU の性能を引き出すにはモデルを 7 MiB より小さくできるとよさそうですが、USB アクセラレータ版を使う場合でも USB 3.0 で接続できれば、思ったより速く動くな、という印象です。この辺は、もうすぐリリースされそうな Mini PCIe / M.2 Accelerator ではもう少し変化すると思われます。USB 3.0 の最大転送速度は 5 Gbps、Mini PCIe or M.2 (A+E or B+M) だと・・・どうなるんでしょうね。

PoseNet を Edge TPU 用にコンパイルする公式手順が公開がされるとありがたいですね。ちなみに、rwightman/posenet-pytorch では TensorFlow.js の PoseNet を PyTorch の世界に引きずり込んだりしているので、不可能ではないかもしれません。

複数のモデルを動かす戦略はいくつかありますが、最もシンプルな戦略は１つの Edge TPU で１つのモデルを実行する、という戦略です。モデルの数を M、Edge TPU の台数を N としたとき、M > N となるならば、同時コンパイルする戦略を検討することになると思います。同時コンパイルする戦略ではさまざまな組み合わせが考えられますが、何をどこまでやる必要があるのかを検討する、問題を単純化して解きやすくする、などをやった後の最後の手段かな、という印象です。ちなみに、複数 Edge TPU、複数モデルの実装例は以下が参考になります。

• C++: ２つの Edge TPU でモデルをひとつずつ実行する例
• Python: １つの Edge TPU で２つのモデル、または２つの Edge TPU でモデルをひとつずつ実行する例

オプティムでは、こうした技術に興味がある・作ってみたい・既に作っている、というエンジニアを募集しています。興味のある方は、こちらをご覧ください。

www.optim.co.jp

• 参考資料
  1. . Edge TPU USB Acceleratorの解析 - Operationとモデル構造 - Qiita
  2. . M.2 - Wikipedia