Coral USB Accelerator について

Coral USB Acceleratorは、Googleが開発していたテンソル・プロセッシング・ユニット(Tensor processing unit、TPU)をUSB接続で簡単に使えるようにしたもので、ディープラーニング系の機械学習モデルの推論のみを高速に行えるデバイスだ。

USBに挿すだけで非力なエッジデバイスでも高度な機械学習のモデルで推論が行えるとの謳い文句になっている。

今回は、自分が用意した動画に写っている物体を学習させ、このデバイスを用いて検知する。

動作環境

  • Ubuntu 16.04 のノートパソコン(Thinkpad X1 carbon)
  • Intel® Core™ i7-8550U CPU @ 1.80GHz × 8
  • メモリ16GB
  • GPUなし

Coral USB Acceleratorのセットアップ

https://coral.withgoogle.com/docs/accelerator/get-started/ を参考にセットアップを行う。

$ echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list
$ curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
$ sudo apt-get update
$ sudo apt-get install libedgetpu1-max
$ # sudo apt-get install libedgetpu1-std # 温度が気になるときは、最大周波数以下で動かすこっちにする

推論を行うには、TensorFlow Liteのランタイムが必要になる。クイックスタートガイドを参考に、pipでインストールする。また、結局あとでサンプル動かすときにTensorFlowも必要になるので、それもインストールする。Python3.6が使いたかったので、virtualenvで仮想環境を作った。Python3.6のインストールは、ここを参照。

$ # 仮想環境の作成 (グローバルで動かす場合は不要)
$ sudo apt install python-virtualenv
$ cd ~ && mkdir -p edgetpu && cd edgetpu
$ virtualenv -p python3.6 venv-coral # venv-coralという名前の仮想環境を作る
$ source venv-coral/bin/activate

TensorFlow Lightだが、ここから自分のアーキテクチャーとPythonのバージョンによって違うWheelをダウンロードする必要がある。

$ pip install tflite_runtime-<version>-<archtechture>.whl
$ pip install tensorflow

サンプルプログラムを動かしてチェック

exampleをgithubからクローンしてきて動かす

$ cd ~/edgetpu && mkdir coral && cd coral
$ git clone https://github.com/google-coral/tflite.git
$ cd tflite/python/examples/classification
$ bash install_requirements.sh
$ 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

これで、鳥の種類を判別するモデルで推論が行われるはずだ。エラーなく動いたら先に進む。

自分で用意した動画に出てくる物体を学習させて検知する

Coral USB TPUのサンプルには、物体のクラス分けと物体の検出が用意されている。

特に、公式チュートリアルにも用意されているように、物体の検出では、高速に検出ができるMobileNet-SSDを簡単に使うことができるようになっている。

ここでは、チュートリアルに従ってCOCOデータセットで90種類の物体を検出できる学習済みMobileNet V1 SSDを転移学習させることで、自分が判別したいデータセットに対応させる。

転移学習とは?

従来、物体識別モデルの学習には数日かかっていたが、転移学習の技術では新しいモデルを作る際にすでに関連する他のタスクのための学習済みモデルを始点として使うことで、おおむね1時間以内に学習を終えることができる。

Coral Tutorial (物体検出モデルの再学習) より和訳

転移学習には、最下層のみを再学習させる方法と、すべての層を再学習させる方法の2種類があるが、ここでは高速に学習を終えたいので最下層のみを再学習させることにする。

データセットの用意

まずは、検出したい物体のデータセットを用意する必要がある。

アノテーションを自分で行う際には、便利なツールがいくつかあるようだ。

LabelImgの評判もよさそうだったが、Microsoftが作っている動画をそのままアノテーションできることが売りのクールなVoTTを使うことにした。

VoTTの使い方は、こちらを参考にした。

ここでは、1分間の動画をそのまま読み込み、1秒5フレームずつアノテーションを行った。

Export Settingで、TensorFlow形式を指定しておくと、画像とアノテーション情報はアノテーションした1フレームずつ.tfrecord形式で、ラベルは.pbtxt形式でエクスポートしてくれる。後処理や変換の必要なく、この後の学習に仕えてすごく便利。

ここで、データセットをTrainとValidation用に分割する。split-folders という便利なPythonモジュールがあるので、pipでインストールして使う。このモジュールだが、Python 3.6と3.7でしか動かないので注意。

$ pip install split-folders tqdm
$ split_folders <path_to_your_tfrecord_exported_directory> --ratio .8 .2

のちの再学習で使う画素数は300×300(多分設定で変えることは可能)だが、大きさ違ってても学習時にリサイズしてくれるみたいなので、画素数は適当でよい。ただ、要求精度や用途によってちゃんと考えるべきではある。

再学習を行うためのDockerイメージの準備

再学習に必要な環境とモデルが用意されたDockerfileが配布されているので、ありがたく使う。(Docker環境はあらかじめ構築しておくこと)

