2019年4月7日日曜日

Coral USB Accelerator(Edge TPU)で犬猫のImage classification(GPU)

Coral USB Accelerator(Edge TPU)でRetrain an image classification modelを試してみる。せっかくなので

  • 他のデータセット(今回は犬と猫の分類)を試してみる。
  • GPUで学習できるようにする。
  • モデルはとりあえずInception v1。
を目指してみた。
Edge TPUのImage classificationのサンプルは「TensorFlow-Slim image classification model library」を使っている。これを簡単にスクリプトから実行できるようにしている。
(Image classificationはKerasが主体になるので、slimを使ってのImage classificationはやったことがなかった。そのへんも理解したい。)

リポジトリはここに作成した。

環境

自分のPC環境
  • CPU AMD Ryzen 7 1700 Eight-Core Processor
  • GPU NVIDIA GeForce GTX 1070
  • OS Fedora 29

Docker準備

まず、Fedora29にCUDA、Docker、NVIDIA dockerをインストールしておく。
  • CUDAのインストールはここ
  • Dockerのインストールはここ(Get Docker CE for Fedora)
  • NVIDIA Dockerのインストールはここ
    (NVIDIA DockerはFedora29はサポートしていないが、CentOS/RHELと同じものをインストールすれば動く)

作業ディレクトリの作成とデータの配置

適当な作業ディレクトリを設定しておく。
学習データはKaggleのDogs vs. Catsで提供されているデータを使用する。
(Kaggleに登録してダウンロードしておく)
ダウンロードしたzipファイルを作業ディレクトリ/data/配下に置く。

$ CLASSIFY_DIR=..../classify && mkdir -p 
$ mkdir {$CLASSIFY_DIR}/ckpt
$ mkdir {$CLASSIFY_DIR}/train
$ mkdir {$CLASSIFY_DIR}/data
$ mkdir {$CLASSIFY_DIR}/models
$ cp dogs-vs-cats-redux-kernels-edition.zip {$CLASSIFY_DIR}/data/

※実は、Dogs vs. Catsのデータはマイクロソフトからもダウンロードできるのだが、一部ファイルが壊れていた。

Docker containerのセットアップ

DockerfileはGPUで学習できるように変更。
あとは、TensorFlowのmodelsをForkしたリポジトリをcloneしている。
(これは、Dogs vs. Catsのデータセットを追加してあるため)
リポジトリのDockerfileを移動してビルドし、起動してあげる。
$ sudo docker build - < Dockerfile --tag classify-dogcat
$ sudo docker run --name edgetpu-classify --rm --runtime=nvidia -it --privileged -p 6006:6006 --mount type=bind,src=${CLASSIFY_DIR},dst=/tensorflow/models/research/slim/transfer_learn classify-dogcat


データセットの準備

学習のためのTF-Recordを作成するためのスクリプトを準備する。
TensorFlow Modelsのdownload_and_convert_flowers.pyを参考にTF-Recordを生成するcreate_tfrecord_dogcat.pyを作成した。
犬、猫それぞれ12500枚から学習用にそれぞれ2,000枚、評価用にそれぞれ500枚を使ってTF-Recordを生成する(学習用:dogcat_train_xxxxx-of-00005.tfrecord、評価用:dogcat_validation_xxxxx-of-00005.tfrecord)。
また、画像とラベルを対応付けるためにslimのdataset_factory.pyを変更し、dogcat.pyを追加した。詳細はこちらを参照。

これで、コンテナ内でTF-Recodeを生成する。

# python edge_tpu/dataset/create_tfrecord_dogcat.py ${DATASET_DIR}/dogs-vs-cats-redux-kernels-edition.zip "${DATASET_DIR}"

学習の準備

Inception V1用にPre-trainedモデルのダウンロードと環境変数の登録を行う。
checkpointのダウンロードは、Tensorflowのmodelsで公開されているものを使った。

(Edge TPUのデモで使用するPre-trainedモデルは異なるリンクを指していた。何が違うのかは不明)

Checkpointのダウンロード
# wget http://download.tensorflow.org/models/inception_v1_2016_08_28.tar.gz
# tar xvf inception_v1_2016_08_28.tar.gz
# mv inception_v1.ckpt transfer_learn/ckpt/
# rm inception_v1_2016_08_28.tar.gz


環境変数の設定
# network_type=inception_v1
# train_steps=300
# quantize_delay=100
# image_size=224
# scopes=InceptionV1/Logits

学習

train_image_classifier.pyで学習を行う。

Pre-trainedモデルを使用する場合、--checkpoint_exclude_scopesフラグを指定して、出力層の重みをロードしないように指定する(Pre-trainedモデルの出力層と学習したい出力層のラベル数が異なるため。--checkpoint_exclude_scopesフラグで指定する層は${scopes}の環境変数で指定したもの)。
(Edge TPUのでもに含まれるconstants.shのscopes_mapを参照するか、コードから出力層がどのように定義されているかを確認すると良い)

--trainable_scopesは、出力層のみを学習したい場合に指定する。

ここまでは、Edge TPUに限らず一般的な学習の方法だが、--quantize_delayフラグだけは、Edge TPU(というかTensorflow Lite)向けのパラメータとなる。これは、モデルを量子化した際の精度をシミュレートするためのフラグである(詳細はここ)。指定するパラメータはシミュレートを開始するstep数で、ある程度学習が収束してからが良いとある。

