Edited at

Tensorflow Object Detection APIで寿司検出モデルを学習するまで


内容

最近DeepLearningの勉強を始めて、[Tensorflow Object Detection API]を使って、自分で集めた画像を使って学習してみました。学習データの収集、ローカルマシンでの学習からCloudでの学習まで自分が経験したことを共有したいと思います。

追記:

学習したモデルをiOSで動かすまでの記事を書きました。

https://heartbeat.fritz.ai/building-a-real-time-object-recognition-ios-app-that-detects-sushi-c4a3a2c32298


Tensorflow Object Detection API

今年6月に公開されたTensorflow Object Detection APIを使うと、画像から物体(Object)とその位置を検出するためのモデルを作ることが簡単にできます。

実際に僕が触って見たときは、PythonやTensorflowの知識が無かったのですがチュートリアルに沿ってすぐに試して見ることができました。


学習データ


寿司画像

今回学習に利用した画像は寿司の画像になります。

なぜ寿司かというと、

オーストラリア出身の同僚が寿司に関するアプリを一緒に作ろうとずっと言われていたんですが

M花の種類を判別するMLのサンプルモデルを見て、専門家ならわかるけど素人にはわからないような問題をアプリで解決できたらいいなと思って寿司にトライしてみました。

image.png

後述のラベリングにも時間がかかるので、今回は中トロ・サーモン・コハダの3種類 100枚づつ集めて試してみました。

寿司以外にも食べ物にフォーカスした画像は比較的集めやすそうなので、今後は他の食べ物も試して行きたいと思います。(インスタとかグルメサイトとか)


ラベリング

画像が集まったら、1枚1枚の画像の 'どこ' に 'なに' があるのかのラベル付けが必要です。

今回はlabellmgというツールで1枚づつ手動でラベルを付けていきました。

image.png


学習の準備


TFRecordの作成

学習データの準備ができたら次はTFRecordというTensorflowに対応した形に画像とラベリングデータを変換する必要があります。 Tensorflowがインストールされている FloydhubのDockerイメージを使って、Object Detection APIをインストールしたコンテナー内で変換スクリプトを実行しました。

詳細はGithubリポジトリを参考にしてみてください。

参考


トレーニングConfigファイルの作成

次にトレーニングの設定ファイルを用意します。サンプルファイルから基本コピペで使いました。

またモデルを選ぶ必要があるのですが、 最終的にスマホアプリに組み込みたいと思っているので ssd_mobilenet_v1_coco というモデルを選択しました。

実際に用意したConfigファイルがこちらです。

参考


ローカルマシンでの学習

さて、やっとこれで学習が開始できます。 さきほど用意したこConfigファイルと学習中のデータを書き出すディレクトリを指定した実行します。

python -m object_detection.train

--logtostderr \
--pipeline_config_path=/tensorflow/data/ssd_mobilenet_v1_sushi.floyd.config
--train_dir=/data/train

こちらもレポジトリのスクリプトを参考にしてみてください。

スクリーンショットを撮り忘れてしまったのですが、数十秒/ステップくらいのスピードで学習していました。

image.png

1000ステップくらいでも実際に寿司を検出できているのがわかります。

他の方の記事を参考に、最低10,000ステップくらい学習したかったのでこのままだと2,3日かかる計算ですね。

ここからはCloudGPUで学習していくことにしました。


Cloudでの学習

自前のGPUマシンを持っていないので、CloudGPUを使う必要があったのですが、ここで問題が発生しました。

チュートリアルで紹介されているGoogle CloudMLが最新のTensorflowに対応していない!涙

そこで見つけたのがFloydhubでした。


Floydhubについて

FloydhubはデータサイエンティストのGithubのようなサービスで、学習データのバージョニング・学習ジョブの実行/管理などがCLI経由できます。


学習ジョブの確認

いきなりGPUを使うまえに、まずはCPUを使って設定を試行錯誤しました。

そこで行き着いたのがこちらのスクリプトです。

floyd run --env tensorflow-1.4 \

--data junji/datasets/sushi_detector/1:data \
"bash ./floyd/setup.sh && cp -R /data/* /output && sh ./floyd/train.sh"

これは



  • /data ディレクトリに学習データをセット


  • setup.shでObjectDetectionAPIのインストール


  • /dataのファイルを/outputにコピー

  • 学習の実行

を行っています。

Floydhubでジョブを実行するときは書き込みが出来るのが /output ディレクトリだけに制限されており、このディレクトリはジョブの終了後も保存されます。

次に学習を再開するときは、下記のように /data ディレクトリに前回のoutputを使うことで途中から学習を再開することができます。

floyd run --env tensorflow-1.4 \

--data junji/projects/sushi_detector/{job_id}/output:data \
"bash ./floyd/setup.sh && cp -R /data/* /output && sh ./floyd/train.sh"

こちらも詳細はレポジトリを参考にしてみてください。

image.png

ローカルマシンと同程度の学習スピードでした。


GPUで実行

全て上手く動いていることが確認できたので、ついに念願のGPUを使って学習をしました。

先程のfloydコマンドのオプションに --gpu をつけるだけす。とても簡単!

floyd run --gpu --env tensorflow-1.4 \

...

そして、結果は

image.png

おおー、もうスピードが全然違いますね 1ステップ1秒以下。20倍以上の速度でしょうか。

今回は10,000ステップを超えたあたりで止めました。


学習結果

ちゃんと出来てる、寿司を認識していますね

中トロ
サーモン
コハダ

image.png
image.png
image.png


まとめ

詳細な説明はかなり省いてしまったのですが、この記事で物体検出モデルの学習が簡単に出来ることが皆さんに伝われば幸いです。

ぼくは普段アプリケーションエンジニアとして仕事をしていますが、これくらいの手間で物体検出が出来るならサービス開発の際も全然使えそうだなと感じました。

もっとトライして、実際のアプリケーションに導入するところまで勧めていきたいと思います!

最後まで読んで頂いてありがとうございました。