近似最近傍探索Indexを作るワークフロー

機械学習 自動化 設計 マイクロサービス バックエンド インフラ Kubernetes Airflow

はじめに

こんにちは。ZOZO研究所のshikajiroです。主に研究所のバックエンド全般を担当しています。ZOZOでは2019年夏にAI技術を活用した「類似アイテム検索機能」をリリースしました。商品画像に似た別の商品を検索する機能で、 画像検索 と言った方が分かりやすいかもしれません。MLの開発にはChainer, CuPy, TensorFlow, GPU, TPU, Annoy、バックエンドの開発にはGCP, Kubernetes, Docker, Flask, Terraform, Airflowなど様々な技術を活用しています。今回は私が担当した「近似最近傍探索Indexを作るワークフロー」のお話です。 corp.zozo.com

目次

画像検索の全体像説明

クラウドはGCPを採用しています。分析基盤にBigQueryを使っていること、KubernetesのマネージドサービスであるGKEが安定していること、TPUを活用していることなどが採用の理由です。図の上から簡単に紹介します。

techblog.zozo.com

Workflow

ワークフローであるComposerは毎日販売中のおよそ300万の商品画像から特徴量抽出を行いIndexを作成しています。今回のお話のメインはここですが、もう少し全体像の説明を続けます。

Develop

アプリケーションのコードはGitHub, CI/CDはCircleCIで管理しており、モデルはGCSに、アプリケーションはDockerとしてGCRに登録します。デプロイはCircleCIやComposerを使っています。なお、ML部分のCI/CDはまだ完全には実現できてません。

Application

ユーザーが指定した画像から似た商品画像を返すAPIのことを 推論API と呼びます。Kubernetesで構成しており、後で紹介するマイクロサービスが連携して動作しています。

推論APIの流れ

Composerの説明をする前に、特徴量Indexの動きを知るために推論APIの流れを紹介します。API、物体検出、特徴量抽出、近似最近傍探索、それぞれをマイクロサービスとして動かしています。商品情報検索だけはk8s外の外部サービスです。

ユーザーが画像検索に画像を送ると画像の中に写っている服などのアイテムを検出します。これが 物体検出 です。ここでトップスやシューズなどに分類します。判別した画像のままでは計算に適さないので、検出した部位を多次元ベクトルの特徴量に変換します。この特徴量で距離や類似度などの計量を計算できるようにします。具体的には512次元のfloatの配列になります。これが 特徴量抽出 です。過去テックブログでも紹介していますのでぜひご覧ください。techblog.zozo.com techblog.zozo.com

特徴量から予め準備していたZOZOTOWNの約300万画像のIndexを使って、似ている商品画像を高速に探します。これが 近似最近傍探索 です。近似最近傍探索については東京大学の松井先生の資料が大変分かりやすいので、興味がある方は御覧ください。 speakerdeck.com

近似最近傍探索の段階では似ている商品の画像までしか分かっていません。最新の商品情報を取得するため、ZOZOTOWNの商品データベースに問い合わせてデータの整合性を保ちます。この部分は画像検索とは直接関係ないですが、実際のサービスでは大事な部分ですので紹介しました。

近似最近傍探索とAnnoy

上でも少し触れましたが、ZOZOTOWNが販売する約300万商品画像の中から最も似ている数十〜数百商品を高速に検索する必要があります。特徴量は多次元ベクトルであり対象商品も大量なため、普通に計算するととんでもない計算時間になってしまいます。ここで利用するのが近似最近傍探索です。これを実装した代表的なPythonのライブラリにSpotifyが開発しているAnnoy, Facebookが開発しているFaissなどがあります。画像検索の開発時に主にこの2つを検討し、Annoyを採用しました。理由は以下の2つです。

  • 実装が容易であること
  • 十分に高速であること

もっと速度が必要ならばFaissとGPUの組み合わせを検討していましたが、AnnoyとCPUの組み合わせで速度・精度ともに十分だったため、現状はFaissにする予定はありません。

github.com

近似最近傍探索Indexを作る

