Edge TPUで(物理的に)動かす独自モデル

NTTドコモの北出です。

普段はロボティクス、画像認識からハードウェアまでHCI全般を扱っています。
こちらはNTTドコモサービスイノベーション部 AdventCalender2019 25日目の記事です。
「北島三郎さんの北、にちなんで」という謎の理由でカレンダーのトリを任されました汗

はじめに

Googleから発売されているEdge TPUが搭載されたキット「Coral Dev Board」で独自に学習したモデルを使って動くものを作ります。この辺りのまとまった日本語情報がなかったので、

• セットアップ
• 学習
• モデルの変換
• Dev boardへの実装

まで少し丁寧に書いてみました。モデルを作るところが少し厚めです。

一通りお読みいただくと、Coral Dev Board上で物体認識を行い、物理的に動く何かが作れるようになると思います。
特に積みボードが溜まっている方、お正月の工作にいかがでしょうか笑

Coral Dev Boardとは？

画像認識などのDeeplearning推論処理を高速にこなしてくれる専用チップ「TPU」が搭載されたシングルボードコンピュータです。TPU自体は色々なGoogle製品に搭載されるようになってきており、最近ではPixel4にも搭載されているそうです。

Dev Boardについては日本でも技適が通っており、2019年の8月頃から2万円ほどで購入できるようになってきました(例：SWITCHSCIENCE)。Dev Boardの細かい仕様はこちら。
サイズやGPIOのピンアサインはRaspberry piシリーズに準拠しており、組み込み機器にも使いやすいようなフォームファクターです。最近では似たようなボードとしてASUSから「Tinker Edge T」も発表されていますね。
下図の左がCoral Dev board、右がRaspberry Pi3(比較用)です。

Edge TPUについては、公式ページに詳しい説明があります。私もこれまでNVIDIAのJetsonシリーズ、IntelのNeural computing stick、SipeedのMAiXシリーズ(K210)などに触れてきましたが、TensorflowやMobilenetなどGoogleの研究成果を活用するには、EdgeTPUは良い選択肢になるのではないでしょうか。とりわけ、コンパクトにDeeplearningを加速させるシングルボードコンピュータとして、Coral Dev Boardは良いと思います。

セットアップ

早速、公式のドキュメントに従ってCoral Dev Boardをセットアップしていきます。
OSをフラッシュするための母艦PCとしてはLinuxの他Macが対応していますので、今回はMacを選択しました。
Dev boardにソフトウェアをインストールするためにscreenコマンドのインストールが案内されていますが、Macでは最初から使えるようになっています。

fastbootのインストール

OSイメージのフラッシュを行うため、fastbootコマンドを使えるようにします。AndroidStudioをインストールしているPCであればすでに導入されていると思いますが、導入されていない場合はAndroidの開発者サイトおよびDev Board公式の手順を参考にインストールします。
Dev boardの手順では、platform-toolsの中からfastbootコマンドのみを取り出して配置する案内がされています。

Mendel Development Tool(MDT)のインストール

Dev boardにはMendel LinuxというOSをインストールしますが、開発ツールとしてMDTをインストールします。

pip3 install --user mendel-development-tool

OSイメージのフラッシュ

まずは、フラッシュを行う前にボードのDIPスイッチの確認を行います。Switch1から順に、ON-OFF-OFF-OFFとなっていればOKです。
※2019年4月10日以前に製造されたボードでは、Serial consoleを使ってfastbootを有効にする必要があるようです。その場合は、こちらの手順に従います。

USB-OTGにMacとの接続ケーブルを、USB-PWRに電源を接続すると、ファンが回りだしLEDが光ります。(最初、電源繋いでもうんともすんとも言わず焦りましたが、全てのケーブルを抜いてUSB-PWR、OTGのみにしたところ起動してホッとしました汗)
電源供給にはiPadPro2018に付属していた充電器を使用しました。

接続がうまくいっていると、Macからfastbootコマンドを打ってシリアル番号が表示されることを確認します。エラーが出る場合は、fastbootコマンドが最新になっていることを確認しましょう。
下記のようなログが出てプロンプトが戻ってくると、Dev board側のHDMIから映像が出力されるようになります。

mendel-enterprise-chef-13 $ bash flash.sh
Sending 'bootloader0' (1006 KB)                    OKAY [  0.051s]
Writing 'bootloader0'                              OKAY [  0.234s]
Finished. Total time: 0.305s
～途中省略～
Writing 'rootfs'                                   OKAY [104.724s]
Finished. Total time: 235.270s
Rebooting                                          OKAY [  0.005s]
Finished. Total time: 0.005s
mendel-enterprise-chef-13 $