$ DETECT_DIR=${HOME}/edgetpu/detection && mkdir -p $DETECT_DIR
$ cd $DETECT_DIR
$ wget -O Dockerfile "https://storage.googleapis.com/cloud-iot-edge-pretrained-models/docker/obj_det_docker"
$ sudo docker build - < Dockerfile --tag detect-tutorial

$ docker run --name edgetpu-detect \
--rm -it --privileged -p 6006:6006 \
--mount type=bind,src=${DETECT_DIR},dst=/tensorflow/models/research/learn_pet \
detect-tutorial

これでTensorFlowとか必要なもの一式が入ったDockerコンテナ内に入るので、そこで学習済みモデルをダウンロードして再学習を行う。

MobileNet SSDの再学習

学習済みモデルのダウンロード、MobileNet V1 SSDを使う。V2を使いたいときは、network_typeをmobilenet_v2_ssdに変えればよい。また、行末のフラグで転移学習の際に最下層のみか全体かを切り替えられる模様。

# From the Docker /tensorflow/models/research/ directory
./prepare_checkpoint_and_dataset.sh --network_type mobilenet_v1_ssd --train_whole_model false

学習ステップ数と検証ステップ数の設定を行う。最下層のみの再学習の場合、500ステップでよいらしい。

$ # If you're retraining just the last few layers, we suggest the following numbers:
$ NUM_TRAINING_STEPS=500 && NUM_EVAL_STEPS=100

$ # If you're retraining the whole-model, we suggest the following numbers:
$ # NUM_TRAINING_STEPS=50000 && NUM_EVAL_STEPS=2000

本家のチュートリアルでは、先のprepare_checkpoint_and_dataset.shで、学習済みモデルとデータセット(ペットの犬と猫の種類)と再学習の設定ファイルを、一括でダウンロードと後処理を行っている。ダウンロード先はローカルにバインドマウントされているはずなので、そのディレクトリにVoTTで用意したデータセットをコピーし、設定ファイルをいじる。

いじるべき設定ファイルは、(このページの通りなら)~/edgetpu/detection/ckpt/pipeline.configにある。必須な変更箇所は2つ(詳細はこちら)。

  • ファイル先頭のnum_classesをアノテーションしたクラス数に適切に変更
  • train_input_reader と eval_input_readerのパスを適切に設定

再学習は、スクリプト一発でいける。

# From the /tensorflow/models/research/ directory
./retrain_detection_model.sh \
--num_training_steps ${NUM_TRAINING_STEPS} \
--num_eval_steps ${NUM_EVAL_STEPS}

再学習の経過は、別途tensorboardで監視できる。

(host system) $ sudo docker exec -it edgetpu-detect /bin/bash
(inside container) # tensorboard --logdir=./learn_pet/train/

http://localhost:6006/ で経過のグラフや検証状況が確認できる。

今回、300枚程度のデータを用意したが、GPUなしのノートパソコンで3時間で再学習がおわった。今後は、GPUも使えるようにすればもっと速くできそう。

Coral USB Accelerator (Edge TPU)用にコンパイル

学習が終わったら、Edge TPU用にモデルをコンパイルする。

モデルの抽出は再学習に使ったDockerコンテナ内で行う。適宜checkpoint_numを調節して一番学習が進んだステップを選ぶ。

# From the Docker /tensorflow/models/research directory
$ ./convert_checkpoint_to_edgetpu_tflite.sh --checkpoint_num 500

コンパイルはホストシステム(Dockerコンテナ外)で行う。

# Install the compiler:
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 update

sudo apt install edgetpu

# Change directories to where the new model is:
cd $HOME/edgetpu/detection/models

# Compile the model:
edgetpu_compiler output_tflite_graph.tflite

実行!

適宜検証用の画像を用意して、以下を実行。画素数は相変わらず適当でもいけちゃう。

# サンプルプログラムで試す
$ cd ~/edgetpu/coral/tflight/python/examples/detection
$ python3 detect_image.py \
--model ~/edgetpu/detection/models/output_tflite_graph.tflite
--labels <path_to_your_tf_label_map.pbtxt>
--input <path_to_your_input_image>

これで、検出時間と検出位置が出力されれば、成功!

自分の環境では、だいたい100msくらいで検出ができていた。

学習済みのモデルや重みは、.tfliteファイル1つに収まっているので、これをもっていけば、どこでも自分用にカスタマイズされた物体検出ができる(はず)。

まとめ

思ったより簡単に自分用のモデルを動かすことができて驚いた。コードを一行も書かずにここまでできるとは思っていなかった。

GPUなしのノートパソコンでも、再学習を数時間で終えられて、10Hz以上の推論実行スピードが出る時代になったようだ。

Googleおそるべし。今後、非力な組み込みエッジ分野で機械学習を使う動きがより加速されそうな予感がする。

コメントを残す

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です

このサイトはスパムを低減するために Akismet を使っています。コメントデータの処理方法の詳細はこちらをご覧ください