Edge TPUを用いて，自作データセットによるClassifierの再学習（転移学習）を試してみた．
Edge TPUでは，サーバで分類モデルを再学習する代わりに，ImprintingEngine APIを使用してEdge TPU上で転移学習を実行することができる．
Edge TPUのImprintingEngine APIはCVPR2018で提案されたLow-shot learning with imprinted weights [1]を用いて転移学習を実現する．
この手法は，再学習の際にbackwardが必要ないため，Edge TPU上で高速に再学習を実行することが出来る．
Edge TPU APIのインストール
$ wget http://storage.googleapis.com/cloud-iot-edge-pretrained-models/edgetpu_api.tar.gz $ tar xzf edgetpu_api.tar.gz $ cd python-tflite-source $ bash ./install.sh
Feature extractorのダウンロード
特徴抽出を行うためのpretrained modelをダウンロードする．
$ wget http://storage.googleapis.com/cloud-iot-edge-pretrained-models/canned_models/mobilenet_v1_1.0_224_quant_embedding_extractor_edgetpu.tflite
データセットの準備
今回はハムスターとモルモットを分類してみる．
以下のようなディレクトリ構造で画像を保存しておく．
サブディレクトリ一つが一つの分類クラスに対応する．
cavy_hamster/
├── cavy
│ ├── 0000006.jpg
│ ├── 000001.jpg
│ ├── 000002.jpg
│ ├── 000003.jpg
│ ├── 000004.jpg
│ ├── 000005.jpg
└── hamster
├── 000001.jpg
├── 000002.jpg
├── 000003.jpg
├── 000004.jpg
├── 000005.jpg
今回は，各クラスで20枚ずつ画像を用意した．
学習の実行
Edge TPU USB AcceleratorをUSB接続した状態で以下を実行する．
$ python3 demo/classification_transfer_learning.py --extractor mobilenet_v1_1.0_224_quant_embedding_extractor_edgetpu.tflite --data /cavy_hamster --output cavy_model.tflite --test_ratio 0.95
- extractor: Edge TPU用にコンパイルした特徴抽出モデル
- data: 用意したデータセット
- output: Edge TPUで転移学習した結果のモデルの保存先
- test_ratio: テストに用いる画像の割合
学習結果
$ python3 demo/classification_transfer_learning.py --extractor mobilenet_v1_1.0_224_quant_embedding_extractor_edgetpu.tflite --data /cavy_hamster --output cavy_model.tflite --test_ratio 0.95 ---------------------- Args ---------------------- Embedding extractor : .//mobilenet_v1_1.0_224_quant_embedding_extractor_edgetpu.tflite Data set : .//cavy_hamster Output path : .//cavy_model.tflite Ratio of test images: 95% --------------- Parsing data set ----------------- Dataset path: .//cavy_hamster Image list successfully parsed! Category Num = 2 W third_party/darwinn/driver/package_registry.cc:65] Minimum runtime version required by package (5) is lower than expected (10). ---------------- Processing training data ---------------- This process may take more than 30 seconds. Processing category: cavy Processing category: hamster ---------------- Start training ----------------- W third_party/darwinn/driver/package_registry.cc:65] Minimum runtime version required by package (5) is lower than expected (10). ---------------- Training finished! ----------------- Model saved as : .//cavy_model.tflite Labels file saved as : .//cavy_model.txt ------------------ Start evaluating ------------------ W third_party/darwinn/driver/package_registry.cc:65] Minimum runtime version required by package (5) is lower than expected (10). Evaluating category [ cavy ] Evaluating category [ hamster ] ---------------- Evaluation result ----------------- Top 1 : 86% Top 2 : 100% Top 3 : 100% Top 4 : 100% Top 5 : 100%
再学習後には，以下が出力として得られる．
- cavy_model.tflite: 再学習後モデル
- cavy_model.txt: ラベルファイル
ラベルファイルは以下のフォーマットで保存されている．
$ cat cavy_model.txt 0 cavy 1 hamster
再学習したモデルを試す
先程学習したモデルをロードして，新しい入力画像に対して分類を行ってみる．
$ python3 demo/classify_image.py --model ./cavy_model.tflite --label ./cavy_model.txt --image test.jpg --------------------------- cavy Score : 0.71953125 --------------------------- hamster Score : 0.28046875
参考
- [1] Qi, Hang, Matthew Brown, and David G. Lowe. "Low-shot learning with imprinted weights." Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. 2018.
