はじめに

AI Lab Creative Researchグループでリサーチエンジニアをしている大田（@ciela）です。主にAI LabリサーチサイエンティストやAI事業本部プロダクト開発チームに対する機械学習エンジニアリングサポートを行っています。

昨年末、弊社よりFEED自動承認AIというサービスがリリースされました。このサービスはGoogle Merchant Center（GMC）にて不承認となった商品画像から不承認理由に該当する領域を除去することで、承認率を向上することのできるサービスです。内部では機械学習モデルが利用されており、不承認画像それぞれに特化した画像編集が施される仕組みとなっています。詳しくは下記プレスリリースをご参照ください。

AIによってGoogleショッピング広告における承認率の大幅改善を図るサービス「FEED自動承認AI」を開発・提供開始

FEED自動承認AIのモデル開発にはAI Labからも下田と私大田が関わっているのですが、主に下田がモデルの開発と学習を行い、私がシステム化してプロダクトチームへ提供する、といった役割分担でプロダクトサポートを進めています。下田はリサーチサイエンティストとしてモデル精度向上のための試行錯誤を継続していく必要があり、並行してシステムへモデルコードをデプロイしていくことがなかなか難しい立場です。しかし、サービスとして実用化するためにはモデルの性能がある程度見極められたタイミングでデプロイしていかなければなりません。そこで私がシステム化に関するエンジニアリングサポートに入ることとなりました。

本記事では、FEED自動承認AIリリースまでに私が行ってきた実験用モデルをシステム導入する上での取り組みについて記したいと思います。ほぼPyTorchとAWS SageMakerに特化した内容ではあるものの、研究で組み立てている機械学習モデルを実用化していく上での経験談としてお役に立てれば幸いです。

モデルリファクタリング

今回のFEED自動承認AIモデルの入力・出力は下記のとおりです:

• 入力: GMCにて不承認となった商品画像
• 出力: 不承認理由に該当する領域が編集・除去された画像

モデル内部ではさらにいくつか特定の処理に特化したサブモデル・画像処理で構成されており、PyTorch、OpenCV、Pillowを利用して実装されています。重みファイルは3つあり合計850MBほどの容量です。実験用にほぼ一直線になっていたコードをモジュール・クラスに分割して構造化し、推論に特化した形に仕立てあげていったのですが、その際大事にしていたポイントをまとめてみます。

不要な計算・重みの削除

モデルコードを編集していくにあたり、まずは無駄なレイヤーや計算が含まれていないかを確認するところからはじめました。PyTorchの場合、 nn.Module クラスを拡張したレイヤークラスを組み合わせてモデルを表現していくことが多いかと思われますが、試行錯誤の過程で現在利用されていないレイヤーがコードに残っていたり、逆に既にコードから削除済みのレイヤーの重みがファイルに残っていたりします。重みファイルをロードする際の nn.Module.load_state_dict のオプションで strict=False と指定することで、こういった現行コードと重みファイルのミスマッチエラーを抑制しているコードをこれまでよく見かけてきました。これは実験においては何かと役に立つオプションなのですが、重みファイルの容量がかさばってしまったままとなってしまいます。

モデルコードと重みファイル内の重みデータに過不足がないようにするために、モデルクラスから不要なレイヤーを削除した後、一度 strict=False 付きで重みをロードしそのまま torch.save にかけることで新しい重みファイルに出力しました。結果、重みファイルのサイズを約200MB削減することができ、モデルのロード時間やGPUメモリ消費量の削減にも繋がりました。さらに、新しい重みファイルは strict=False オプションを付けずに読み込むことが可能となりコードもすっきりすることができました。

CPU/GPU間データ移動の除去

GPU上のモデルでデータを処理する際、理想的には一度GPUに送ったデータは出力時までずっとGPU上で計算させたいものです。今回のモデルにおいては、サブモデル同士の出力と入力の間にOpenCVとPillowを用いた画像処理が挟まれており、このためにCPU/GPU間でのデータ移動が発生していました。そのため、この処理をGPU上で行うように置き換えデータ移動が発生しないようにできないかを考えました。

嬉しいことに最近のtorchvisionではOpenCVやPillowとほぼ似たようなパラメータを扱えるAPIで torch.Tensor に対する画像処理が行えるようになっています。無事PyTorchのみで画像処理部分をCUDA化することができ、CPU/GPU間データ移動のオーバーヘッドを減らすことによって処理を高速化・効率化することができました。

