Dockerコンテナ上でEdge TPUコンパイラを使う

2020年1月に Edge TPU が macOS/Windows でも動作するようになりましたが、Edge TPU コンパイラはまだなので、macOS/Windows から手軽に利用できるように Docker でコンテナ化してみました。R&D チームの奥村（@izariuo440）がお送りします。

ベースイメージ選定
Dockerfile を書く
Docker イメージのビルド
EfficientNet-lite のモデルの準備
Docker でコンパイル
おわりに
おまけ

Edge TPU に関しては、過去記事も参考になるかもしれません。

ベースイメージ選定

まずはベースイメージを決めます。Edge TPU コンパイラは Debian 6.0 (squeeze) 以降で動作するとありましたが、実際に Debian 6.0 (squeeze) を試してみたらapt-get update がこけました。2016年に LTS サポートが切れていますしね。Debian 9.0 (stretch) も通常サポートは今年で切れます。この辺は BASE_IMAGE という引数で変更できるようにしておいて、既定では Debian 10.0 (buster) を使うようにしました。これで通常サポートでも 2022年までは大丈夫のはず。

Dockerfile を書く

公式の手順に従って書くだけです。Edge TPU のためのパッケージレポジトリを追加するために一時的に curl などが必要になります。使い終わったら削除するようにしてイメージが大きくならないようにしておきます。

ARG BASE_IMAGE=debian:buster-slim
FROM ${BASE_IMAGE}

RUN export DEBIAN_FRONTEND=noninteractive \
 && apt-get update \
 && apt-get install --no-install-recommends -y \
    apt-transport-https \
    ca-certificates \
    curl \
    gnupg \
 && curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - \
 && echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" > /etc/apt/sources.list.d/coral-edgetpu.list \
 && apt-get update \
 && apt-get install --no-install-recommends -y \
    edgetpu \
 && apt-get purge --autoremove -y \
    apt-transport-https \
    ca-certificates \
    curl \
    gnupg \
 && apt-get clean \
 && rm -rf /var/lib/apt/lists/* /var/log/apt/*

WORKDIR /work

ENTRYPOINT [ "edgetpu_compiler" ]

エントリポイントに edgetpu_compiler を指定して、docker run が短くなるようにしています。また、カレントディレクトリを作業ディレクトリとしてマウントすることを想定して /work を作業ディレクトリにしておきます。

Docker イメージのビルド

上記を Dockerfile として保存し、ターミナルで以下を実行します。

docker image build -t edgetpu-compiler .

docker images でサイズを確認すると、88.6 MB でした。

以下のように Debian 9.0 (stretch) ベースにすると、74.7 MB でした。

docker image build --build-arg BASE_IMAGE=debian:stretch-slim -t edgetpu-compiler .

EfficientNet-lite のモデルの準備

試しに EfficientNet-lite を Edge TPU 向けにコンパイルしてみましょう。いくつかバリエーションがありますが、最軽量の efficientnet-lite0 で試します。

ターミナルで以下を実行してモデルをダウンロード・展開します。

export MODEL=efficientnet-lite0
curl -fsSLO https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/lite/${MODEL}.tar.gz
tar zxf ${MODEL}.tar.gz

展開結果は以下のようになります。

$ ls -1 ${MODEL}
checkpoint
efficientnet-lite0-fp32.tflite
efficientnet-lite0-int8.tflite
model.ckpt.data-00000-of-00001
model.ckpt.index
model.ckpt.meta
result.txt

これらのファイルのうち efficientnet-lite0-int8.tflite が INT8 量子化済みで Edge TPU 互換な TFLite 形式のモデルです。

Docker でコンパイル

ターミナルで以下を実行します。

docker run --rm -it -v $(pwd):/work edgetpu-compiler ${MODEL}/efficientnet-lite0-int8.tflite

実行すると以下が出力され、無事にコンパイルできたことが分かります。

Edge TPU Compiler version 2.0.291256449

Model compiled successfully in 688 ms.

Input model: efficientnet-lite0/efficientnet-lite0-int8.tflite
Input size: 5.17MiB
Output model: efficientnet-lite0-int8_edgetpu.tflite
Output size: 5.47MiB
On-chip memory available for caching model parameters: 6.91MiB
On-chip memory used for caching model parameters: 5.32MiB
Off-chip memory used for streaming uncached model parameters: 0.00B
Number of Edge TPU subgraphs: 1
Total number of operations: 64
Operation log: efficientnet-lite0-int8_edgetpu.log
See the operation log file for individual operation details.

efficientnet-lite0-int8_edgetpu.log というログファイルを見ると、すべてのオペレーターが Edge TPU にマッピングされたことが分かります。こうなった上で、上記の出力でモデルがオンチップメモリに乗り切る場合、高速推論が期待できます。

