Edge TPU を試す

　Beeは組込み会社では知見の少ないAIにも取り組んでおり、組込みソフト開発専門の知見を活かし組込みAI(エッジAI)にも積極的に取り組んでいます。エッジデバイスは性能が低いため、推論に時間がかかってしまう欠点がありますが、アクセラレータを使用することでエッジデバイスでもリアルタイムに推論することができます。
　今回は、Google から発売された USB Accelerator を使用してモデルの実行を試してみます。

　Edge TPU はエッジデバイス向け機械学習/ディープラーニングのフレームワークTensorFlow Lite で作成したモデルの推論部分を高速化するためのコプロセッサです。 Edge TPUが搭載されたデバイスとしては、開発ボードとしてはGoogle Coral Dev Board、USB AcceleratorのColal USB Acceleratorが発売されています。

USB Accelerator

高速 ML 推論
オンボードの Edge TPU プロセッサは、毎秒0.5ワットを使用して、1秒間に4兆回の操作を実行できるそうです。
以下は、CPU と Edge TPU の性能を比較したベンチマークです。
ベンチマーク: https://coral.withgoogle.com/docs/edgetpu/benchmarks/
Debian Linux で動作
付属の USB 3.0 Type Cケーブルで、Debian ベースの Linux システムに接続。
TensorFlow Lite をサポート
TensorFlow Lite モデルを Edge TPU 上で動作するようにコンパイル可能。

Dev board

Edge TPUコプロセッサを搭載したシングルボードコンピュータ。

環境

Raspberry Pi 3 Model B
OS:Raspbian Stretch
現在すでにRaspbianの新バージョンのBusterがDL出来ますが、まだ動作が不安定な所や、ライブラリ類が追いついていないこともあり、今回はStretchを利用しています。過去バージョンは以下から取得しました。
過去バージョン : http://ftp.jaist.ac.jp/pub/raspberrypi/raspbian/images/
Google EdgeTPU USB Accelerator

対応しているデバイス

USB ポート付きの Linux マシンに対応
- Debian 6.0 以上（Ubuntu 10.0+）
- x86-64 or ARM32/64 with ARMv8

公式にサポートしているデバイスは、Raspberry Pi 2/3 Model B/B+ です。
非公式ですが Raspberry Pi Zero でも利用可能です。

デバイス設定

Edge TPU ランタイムをインストールします。


cd ~/
wget https://dl.google.com/coral/edgetpu_api/edgetpu_api_latest.tar.gz -0 edgetpu_api.tar.gz --trust-server-names
tar xzf edgetpu_api.tar.gz
cd edgetpu_api
bash ./install.sh

Edge TPU は Python3.5対応のため Python3.6以上だとデモが動かないです。
Raspbian Stretch は標準で Python3.5 のため Stretch にしておくとすぐにデモ試せます。
参考: https://qiita.com/iwatake2222/items/e8e29babc7e60d983093

付属のUSB3.0ケーブルを使用してアクセラレータを接続します。（すでに接続している場合は、再接続）

Edge TPU 上でモデルを実行する

サンプルの鳥分類モデルをダウンロードします。


cd ~/Downloads/
wget https://dl.google.com/coral/canned_models/mobilenet_v2_1.0_224_inat_bird_quant_edgetpu.tflite \
https://dl.google.com/coral/canned_models/inat_bird_labels.txt \
https://coral.withgoogle.com/static/docs/images/parrot.jpg

オウムの画像を使用して画像分類を実行します。


# on Debian/Ubuntu Linux:
cd /usr/local/lib/python3.6/dist-packages/edgetpu/demo
# on Raspberry Pi :
cd /usr/local/lib/python3.5/dist-packages/edgetpu/demo
python3 classify_image.py \
--model ~/Downloads/mobilenet_v2_1.0_224_inat_bird_quant_edgetpu.tflite \
--label ~/Downloads/inat_bird_labels.txt \
--image ~/Downloads/parrot.jpg

画像分類モデルを試してみる

約11msでリアルタイムに分類しています。

環境

Jetson Xavier
Google EdgeTPU USB Accelerator

手順

edgetpu のダウンロード、インストールまではデモと同じです。
Ubuntu の場合は Python3.6 上で edgetpu は動作します。 # on Debian/Ubuntu Linux: cd /usr/local/lib/python3.6/dist-packages/edgetpu/demo
EdgeTPU実行時にエラー ModuleNotFoundError: No module named '_edgetpu_cpp_wrapper' が出る場合。
- install.shの実行時に_edgetpu_cpp_wrapper.so のファイル名を設定してくれないことがあります。
- 手動でファイル名を変更することで解決できます。
```
cd /usr/local/lib/python3.6/dist-packages/edgetpu/swig/
sudo cp _edgetpu_cpp_wrapper.cpython-35m-aarch64-linux-gnu.so _edgetpu_cpp_wrapper.so
											
```