システムへのデプロイ

FEED自動承認AIのシステムはAWS上に構築されています。そのため、今回のモデルはAWSの機械学習統合プラットフォームであるSageMakerにデプロイしていくことになりました。SageMakerは以前から何度か研究用途で利用してはいたものの、プロダクトシステムへと提供される推論エンドポイント構築のために利用するのは以前参加したワークショップ以来でした。

学習済みモデルの推論エンドポイント化

モデルの開発・学習からSageMaker上で行っている場合は推論エンドポイントへのデプロイが非常にスムーズなことをワークショップでも体験させてもらっていたのですが、今回は学習済みモデルが既に手元にある状態です。そのため、下記のSageMakerドキュメントページに記載されている外部モデルのデプロイ方式を採用しようと考えました。

I have a PyTorch model that I trained outside of SageMaker, and I want to deploy it to a SageMaker endpoint
Use PyTorch with Amazon SageMaker

しかしながら、2022年2月現在、このデプロイ方式に記されているリンク先でのアンカリングが効いておらず、ページ内の該当項目にナビゲートされないという状況に陥ってしまいます…。おそらくリンク先のSphinxドキュメントのみ更新されたままなのでしょう。手探りで色々調べたり普段からお世話になっているAWS Solution Architectの方々に尋ねたりしつつ、下記のImmersion Dayのリポジトリに含まれているsickit-learnモデル向けのハンズオンノートブックに従って作業することで、なんとか推論エンドポイントとしてデプロイすることができました。

Lab: Bring your own custom container with Amazon SageMaker

もし今回のように既に学習済みモデルがある状態でのSageMaker推論エンドポイントへのデプロイを考えている方は参考にしていただけたらと思います。この作業において重要だと感じたポイントを下記に記しておきます。

• 重みファイルは model.tar.gz という名前でアーカイブしてS3にアップロードしておく
• serve というファイル名のスクリプトを実行可能状態でDockerイメージ内に含めパスを通す（今回は推論なので serve のみだが学習の場合は train も準備する）
• ヘルスチェック用の /ping と実際の推論を行う /invocations の2種類のAPIを準備したWebサーバを記述し serve スクリプトから起動する

ローカルでの動作確認

推論エンドポイントとしてのモデルの動作確認をWebリクエスト経由で行う場合、最終的にはSageMakerにデプロイして確認しなければなりません。上記Immersion Dayのノートブック通りに作業した場合、FlaskとGunicornを使ったWebサーバが推論エンドポイントとなりますが、それらの設定ミスなどの些細なことで発生したエラーによってデプロイにかかった時間を無駄にするのはもったいありません。

SageMakerではローカルモードというローカルマシン上にSageMaker環境を構築できる機能も提供しています。SageMakerにデプロイする前には必ずこのローカルモードで推論エンドポイントの動作確認を事前に行うことで、デプロイ後の手戻りをなるべく少なくするように心がけていました。

ノートブックインスタンスのトレーニングに Amazon SageMaker のローカルモードを使用する
オンプレミス環境から Amazon SageMaker を利用する

その他

基礎的ですがシステムへの導入にあたって他に考慮していた項目を挙げておきます:

• モデル推論時には nn.Module.eval で推論モードにし、 torch.no_grad でガードして勾配計算をカットする
• SageMakerインスタンスタイプに合わせてGunicornパラメータ（ワーカー数、スレッド数など）を調節する
• SageMakerバッチ推論における同時並列リクエスト数を調節する

特にGunicornパラメータ周りについては、パフォーマンス確認の際にいちいちコードから変更する必要がないように、SageMaker推論エンドポイント起動時に渡すことのできる環境変数値を介して調節が出来るようにスクリプトを記述しました。

おわりに

今回の記事ではFEED自動承認AIのリリースに向けて、私が実験用コードを推論エンドポイントとして実用化するまでの取り組みについてまとめてみました。下田とうまく分業してAI Labから事業への貢献事例をつくれたのではないかと思っています。

ちなみにJIT化による高速化などモデル自体のパフォーマンス向上や、PyTorch推論アプリケーションサーバであるTorchServeについての調査など、今後FEED自動承認AIモデルの推論エンドポイントに関して検証したいタスクはたくさんあります。もしAI Labにてプロダクトサポートや研究のためのエンジニアリングに興味のある方がいらっしゃれば、どうぞお気軽にご連絡いただければと思います。