ZOZO研究所の後藤です。本記事ではGoogle Cloud TPUを使った計量学習の高速化の事例を紹介します。
はじめに
深層学習を用いたプロダクトを開発・運用する上で、モデルの学習にかかる膨大な時間はボトルネックの1つです。 ファッションにおける深層学習を用いた画像認識技術にも同じことが言えます。 今回はファッションの分野において定番のタスクであるStreet2shopの課題設定に対し、Google Cloud TPUを用いて計量学習の高速化を試みます。 Street2shopは、スナップ画像から商品部分を切り出す物体検出のパートと、切り出した画像と類似した商品を検索するクロスドメイン画像検索のパートに分けられます。 今回の取り組みでは、後者のパートで利用する画像間の距離を測るためのモデルの学習の高速化を行います。
目次
Google Cloud TPUとは
Tensor Processing Unit（以下、TPU）は、機械学習で頻繁に行われる行列演算を高速に処理することを目的にGoogleが開発したハードウェアです。 CPUやGPUは様々なソフトウェアやアプリケーションに対応するための汎用プロセッサとして設計されている一方で、TPUは汎用性を犠牲にして行列演算に特化した設計がなされています。 そのため、行列演算を多用する機械学習の計算にTPUを用いることによって、モデルの学習や推論にかかる時間を大幅に短縮することが可能になります。 TPUがなぜ高速に演算できるのかを理解するには、公式のブログ記事が詳しいです。
cloudplatform-jp.googleblog.com
速いことは良いことです。 学習速度が速いと、学習が収束したモデルをより短時間で獲得できるようになります。 また、限られた時間の中での試行錯誤の回数を増やすことができるため、ハイパーパラメータチューニングが捗り、最終的に得られるモデルがより良いものになる可能性が高くなります。Cloud TPUはTPUをクラウド上のインスタンスとして利用できるようにしたサービスです。TPUの強力な演算能力を時間単位で誰でも利用できます。
注意点として、TPUは設計上、得意な演算が限られるため、あらゆるモデルを高速化できるわけではありません。 適切な計算環境の選び方に関しては、以下の記事が参考になります。
例えば、以下の4つの項目に当てはまる場合、TPUによる高速化が達成可能かもしれません。
- ほとんどの計算が行列計算である
- 学習ループにカスタムTensorFlow operationsが含まれない
- 学習に数週間から数ヶ月かかる
- 大きなバッチサイズで学習する
一方で、TPUを使うのに向いていない場合もあります。
- 頻繁なブランチングや要素ごとの積が多いプログラム
- スパースなメモリアクセス
- 高精度演算
- カスタムTensorFlow operationを含むネットワーク
タスク
Street2shop
Street2shopはストリートで撮影されたスナップ写真から、商品部分を検出して、ショップの商品と対応づけるタスクです。 ECサイトの商品検索の文脈においてこのタスクを応用すると、言葉では表現しにくい直感的な検索を画像を使って行うことができます。 弊社では過去に、この問題設定を解決するシステムの解説をブログに書いています。
計量学習
写り方の異なる同一商品の画像や、見た目が似ている商品を検索するには、画像同士の距離や類似度が適切に測れる空間が必要です。画像をそのような空間に埋め込むモデルを獲得するタスクが計量学習です。今回は損失関数にN-pair loss、ベースとなるネットワークにResNet-50、最適化の手法にAdam、空間の次元数は256を採用しています。問題設定の詳細な解説は過去のテックブログを参考にしてください。
実装のフレームワークにはTensorFlowを用います。 TensorFlowのTPUEstimator APIを利用して、CPU、GPU、TPU全てのチップで動作する汎用的なコードが得られます。 Googleが用意しているリファレンス実装を元に、自前のタスクに関する部分を書き換えることでTPUを使って動作するコードを得られます。
学習データ
今回の計量学習のために、商品画像と着用時の画像のペアを120万組ほど用意しました。 ストリートで取られたスナップ画像に対して商品検出器をかけ、商品部分を切り出しておきます。今回対象としたカテゴリはトップス、アウター、ドレス、パンツ、スカート、シューズの6種類で、女性ファッションのみに絞りました。
学習
プロファイリング
Cloud TPUによる学習を高速化するためには、速度のボトルネックとなっている部分を特定して、適宜チューニングしていくという作業が必要になります。その際、cloud_tpu_profilerというツールを利用すると、ボトルネックの特定が容易になります。この辺りの作業は、一度Brennan Saetaさんによる実演を見ておくと理解がスムーズかと思います。
www.youtube.com cloud.google.com
cloud_tpu_profilerを実行すると、TensorBoardでプロファイリングの結果が閲覧できるようになります。 サマリー画面は以下のようなものです。 この画面の3段目「Recommendation for Next Steps」に、学習のパフォーマンスを上げるためのアドバイスが書かれています。
TPUの学習速度に対してデータのインプットが追いついていない場合は、TPUに待ち時間が発生しますし、その逆も考えられます。プロファイラの各項目を深掘りすることで、どこが計算のボトルネックになっているのかを特定できます。
各オペレーションとメモリ使用量の関係
各オペレーションとTPUリソースの関係
今回の学習の場合、リファレンス実装をそのまま利用することにより高いパフォーマンスでCloud TPU動作させることができたため、特別なチューニングは行なっていません。 TPUはフィードフォワードとバックプロパゲーションが高速であり、データの読み込みがボトルネックになることがあります。 例えば、データの読み込み速度と学習速度のバランスをみてバッチサイズを変えるというチューニングが考えられます。 今回は、1バッチあたり1024枚と大きなミニバッチサイズを設定します。
学習結果
TPUにはコアが8つあるため、シングルGPUとの単純な比較は不公平ですが、現状の計算環境とTPUを使った場合の差分をはっきりさせるために敢えて比較します。 以下の表は、同じタスクを5日間ほど回した際の各評価値です。 学習の速度の観点では、P100で学習した場合、1秒あたり20組の画像を学習するのに対し、TPU v2では1100枚の画像を学習できました。 Cloud TPUを使うことにより、これまでに比べて55倍の学習の高速化が達成できたと言えるでしょう。 精度の観点では、同じ期間内にイテレーションをより多く回せたため、TPUで学習したモデルのAccuracyがGPUで学習したモデルのものより高い値になっています。
| Unit | Batch Size | Top1 Accuracy (batch size 64) | Top5 Accuracy (batch size 64) | Train Loss | Val Loss | Elapsed Time / 1000 step | N image / sec |
|---|---|---|---|---|---|---|---|
| Nvidia P100 | 64 | 0.6511 | 0.9085 | 1.784 | 1.355 | ~ 54 min | ~ 19.75 image |
| Cloud TPU v2 | 1024 | 0.8070 | 0.9573 | 1.155 | 0.7812 | ~ 15.5 min | ~ 1100 image |
計算にかかる料金
2019年1月31日の時点でCloud TPU v2の単位時間あたりの料金は$4.50 USDです。
cloud.google.com 一方で、NVIDIA P100のGPUの料金は単位時間あたり1GPUあたり$1.46 USDです。
具体的な料金はCPUの数やメモリ、ディスクの容量にもよりますが、料金はせいぜい数倍の差です。その一方で、55倍の速さで学習ができたと考えると、Cloud TPUを使った方が最終的にかかる料金が低く抑えられるかもしれません。
モデルの活用
今回学習したモデルはインプットとして、「スナップ画像から商品部分を切り出した画像」と「商品画像」の両方を受け付け、256次元の数値をアウトプットします。 Shop2shopという単なる類似画像検索と、Street2shopという問題設定の両方を定性的に評価してみます。 評価時の入力と検索対象には、学習には用いていないバリデーション用データセットを利用しています。
Shop2shop
左側の赤枠で囲んだ画像を空間に埋め込み、ユークリッド距離が近い順に商品を選び出すという検索を行なっています。商品画像同士の類似度を直接学習したわけではありませんが、多くの場合で、形、色、柄、素材の観点で似ている商品が検索できています。以下、カテゴリ別の検索結果になります。
トップス
- 色、形はもちろんのこと、中央に写っているキャラクターも一致している
- ネック部分や袖の形の一致、花の刺繍部分の特徴も捉えることができている
shop2shop_tops
アウター
- 丈、色、カテゴリーが近いものを検索できている
- スカジャンの入力画像のポーズが独特だが、正面を向いたものが検索できている
shop2shop_outer
ドレス
- 丈、色、形、ネックの一致が見られる
- シースルーの素材感や柄の特徴も捉えられている
ボトムス
- 色、形、柄の一致が見られる
- 例には載せていないが、ダメージジーンズのダメージ具合やダメージ位置の特徴の一致も検索できる
shop2shop_bottoms
スカート
- 丈、素材の一致が見られる
- デニムスカートのボタンの位置の特徴を押さえることが出来ている
shop2shop_skirt
シューズ
- 色、形、丈、素材感の一致が見られる
- 検索結果に着用画像と商品単体画像が混じり、その両方をある程度公平に評価できていることがわかる
shop2shop_shoes
Street2shop
ストリートの環境で撮影された画像の商品部分を切り出して入力に使うタスクです。 自然光のもとで撮られたり、影が発生したりと複雑なシチュエーションが多くShop2shopよりも難しいタスクです。
スナップ1
- トップスの文字部分の特徴を捉えている
- フォントの形も似ているものが選ばれている
- パンツは形が異なる商品も含まれる一方で、色、幅、丈が一致する商品も含まれる
street2shop_1
スナップ2
- ロゴが含まれたTシャツを検索することができる
- スカートのボタンのポイントは押さえているが、丈や形の面で間違えるパターンがある
- シューズはつま先とヒール部分を押さえられている
まとめ
Cloud TPUを使うことで今回の計量学習に関しては、単体のGPU（Tesla P100）を使って学習した場合に比べて、55倍の学習速度の改善ができました。 55倍の速度改善は、1週間かかっていたモデルの学習が3時間で終わることを意味します。 Cloud TPUによる高速化が可能かどうかの見極めは必要になりますが、時間短縮のためにTensorFlow + Cloud TPUを使うことは一考に値するでしょう。
最後に
弊社では次々に登場する新しい技術を使いこなし、プロダクトの品質を改善できる機械学習エンジニアを募集しています。 ご興味のある方は、以下のリンクからぜひご応募ください。