Edge TPU Compiler version 2.0.291256449
Input: efficientnet-lite0/efficientnet-lite0-int8.tflite
Output: efficientnet-lite0-int8_edgetpu.tflite

Operator                       Count      Status

RESHAPE                        1          Mapped to Edge TPU
SOFTMAX                        1          Mapped to Edge TPU
FULLY_CONNECTED                1          Mapped to Edge TPU
ADD                            9          Mapped to Edge TPU
AVERAGE_POOL_2D                1          Mapped to Edge TPU
QUANTIZE                       2          Mapped to Edge TPU
CONV_2D                        33         Mapped to Edge TPU
DEPTHWISE_CONV_2D              16         Mapped to Edge TPU

おわりに

Edge TPU ランタイムは macOS/Windows に対応したものの、Edge TPU コンパイラはまだなので、このような方法で macOS/Windows でもコンパイルできるようにしてみました。実は先駆者がいるのですが、Docker Hub には見当たらなかったので自分で書いた次第です。

github.com

Google の Cloud Run でやる方法もあります。社内の量子化警察から教えてもらいました。

github.com

オプティムでは、日々の作業をサクッと効率化していけるエンジニアを募集しています。興味のある方は、こちらをご覧ください。

www.optim.co.jp

おまけ

なお、上記の Dockerfile について色々とツッコミを受けたやりとりが面白かったので紹介しておきます。

T: Dockerfile に関して、イメージをlatestのみ使う場合はまとめたほうが一番小さいですが、Updateごとに変更追跡する場合はRUN分けるほうが全体サイズ小さくなります。今のとこ apt-get install edgetpu が更新ターゲットなので、まるっと差分走るのは煩雑だなと。
K: curl -fSLO の -S って -s がついていないときは何の意味もないはずですよ．
I: tar に z 付ける派ですか？私は付けない派なので、tar xf でいいじゃないと思います。
K: 私もそれ思いましたが，元のページで z ついていたのでそれに合わせるで良いんじゃないかと思いました．
I: なるほど。
K: あと docker run --rm -it -v $(PWD):/work edgetpu-compiler ${MODEL}/efficientnet-lite0-int8.tflite の $(PWD) は $(pwd) もしくは環境変数 $PWD にしておかないと大文字小文字を区別するファイルシステム上で動かした時に PWD という実行ファイルは見つからないというエラーを吐かれると思いますよ．
I: tee /etc/apt/sources.list.d/coral-edgetpu.list は > /etc/apt/sources.list.d/coral-edgetpu.list で良いのではないでしょうか？
M: (tee は一般ユーザで書き込みだけ sudo したいときによく使いますね)
O: 公式手順をベースにしたのが残ってましたね。
S: rm を2回実行しているのが気になります。1個に繋げられると思います
M: buildkitじゃないとだめということはないので DOCKER_BUILDKIT は外してもいいんじゃないですか
I: apt-key の man を読むと、apt-key は使う必要が無いことが分かります。https://manpages.debian.org/unstable/apt/apt-key.8.ja.html curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - の代わりに curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg -o /etc/apt/trusted.gpg.d/google.gpg や ADD https://packages.cloud.google.com/apt/doc/apt-key.gpg /etc/apt/trusted.gpg.d/google.gpg でいいんです。前者だとgpg、後者だとgpgとcurlが不要になります。
S: ADD ってレイヤー増えませんでしたっけ？そこがちょっと気になります
T: 増えます
I: そうなんです。なので、気にするなら curl
O: このままにします。
I: 最後の apt-get clean も不要ではないかと思っています。/etc/apt/apt.conf.d/docker-clean というファイルがあって、そこで apt-get clean 相当のことをやっているという話がありまして。
M: ベースイメージを差し替え可能にする以上、そこは仮定しきれないんじゃないですかね
I: あーたしかに。DebianとUbuntuの公式以外はこれがあるかどうか分からないですね。
K: って言っても，/var/lib/apt/lists/* は削除しているので，/var/cache/apt/* も自分で削除するなら apt-get clean 不要だと思いますけどね．
T: プロダクションに乗せるイメージであれば削減検討必要と思いますが、開発ビルド動かすところに厳密にしなくてもとは思います。やってもやらなくても100MB前後です。用途的には「Edge TPU 何たら Updates」が出たときに docker build が短時間で終わって試せることのほうが重要かなと思ってます。
N: 個人的にはあまり RUN コマンド1つにまとめると、何かの変更で途中でコケたときにぐぇーとなるので悩ましさを感じます。
K: 早く --squash が experimental から外れて多くの人が救われると良いなと思います．