課題

RuntimeError: Failed to allocate tensors でデモ実行できず。

手順

　今回は、Quantization-aware training ※1 をしていないモデルで試してみます。

※1 Quantization-aware training
　量子化を意識したトレーニング。モデルへの入力、重みパラメータを uint8 にします。EdgeTPU の入力が uint8 のため、Quantization-aware training をしたほうがいいです。
　今回は学習済モデルの重みパラメータを uint8 に変換したモデル(Post-training quantization)を使用します。

学習済モデルの準備

下記モデルのどれかを用意。

SavedModel directories
Frozen GraphDef (models generated by freeze_graph.py)
Keras HDF5 models
Models taken from a tf.Session

tfliteモデル作成

tflite_convertコマンドでモデルを tflite モデルに変換します。
Keras HDF5 models を変換する場合の例：


tflite_convert ^
  --output_file=towei_tpu_test.tflite ^
  --keras_model_file=result/towei_out_cnl_WH9_tDA9.h5 ^
  --inference_type=QUANTIZED_UINT8 ^
  --input_shapes=1,94,110,3:1,5 ^
  --input_arrays=input_1,input_2
  --output_arrays=l4/BiasAdd ^
  --default_ranges_min=0.0 ^
  --default_ranges_max=9310.0
  --mean_values=0,0 ^
  --std_dev_values=255,0.027389903329752954

EdgeTPU 用モデル作成

作成した tflite モデルを EdgeTPU 用に変換します。
変換には Google が提供している Web コンパイラを使用します。
https://coral.withgoogle.com/web-compiler/

実行

EdgeTPU API で推論コードを実装、実行します。


#_*_ coding: utf-8 _*_
import numpy as np
import time
import sys,os
import pickle

from edgetpu.basic.basic_engine import BasicEngine

sys.path.append('/home/nvidia/AI/chainer/')

# private library
import lib_data_handle as dh

class MODEL_DATA(object):
	MODEL_FILE = 'result/towei_out_cnl_WH9_t7.tflite'
	TARGET = 'WH9'

def load_test_case():
	with open('test.binaryfile', 'rb') as web:
		l = pickle.load(web)
		xtest, ytest, ztest, ttest = dh.get_data(l, mode='amptime', dtype='CWH', max_weight_lim=100000.)
		# 入力層の形式に合わせる
		xtest = np.transpose(xtest, (0,3,2,1))
		return xtest, ytest, ztest, ttest

def edgetpu_infer():
	# load test data
	x, y, z, t = load_test_case()

	# Initialize engine.
	engine = BasicEngine(MODEL_DATA.MODEL_FILE)
	print('Initialize engine.')

	 # Run inference.
	array_size = engine.required_input_array_size()

	for i in range(x.shape[0]):
		print('\n------------------------------')
		print('Run infrerence.')
		print('  required_input_array_size: ', array_size)
		print('  input shape: ', x.shape)

		input = x[i].flatten()
		input_uint8 = input.astype(np.uint8)

		result = engine.RunInference(input_uint8)

		print('------------------------------')
		print('Result.')
		print('Inference latency: ', result[0], ' milliseconds')
		print('Output tensor: ', result[1])
		print('Inference: ', np.argmax(result[1]))
		print('Label    : ', y[i])

args = sys.argv
if __name__ == "__main__":
	MODEL_DATA.MODEL_FILE = args[1]
	edgetpu_infer()

課題

EdgeTPU の推論は1入力のみ？
- 独自モデルを EdgeTPU で推論させるには、edgetpu.basic.basic_engine を使用します。
- 推論用の API RunInference() の第1引数に入力データ渡します。この引数が uint8 の 1次元配列のみ受け付けています。
推論結果が常に0
- 2入力のモデルに対して1入力のみ与えたせい？原因不明。

　Google Edge TPU USB Accelerator を使ったデモを試してみました。デモはPythonのバージョンを気をつけていれば難なく動かせることができました。画像分類モデルのデモでは性能の低いエッジデバイスでもリアルタイムに推論できる処理能力の高さを実感できました。
　複雑なモデルを動かすことは難しいですが、Edge TPU APIを使って独自モデルを動かすことができるため、応用ができるのではないかと思います。