全部の層を学習する場合
# python train_image_classifier.py\
     --train_dir="${TRAIN_DIR}"\
    --dataset_name=dogcat\
    --dataset_split_name=train\
     --dataset_dir="${DATASET_DIR}"\
     --model_name="${network_type}"\
     --checkpoint_path="${CKPT_DIR}inception_v1.ckpt"\
     --max_number_of_steps="${train_steps}"\
     --batch_size=32\
     --learning_rate=0.01\
     --learning_rate_decay_type=fixed\
     --save_interval_secs=60\
     --save_summaries_secs=60\
     --log_every_n_steps=20\
     --optimizer=sgd\
     --weight_decay=0.00004\
     --clone_on_cpu\
     --train_image_size="${image_size}"\
     --checkpoint_exclude_scopes="${scopes}"\
     --input_queue_memory_factor=16

最終層のみを学習する場合
# python train_image_classifier.py \
    --train_dir="${TRAIN_DIR}" \
    --dataset_name=dogcat\
    --dataset_split_name=train \
    --dataset_dir="${DATASET_DIR}" \
    --model_name="${network_type}" \
    --checkpoint_path="${CKPT_DIR}inception_v1.ckpt"\
    --max_number_of_steps="${train_steps}" \
    --batch_size=100 \
    --learning_rate=0.01 \
    --learning_rate_decay_type=fixed \
    --save_interval_secs=60 \
    --save_summaries_secs=60 \
    --log_every_n_steps=20 \
    --optimizer=sgd \
    --weight_decay=0.00004 \
    --quantize_delay="${quantize_delay}" \
    --clone_on_cpu \
    --train_image_size="${image_size}" \
    --checkpoint_exclude_scopes="${scopes}" \
    --trainable_scopes="${scopes}"

学習後の確認

テストデータセットを使って評価する。

# python eval_image_classifier.py\
   --checkpoint_path="${TRAIN_DIR}"\
   --eval_dir="${TRAIN_DIR}"\
   --dataset_name=dogcat\
   --dataset_split_name=validation\
   --dataset_dir="${DATASET_DIR}"\
   --model_name="${network_type}"\
   --eval_image_size="${image_size}"\
   --quantize


全部の層を学習した場合
eval/Accuracy[0.994]
eval/Recall_5[1]
最終層のみを学習した場合

eval/Accuracy[0.976]
eval/Recall_5[1]

この場合も--quantizeフラグは量子化をシミュレートするためのものである。
また悔過kのRecall_5はTop5の再現率であるので2クラス分類では必ず1である。

モデルのエクスポートとEdge TPU向けのモデルの作成

まず、tflite向けのモデルを生成する。

GETTING frozen graph template ...
# mkdir transfer_learn/output
# OUTPUT_DIR=/tensorflow/models/research/slim/transfer_learn/output
# python export_inference_graph.py  \
  --alsologtostderr   \
  --model_name="${network_type}"  \
  --image_size="${image_size}"  \
  --output_file="${OUTPUT_DIR}/graph_template.pb" \
  --dataset_name=dogcat \
  --quantize

CONVERTING from checkpoint to frozen graph ...
# output_tensors=InceptionV1/Logits/Predictions/Softmax
# freeze_graph\
   --input_graph="${OUTPUT_DIR}/graph_template.pb"\
   --input_checkpoint="${TRAIN_DIR}model.ckpt-300"\
   --input_binary\
   --output_graph="${OUTPUT_DIR}/frozen_graph.pb"\
   --output_node_names="${output_tensors}"

CONVERTING from frozen graph to TFLite graph ...
# input_tensors=input
# tflite_convert\
   --output_file="${OUTPUT_DIR}/inception_v1_dog_cat_tflite_graph.tflite"\
   --graph_def_file="${OUTPUT_DIR}/frozen_graph.pb"\
   --inference_type=QUANTIZED_UINT8\
   --input_arrays="${input_tensors}"\
   --output_arrays="${output_tensors}"\
   --mean_values=128\
   --std_dev_values=128\
   --input_shapes="1,${image_size},${image_size},3"

tflite_convertを使うことで、tflite向けのモデルが生成できる。--mean_values、--std_dev_valuesを変えるとどのように変化するのかはまだ調査できていない。
生成した.tfliteファイルをEdge TPU Model CompilerでEdge TPU向けのモデルとして生成し、ダウンロードすればよい。

確認

Raspberry Pi+Edge TPUでモデルを動かしてみる(モデルは全部の層で学習)。
まずは、Dogs vs. Catsから適当なデータ。
$ python3 classify_image.py --model ~/dog_vs_cat/inception_v1_dog_cat_tflite_graph_1554209169_edgetpu.tflite --label ~/dog_vs_cat/labels.txt --image ~/dog_vs_cat/dog.1320.jpg
---------------------------
dog
Score :  0.996094

$ python3 classify_image.py --model ~/dog_vs_cat/inception_v1_dog_cat_tflite_graph_1554209169_edgetpu.tflite --label ~/dog_vs_cat/labels.txt --image ~/dog_vs_cat/cat.9276.jpg
---------------------------
cat
Score :  0.996094

実家の猫(もう13年前の写真。でもまだ生きている。)
これも推論させてみる。
$ python3 classify_image.py --model ~/dog_vs_cat/inception_v1_dog_cat_tflite_graph_1554209169_edgetpu.tflite --label ~/dog_vs_cat/labels.txt --image ~/012.JPG
---------------------------
cat
Score :  0.996094

おぉ、猫である。