Annoyを使って近似最近傍探索を行うには、予め検索対象である約300万画像の特徴量を抽出して Index を作っておく必要があります。ビルド処理自体は数行のコードで実現できるのですが、画像を準備するのがなかなか大変です。以下の手順で作成しています。

  1. BigQueryから現在販売中の画像情報を取得する
  2. 商品画像をZOZOTOWNストレージから取得する
  3. 画像から特徴量を抽出する
  4. 特徴量からAnnoyのBuildを行う
  5. IndexをGCSに保存する
  6. 推論APIで利用する

これらをバッチプログラムを書いて実装することも可能ですが、上記それぞれで必要なCPUリソースもバラバラで、特に特徴量抽出は多くの計算資源を必要とします。エラー時のリトライ、Slackなどへの通知、ロギングの仕組みなどを考えるととても2から作るのは大変です。

そこで利用したのがAirflowなどに代表されるワークフローツールです。

ワークフローツールの説明

上でも触れましたが、バッチ処理を自前で書くとエラー処理、リトライ、ログ、通知処理など実装・運用コストが高いです。この辺りの面倒を見れくれるワークフローツールを検討しました。Airflow, Digdagなどがありますが、GCPのAirflowマネージドサービスであるComposerを採用しました。マネージドサービスだったのが一番の理由です。

AirflowなどのワークフローはDAG（Directed Acyclic Graph, 有向非巡回グラフ）という概念でタスク同士の依存関係を定義しています（詳しくは公式サイトに委ねます）。上記の1〜6の流れを簡潔に書けるようになり、途中からの実行などが容易になります。

airflow.apache.org

デメリット

ComposerのAirflowバージョンはGCPが管理しているため、最新のAirflowより遅れています。そのため、解決されていないバグや対応していないGKEオペレーションなどがあり、使い勝手・保守性に問題がありました。最近はバージョンでは追従できており、安定して稼働しています。

画像検索のワークフローの全体像

ざっと流れを紹介します。BigQueryから現在販売中の商品の画像URLをすべてダウンロードし、Filestoreに保存します。実際は前日までの画像がすでに存在するので、前日までの分と本日分の差分だけをダウンロードしています。

準備した画像を元に物体検出、特徴量抽出、近似最近傍探索Indexのビルドを行うのですが、大変重い処理なうえ、GPUも使うのでComposerのインスタンスで実行することはできません。そのため、別途GKEのクラスターを準備し、重い計算部分はそちらで計算しています。トップス、ボトムスなどのカテゴリ単位でPodを並列で動かし高速化しています。それでもまだ300万画像すべてを計算すると全体で24時間以上かかるので、日々改善を行っています。

特徴量は一部キャッシュとして保存しており、推論API側で利用しています。ZOZOTOWNの画像で画像検索する場合はこのキャッシュを流用できるため、推論APIのGPU資源を軽減させています。

ビルドしたIndexはモデルと同じGCSバケットに保存します。推論APIのGKEクラスターに対して近似最近傍探索PodのRollingUpdateを指示し、新しいIndexで動作する近似最近傍探索Podを起動します。ユーザーは更新された事に気が付かないまま、新しい結果を得ることができます。

エラーが起きた場合はSlackに通知しており、開発者がいつでも対応できるようにしています。

日々のバージョンアップ

ワークフローは日々改善を行っています。リリース後に行った改善を紹介します。

差分更新

当初、早急にリリースする事を優先しており、ワークフローを簡素化していたため毎日販売中の画像300万画像すべての特徴量を計算しIndexを作成していました。これでは大変時間がかかりコストも高いので「昨日と比較して新しく追加された画像だけ計算する」ように変更しました。コストは大幅に下がり、数時間で計算が終わるようになりました。

特徴量抽出のキャッシュ

開発当初はユーザーが画像をPostして、その画像から似た商品を返却する予定で開発をしていました。ですが、ZOZOTOWNにある商品と似た商品を返却した方が良さそうとのことで、ZOZOTOWNの画像を推論するようになりました。その場合、ZOZOTOWNに既にある商品画像はIndexを作成する際に特徴量抽出を一度行っているため、キャッシュに使えることが分かりRedisを追加しました。これにより、推論APIで物体検出と特徴量抽出を行わなくてもよい状況が増えたので、GPU負荷が大幅に下がりコストが削減できました。

