KerasモデルをEdge TPU用モデルに変換メモ @ TensorFlow 2.0

概要

• Kerasモデル（MobileNet）を、量子化したTensorflow Liteモデルに変換
• 変換したTensorflow LiteモデルをEdge TPU向けにコンパイル
  • ※ Debian系の実行環境が必要
• Coral USB Acceleratorで動作チェック
• （余談）VirtualBoxやDockerコンテナ上から動かす

Keras → Tensorflow Lite

Python API を使って変換する。公式ドキュメントはこちら。https://www.tensorflow.org/lite/convert/python_api?hl=ja

以下は、KerasのMobileNet V2を、TensorFlow Lite用に変換するサンプル。

ポイントは、Full integer quantization（訳すとしたら、全量子化？）として変換すること。公式ドキュメントはこちら。https://www.tensorflow.org/lite/performance/post_training_quantization#full_integer_quantization_of_weights_and_activations

この変換には、converter.representative_datasetに必ずジェネレーターをセットする必要がある。8-bit整数への量子化時に行うキャリブレーションのために利用されるサンプルを指定するという目的があるみたい。

ドキュメントに、"Get sample input data as a numpy array in a method of your choosing." とあるので、推論時にモデルに与える入力データのいくつかをサンプルとして与えてあげる。

Keras → TensorFlowLite

import pathlib
import tensorflow as tf
from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2

model = MobileNetV2(weights='imagenet')
converter = tf.lite.TFLiteConverter.from_keras_model(model)

def representative_dataset_gen():
    yield モデルへの入力データとなるサンプル

converter.representative_dataset = representative_dataset_gen
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
converter.optimizations = [tf.lite.Optimize.DEFAULT]

tflite_model = converter.convert()
pathlib.Path('./mobilenet_imagenet.tflite').write_bytes(tflite_model)

Edge TPU で実行可能へ・・

上で変換したモデルは、そのままではEdgeTPUで処理してくれない。
EdgeTPUコンパイラでさらに変換が必要。

コンパイラの最新のインストール方法は、次のURLを参考にすればよい。
→Download項のコマンドを実行するだけ。あと、Debian系の環境が必須。
https://coral.ai/docs/edgetpu/compiler/

EdgeTPUコンパイラのインストール

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

コンパイルを実行すると、_edgetpuがファイル名の末尾についた目的のモデルが生成される。
量子化できていなかった場合、"Model not quantized" と怒られ変換できない。

コンパイルコマンド

edgetpu_compiler mobilenet_imagenet.tflite
# → mobilenet_imagenet_edgetpu.tflite が生成される

ちなみに、コンパイルで変換後のモデルは、普通にロードするとエラーが発生するので、Interpreter生成時にexperimental_delegatesの指定が必要になる。

動作チェック

Interpreterの生成部分にexperimental_delegatesを追加するだけで、その後の使い勝手は特に変わらない。

Interpreter生成

# CPU Only
interpreter = tf.lite.Interpreter(model_path='./mobilenet_imagenet.tflite')

# with Edge TPU
interpreter = tf.lite.Interpreter(model_path='./mobilenet_imagenet_edgetpu.tflite',
                                  experimental_delegates=[tf.lite.experimental.load_delegate('libedgetpu.so.1'),])

# あとは、同じ流れ
# 【割当】allocate_tensors()
# 【入力】set_tensor(入力層のインデックス情報（get_input_details()参照）, ...)
# 【推論】invoke()
# 【出力】get_tensor(出力層のインデックス情報（get_output_details()参照）)

EdgeTPUで処理できていると、Coral USB Acceleratorの白色LEDがピカピカと点滅する。

ちなみに、動作確認で作ったサンプルは、GitHubに。

余談1. VirtualBoxで

検証環境

• Virtual Box 6.0
• ホストOS：macOS Catalina
• ゲストOS：Ubuntu Server 18.04.3 LTS

やり方

• VirtualBox Extension PackでUSB3.0を利用可能にしておく。
• VMにRuntime環境を入れておく。
  • https://coral.ai/docs/accelerator/get-started#1-install-the-edge-tpu-runtime
• Settings→USB に設定追加
  • Coral USB Acceleratorがホストで認識されてれば、Global Unichip Corp.が選択できるので、2回追加。
  • 一方のUSBの設定内容を、次のように変更する。
    • Name: Google Inc.
    • Vendor ID: 18d1
    • Product ID: 9302
• VMを起動して、サンプルなどを動かす。

起動したてにlsusb打つとGlobal Unichip Corp.のままだったりするが、実行後（ライブラリアクセス後？）にGoogle Inc.になるのであまり気にしない。

Runtime環境　インストール別方法

git clone https://github.com/google-coral/edgetpu
sudo ./edgetpu/scripts/runtime/install.sh

余談2. Dockerで

検証環境

• OS: Ubuntu 18.04.3 LTS
• Docker 19.03

※物理環境のLinuxを用意したほうが良い。ホストのUSBをマウントする形になるので、Docker for Windows などハイパーバイザでは不向き。

やり方

Dockerファイルの例

EdgeTPU-runtime.dockerfile

FROM debian:buster

WORKDIR /workspace

RUN apt update && \
    apt install -y curl gnupg ca-certificates zlib1g-dev libjpeg-dev \
                   git python3 python3-pip && \
    apt clean && \
    rm -rf /var/lib/apt/lists/*

RUN echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | tee /etc/apt/sources.list.d/coral-edgetpu.list && \
    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \
    apt update && \
    apt install -y libedgetpu1-std

RUN curl https://dl.google.com/coral/python/tflite_runtime-1.14.0-cp37-cp37m-linux_x86_64.whl > tflite_runtime-1.14.0-cp37-cp37m-linux_x86_64.whl && \
    pip3 install tflite_runtime-1.14.0-cp37-cp37m-linux_x86_64.whl

以下のように、USBが使えるようにマウント指定、privilege付与して、コンテナを起動する。

docker run -it --privileged --rm -v /dev/bus/usb:/dev/bus/usb edgetpu-runtime bash