MDTを使ってDev boardに接続する

mdtコマンドを使って、USBーOTG経由でDev boardにアクセスが可能です。・・・と公式ページに書いてあったのですが、一向に認識されないので、諦めてDev board本体のコンソールから直接操作することにしました。デフォルトではパスワード認証が切られているので、これを有効にしてしまいます。
/etc/ssh/sshd_configの中にあるPasswordAuthenticationをyesにし、sshdを再起動します。
有線LANを接続した上でIPアドレスを調べて接続すると、無事ログインができます。
(ちなみに、無線LANを使った接続はnmtuiコマンドを使って設定できます。)

その後、公式の手順にしたがってソフトウェアアップグレードをかけておきます。

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 dist-upgrade

ここまで終われば、Dev boardとしての準備は完了です。

独自データセット学習からモデルデプロイまで

正直なところ、この部分の情報が少なく試行錯誤しながら進めることになりました。
Edge TPUはTensorflow Liteのモデルにのみ対応しているため、TensorFlowで学習した後に色々な変換処理が入ってきます。

ワークフロー

公式のドキュメントページにも、TensorFlowモデルを学習しEdge TPU用のモデルを作り出すまでのフローが記載されています。
(The basic workflow to create a model for the Edge TPU(公式ページより引用)!

それぞれの手順は自分で何とかしましょう、感が漂っています。Edge TPUのチュートリアルに「Retrain an object detection model」という項目もあったのですが、GPUを使って学習化を高速化させたかったので別の方法を取ります。この記事では、以下にコマンドレベルで色々と記載していきたいと思います。流れとしては次の3つ。

1. TensorFlow Object Detection APIを使った学習

• 入力：画像データセット(Pascal VOCをtfrecord形式に変換したもの)
• 出力：Tensorflow model(saved_modelディレクトリにできる生成物)

1. TensorFlowモデルからTensorFlow liteモデルへの変換

• 入力：TensorFlow model(Frozen graph形式にエクスポートされた.pbファイル)
• 出力：TensorFlow Liteモデル(.tfliteファイル)

1. Tensorflow liteモデルからEdgeTPU向けモデルへの変換

• 入力：TensorFlow Liteモデル(.tfliteファイル)
• 出力：TensorFlow Liteモデル(Edge TPU向けに最適化された.tfliteファイル)

TensorFlowのバージョンは1.14(pip install tensorflow-gpu==1.14)を使いました。
すでにTensorFlow2.0がリリースされておりTensorFlow1.xからのコード変換にも対応されていますが、contribについては変換に対応しておらず、今回学習に使用したObject Detection APIでもエラーが出てしまいます。

TensorFlow Object Detection APIを使った学習

TensorFlowに付属している、Object Detection APIを使って学習を進めて行きます。大枠としては、からあげさんの記事を参考にさせていただきました。リポジトリのCloneからmodel_builder_test.pyの実行までを済ませておきます。
以下、Object Detection APIはmodels/配下に、からあげさんのツールはmodels/research/object_detection_tools/に配置されているものとします。

データセット変換

今回は以前生成したじゃんけん認識用のデータセットを活用するため、まずはPascal VOC形式からTFrecordの形式に変換します。形式の変換についてはこちらの記事の内容をほぼそのまま活用させていただきました。
Pascal VOC形式のサンプル：

<annotation>
  <filename>../Data/0_bg_IMG_1118.jpg</filename>
  <object>
    <name>guu</name>
    <pose>Left</pose>
    <difficult>0</difficult>
    <bndbox>
      <xmin>36</xmin>
      <ymin>18</ymin>
      <xmax>100</xmax>
      <ymax>100</ymax>
    </bndbox>
  </object>
  <object>
    <name>guu</name>
    <pose>Left</pose>
    <difficult>0</difficult>
    <bndbox>
      <xmin>118</xmin>
      <ymin>0</ymin>
      <xmax>200</xmax>
      <ymax>64</ymax>
    </bndbox>
  </object>
〜以下省略〜

TFrecord形式では、シリアライズされた画像データとアノテーションデータをまとめて保存してあります。中身を覗いてみると、こんな感じ。

features {
  feature {
    key: "image/encoded"
    value {
      bytes_list {
〜途中省略〜
  feature {
    key: "image/object/bbox/xmax"
    value {
      float_list {
        value: 0.46000000834465027
        value: 1.0
        value: 0.28999999165534973
        value: 0.9150000214576721
      }
    }
  }
〜途中省略〜
  feature {
    key: "image/object/bbox/ymin"
    value {
      float_list {
        value: 0.20999999344348907
        value: 0.03999999910593033
        value: 0.5849999785423279
        value: 0.5
      }
    }
  }
〜途中省略〜
  feature {
    key: "image/object/class/text"
    value {
      bytes_list {
        value: "paa"
        value: "paa"
        value: "paa"
        value: "paa"
      }
    }
  }
〜途中省略〜
  feature {
    key: "image/source_id"
    value {
      bytes_list {
        value: "../Data/5_bgpress_IMG_5933.jpg"
      }
    }
  }
〜以下省略〜

Pythonインタプリタから、このようなコマンドで確認することができます。

>> OUTPUT_TFRECORD_NAME = "test.record"
>> example = next(tf.python_io.tf_record_iterator(OUTPUT_TFRECORD_NAME))
>> tf.train.Example.FromString(example)

TFrecordの取扱については、こちらの記事が参考になりました。
TFrecord形式の基本的な構造としては、

• features
  • feature{シリアライズされたJPEGイメージ}
  • feature{xmax(出現順)}
  • feature{xmin(出現順)}
  • ‥
  • feature{クラス名(出現順)}
  • ‥

という感じで、一つのアノテーション(features)の中にアノテーション情報(座標やクラス)が入っています。バウンディングボックスごとにfeatureがあるのではなく、xmax,xmin,,,クラス名がそれぞれのfeatureに順に格納されている(順番に意味がある)ようです。

データの準備としては、models/research/object_detection_tools/data配下に、tf_label_map.pbtxt、train/frame0000.tfrecord、val/frame0000.tfrecordが配置されるようにします。

学習の準備

学習させたいアルゴリズムに応じて、configファイルを編集します。今回は、mobilenet_v2_ssdを選択しました。finetuningに使う重みファイルとconfigは、こちらから取得することができます。pipeline.configを編集し、わかりやすい名前にて保存しておきます。ここでは、からあげさんのツールを使いましたのでmodels/research/object_detection_tools/configにssd_mobilenet_v2_quantized_coco.configとして保存しました。デフォルトから編集すべき内容としては下記のような感じです。

• num_classes
  • 認識させたいクラス数、今回はguu, tyoki, paaの3クラス
• fixed_shape_resizer
  • 入力画像サイズ
  • 学習時の使用メモリに影響します
• batch_size
  • 複数の画像をまとめて読み込ませて学習を進めます
  • 学習時の使用メモリに影響します
    • 参考までに、quantizationを有効にした状態でGeforce Geforce 1050Ti(4GB)では「4」までしかうまく設定できませんでした
• fine_tune_checkpoint
  • ファインチューニングに使う学習済みモデル
  • “./object_detection_tools/models/ssd_mobilenet_v2_quantized_300x300_coco_2019_01_03/model.ckpt”、としました
• label_map_path
  • 認識させるラベル
  • "./object_detection_tools/data/tf_label_map.pbtxt"としました
• train_input_readerの中のinput_path
  • 学習用TFrecordのパス
  • "./object_detection_tools/data/train/frame????.tfrecord"としました
• eval_input_readerの中のinput_path
  • 評価用TFrecordのパス
  • "./object_detection_tools/data/val/frame????.tfrecord"としました
• graph_rewriter
  • Quantization aware trainingにするための設定(有効にすると学習時に消費するメモリが増大します。)
  • 次のようにしました
    graph_rewriter {
    quantization {
    delay: 48000
    weight_bits: 8
    activation_bits: 8
    }
    }

学習の実行

データとconfigの準備が終わったら、いよいよ学習です。次のコマンドで学習を実行しました。

model/research$ python3 object_detection/model_main.py --pipeline_config_path="./object_detection_tools/config/ssd_mobilenet_v2_quantized_coco.config" --model_dir="./saved_model_ssd_mobilenet_v2_quantized_coco" --num_train_steps=10000 --alsologtostderr

model_dirに学習結果が保存されていきます。

学習の様子を可視化する

tensorboardを使って学習が進んでいるかを逐一確認することができます。

model/research$ tensorboard --logdir="saved_model_ssdlite_mobilenet_v2_quantized_coco"

http://localhost:6006/ にアクセスすると、tensorboardの内容を確認することができます。学習が進んでいるかの目安としてはmAPが上がっているか・lossが下がっているか、があります。

mAPのすべての項目が-1で固まっている場合、何かがおかしいですので学習の設定を確認してみましょう。自分がやってしまった間違えとしては、

• tf_label_map.pbtxtに指定するクラス名が間違っている
• TFrecordに取り込むべきデータがおかしい

みたいなものがありました。

「IMAGES」を選択すると、学習中に行われる評価の様子も確認できます。(オーグメンテーションの都合で気味の悪い画像が表示されています笑)各画像右側に正解データ、左側に推論結果が表示されています。スライダを動かすと、各世代ごとの様子が確認できます。

TensorFlowモデルからTensorFlow liteモデルへの変換

まずは、学習済みデータディレクトリからグラフをエクスポートします。

models/research$ python3 object_detection/export_tflite_ssd_graph.py --pipeline_config_path object_detection_tools/config/ssd_mobilenet_v2_quantized_coco.config --trained_checkpoint_prefix saved_model_ssd_mobilenet_v2_quantized_coco/model.ckpt-10000 --output_directory saved_model_ssd_mobilenet_v2_quantized_coco/exported_graphs --add_postprocessing_op True

object detection APIに同じく含まれるexport_inference_graph.pyを使う方法も試しましたが、イマイチうまくtflite_converterの手順までたどり着けなかったためexport_tflite_ssd_graph.pyの方を使っています。

続いて、TensorFlow Lite形式(.tflite)に変換します。

models/research/saved_model_ssd_mobilenet_v2_quantized_coco$ tflite_convert --output_file ssd_mobilenet_v2_quantized_janken_10000.tflite --graph_def_file exported_graphs/tflite_graph.pb --input_arrays=normalized_input_image_tensor --output_arrays='TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3' --input_shape=1,300,300,3 --allow_custom_ops --inference_input_type=QUANTIZED_UINT8 --mean_values=128 --std_dev_values=127 --default_ranges_min=0 --default_ranges_max=255 --inference_type=QUANTIZED_UINT8

ここで、次の手順で示すedgetpu_compilerで変換できるようにquantizeを行います。元はFLOAT32、quantaize後はUINT8で処理されます。
また、学習時にconfigファイルに設定をしてquantize awareにしておく必要があります。(quantize awareでなくても変換自体はできますが、scoreが0.5に張り付いたりと推論結果がめちゃくちゃになってしまいました。)

Tensorflow liteモデルからEdgeTPU向けモデルへの変換

Edge TPU Compilerを使って、tfliteモデルをEdgeTPU用のtfliteに変換します。
コンパイラ自体はx86_64のUbuntuマシンでも動作します。

$ edgetpu_compiler ssd_mobilenet_v2_quantized_janken_10000.tflite -o .
Edge TPU Compiler version 2.0.267685300

Model compiled successfully in 302 ms.

Input model: ssd_mobilenet_v2_quantized_janken_10000.tflite
Input size: 4.52MiB
Output model: ./ssd_mobilenet_v2_quantized_janken_10000_edgetpu.tflite
Output size: 5.17MiB
On-chip memory available for caching model parameters: 7.62MiB
On-chip memory used for caching model parameters: 5.05MiB
Off-chip memory used for streaming uncached model parameters: 0.00B
Number of Edge TPU subgraphs: 1
Total number of operations: 99
Operation log: ./ssd_mobilenet_v2_quantized_janken_10000_edgetpu.log

Model successfully compiled but not all operations are supported by the Edge TPU. A percentage of the model will instead run on the CPU, which is slower. If possible, consider updating your model to use only operations supported by the Edge TPU. For details, visit g.co/coral/model-reqs.
Number of operations that will run on Edge TPU: 98
Number of operations that will run on CPU: 1
See the operation log file for individual operation details.

上のように特にエラーなく変換が完了すれば成功です。
On-chip memoryと書かれている部分はTPU側で動作する領域、Off-chip memoryはDev board側にキャッシュされる領域になります。Off-chipの方に消費メモリがあると速度の低下が想定されますが、今回は大丈夫そうです。

また、Number of operationsの部分はTPUもしくはCPUで実行されるオペレーション数です。.logファイルで詳細を確認すると下記のような感じでした。

Operator                       Count      Status

LOGISTIC                       1          Mapped to Edge TPU
CUSTOM                         1          Operation is working on an unsupported data type
ADD                            10         Mapped to Edge TPU
CONCATENATION                  2          Mapped to Edge TPU
CONV_2D                        55         Mapped to Edge TPU
DEPTHWISE_CONV_2D              17         Mapped to Edge TPU
RESHAPE                        13         Mapped to Edge TPU

CUSTOMってなんだろう。。とりあえず今回はCONV処理がTPUにすべて乗ったので良しとします。(サポートされているレイヤ一覧はこちら。)

念の為TensorFlow Liteに付属するツールvisualize.pyを使って可視化してみると、最終出力の直前にある層がCUSTOMという名前に当たるようです。TFLite_Detection_PostProcess、すなわちバウンディングボックス及びクラス、スコア、検出数を出力する部分になりますので、TPUとCPUを往復することもなく、パフォーマンスには大きく影響しないことがわかりました。

変換したモデルをEdgeTPUで動かす(そしてGPIOへ)

こちらのリポジトリにあるサンプルを改造して使います。OpenCVのサンプルを利用したかったのですが、デフォルトではapt-getもpipもビルド済みのOpenCVを引っ張ってくることができませんでした。セルフビルドは面倒だったので、Pygameを使ったサンプルをベースに進めます。UVC対応のWebカメラも普通に使えます。(/dev/video1と認識されます。)
サンプルを実行するためには、TensorFlow Lite interpreterが必要なので忘れずにインストールしておしておいてください。

今回は画像認識結果を使ってサーボモータを駆動させますが、Coral Dev BoardのPWM出力のレベルが低い(3.3Vにも満たず、実測2.5Vくらい)ので少し使いづらいため、PCA9685を経由してI2Cで制御することにしました。なお、Coral Dev BoardのGPIOに関する仕様はこちらにあります。
PCA9685のボードとして、Adafruit 16-Channel 12-bit PWM/Servo Driver(の互換品)を使いますので、必要なライブラリをインストールしておきます。

I2C経由でのサーボモータ制御は、基本的にはAdafruit_Python_PCA9685にありますexamples/simpletest.pyを参考にすればよいのですが、そのまま使うとデフォルトのI2Cバスが見つからずエラーになります。PWMドライバ初期化時に下記のようにアドレスとバスを明示してやります。

# Initialise the PCA9685 using the default address (0x40).                          
pwm = Adafruit_PCA9685.PCA9685(address=0x40, busnum=1)

もしもデバイスがビジーになってしまう場合は、Dev boardの電源を入れてからPCA9685ボードを接続するのが良いと思われます。

肝心な画像認識部分ですが、サンプル(examples-camera/pygame/detect.py)に上で学習したオリジナルモデルを読み込ませるだけですので、細かい内容は省略します。

python3 detect.py --model ~/workspace/models/ssd_mobilenet_v2_quantized_janken_10000_edgetpu.tflite --labels ~/workspace/models/label.txt

推論時間は7ms前後です。かなり早いですね。

この記事の集大成としてdetect.pyをちょっとだけ改造しまして、じゃんけんに負けるとピコピコハンマーが飛んでくるようにしました。

…ちょっとシュールですね笑
1年間の煩悩が吹き飛ぶと良いですね。
仕組みは簡単で、画像認識を使ってじゃんけんに負けるとサーボモーターが駆動し、ピコピコハンマーが解き放たれる仕掛けです。
ちょっと気を取り直して、リプレイを見てみましょう。

動画を29.97fpsで撮影したのでそれ以上の分析はできないのですが、

• じゃんけんの手を出したフレーム(0ms)
• サーボモーターが動き始めたフレーム(167ms)
• 画像認識結果が表示されたフレーム(267ms)
• ハンマーにやられるフレーム(567ms)

の順にフレームを抜き出しています。
…あれ？画像認識結果が表示される前にサーボモーターが動き始めていますね。これは、表示遅延によるもののようです。

EdgeTPUの恩恵に注目すると、 エンドトゥーエンドで167ms以内にシステムが反応できている ことになりますね。これはかなり高速ですね。
※エンドトゥーエンドとは、USBカメラキャプチャ→EdgeTPUによる推論→GPIO(I2C)信号発出→サーボモータ動作、までを指します。

USBカメラによる遅延もそれなりにあると考えられますので、この部分をMIPIなどを経由した高速なものに取り換えるとさらに応答速度が向上しそうです。

最後に、推論システム全体の消費電力を見てみましょう。図の左側がアイドル状態、右側が推論状態です。

計算すると、 推論時でも7W以下で動作している ことがわかります。モバイルバッテリーで十分駆動できそうですね。
※測定値は、EdgeTPU本体およびUSB周辺機器(カメラ、キーボード・マウス)を含んだ値です。サーボモータへの電力は含みません。

終わりに

いかがでしたでしょうか。ちょっと内容を盛り込みすぎた感じもしましたが、特にEdge TPU内部の動作や量子化など機会があればまた記載したいと思います。

それでは、皆様良いお年を！
(あれ？まだ何かある？？)