※一部簡略化しています。すべてのリクエストでキャッシュを利用してるわけではありません。

新しい特徴量抽出モデルへの変更

現在トップスを始めとした8つのカテゴリに対応しており、日々新たなカテゴリに追加するため物体検出、特徴量抽出のモデルを開発しています。そこで特徴量抽出のモデルを更新する際にいくつかの問題が分かってきました。

  • 新たなカテゴリに対応すると、今までのカテゴリの精度を劣化させる可能性がある
  • 新たなカテゴリの学習のために今までのカテゴリ分も学習する必要があり、大幅に時間がかかる

ここで、カテゴリ毎に特徴量を抽出するようにモデルを作り変える決断をしました。さらに、学習を高速化をするためTPU x TensorFlowへの変更も行いました。モデルが大きく変わったため特徴量は現版（v1）と新版（v2）でまったく異なります。v2でIndexを1から作り直すのに数日かかり、その間もv1でAPIは稼働し続けないといけないため、v1, v2のIndex作成ワークフローを並列で行う必要があります。今回は推論APIのPod、キャッシュストレージなどをv2用に予め作り、v2のDAGがIndexを作り終わったタイミングで、推論APIのPodの向き先をv2に切り替えるという対応を行いました。

※一部簡略化しています。

これから

今後も機能追加、精度向上、コスト削減など様々な施策を実施していく予定です。ご期待ください。

失敗談

みんな大好き失敗談を紹介します。

画像ダウンロードしまくって負荷をかけてしまう

300万近くの画像はZOZOTOWNの画像サーバーから取得しています。300万画像を1つずつ取得していたら数週間かかるので、CloudFunctionsを使って無尽蔵に並列化して取得しました。案の定、サーバーに負荷をかけて怒られました。今は負荷をかけない程度の並列処理を実行しています。

開発期間の見積もりが難しい

実装自体は大したことないなと思っても、動作検証に大変時間がかかるため、実装に実質3日、検証に1か月みたいな事が平気で起きたりします。開発を効率化するためのスクリプト作り、安定した基盤づくり、可能な限り高速でタスクが完了するための高速化の工夫などがとても大事になります。

課題

これから解決していかなくてはならない課題たちです。

特定サービスが止まるとフローも止まる

例えばネットワークなどの都合で「特定のPythonパッケージがダウンロードできない」などが発生するとワークフローが止まってしまい、Indexを構築できません。これはクラウド時代には仕方がない部分ですが、ワークフロー中に外部依存する処理を減らしていく必要があります。

データはいつも同じではない

昨日まで動いていたのに動かなくなることが当たり前のようにあります。自分たちのコードのバグ、Airflowのバグ、商品画像数が増えたことでタイムアウトが発生した、依存している外部サービスに変更があった、など様々な要因があります。いつでも柔軟に対応できるよう、日々改善していくことが大事です。

ML部分のワークフロー化

現在Kubeflowを検証中で、モデルの精度が上がったら自動でデプロイされるような仕組みを検討しています。

まとめ

当初ワークフローの開発はそこまで大したこと無いだろうと考えていたのですが、API開発の10倍（個人の感想）くらい困難なものでした。これから新たに機械学習プロジェクトでワークフローを作る方は開発の初期段階からざっくり作り始めることをおすすめします。

私が携わった次期プロジェクトではこの知見が活き、迅速に開発・リリースができました。 lab.wear.jp

このワークフローは一人で作ったわけではありません。研究所の研究者、開発者、特にMLOpsチームの協力があり安定したワークフローに成熟していきました。今後もZOZOテクノロジーズのメンバーで様々なサービスを提供していきますので、今後とも宜しくお願いします。

登壇資料や関連サイト

techblog.yahoo.co.jp techblog.zozo.com

さいごに

ZOZOテクノロジーズでは福岡研究所のMLエンジニア、バックエンドエンジニア、MLOpsチーム（東京）のメンバーを募集しております。 https://hrmos.co/pages/zozo/jobs/0000029hrmos.co hrmos.co hrmos.co

カテゴリー