Vertex AI で AnimeGAN やってみた
この記事は Google Cloud Japan Advent Calendar 2021 の 8日目の記事です。
TL;DR,
本記事では Vertex AI Prediction のカスタムコンテナを使って、顔写真をアニメ顔にする PyTorch版 AnimeGANv2をサービングする方法を解説します。モデルそのものの詳細や、学習方法については取り上げません。
Vertex AI Prediction とは
Vertex AI は、機械学習に関わるワークフロー（学習、推論だけではなく、データ前処理やモデルのモニタリング、パイプラインなど）を効率的に構築する統合環境です。Vertex AI Prediction は、スケーラブルな推論の実行環境を自動的に構築できます。単にスケールする VM やコンテナ環境というわけではなく、以下のような ML Ops で必要になる機能を備えています。
- バッチ、オンライン推論
- リクエストの認証
- トラフィック分割
- モニタリング
Motivation — とまぁ、なんとなく面倒くさいところをやってくれる便利な機能のように見えますが、少し込み入ったこと（例：カスタムコンテナ）をやろうと思うとハマりどころがいくつかあります。本来、ML Eng. にモデルの開発に専念してもらい、面倒なことは俺らがやるぜ！なプロダクトなはずなので、そんなハマりどころに時間を取られたくないですよね。本記事ではそんな苦労どころをなるべく簡単にできるよう、工夫してみました。
STEP 0: AnimeGAN を Colab で使ってみる
AnimeGAN は先駆者が色々前準備をしてくれたおかげで、使うこと自体はとても簡単です。本編の Vertex AI で使う事とは全く関係ありませんが、Google Colaboratory で AnimeGAN を使ってみましょう。ここは知ってるよ、という方は STEP1 に進んでください。
Colab のノートブックは PyTorch 版のデモがあるので、それを実行してみるのが手っ取り早いです。重要なのは以下の3行です。
model = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", device=device).eval()face2paint = torch.hub.load("bryandlee/animegan2-pytorch:main", "face2paint", device=device)...im_out = face2paint(model, im_in, side_by_side=False)
上2行がモデルのロード、最後の行が実際の推論（アニメ化）です。とても簡単ですね！これは PyTorch Hub から学習済みのモデルをロードしているので、複雑なコードを書かなくても使え便利です。一旦ロードしたモデルはキャッシュされるので、実行の度にロードされることはありません。
STEP 1: サービング用のコードに AnimeGAN を埋め込む
今回使う AnimeGAN は PyTorch 版なので、そのまま Vertex AI にモデルを読み込むことはできません。そこで依存ライブラリを内包したカスタムコンテナに、モデルの推論コードを埋め込む必要があります。するとモデルの I/F どうするの？となるわけですが、HTTP サーバーを立てて Vertex AI の REST API の形式に合わせたリクエストを受け取れるようにします。HTTP サーバーの実装は何でもOKなのですが、特に複雑な機構は要りませんから、Flask で簡易的に実装し、リクエストの形式だけ合わせます。例えば、画像データを入力とするモデルの場合、以下のような JSON の形式で推論時の入力を受け取ります。
{'instances': [
{'b64': base64.b64encode(jpeg_data).decode()}
]}ここでは JPEG の画像データを BASE64 でエンコードして、JSON に埋め込んでいます。推論時にはHTTP Request Payload からこのデータを抽出し、BASE64 をデコードして画像データをモデルに入力する部分を自分で記述しなければなりません。ただこの形式は毎回変わるわけではないので、1度書いたらコードスニペットにでも入れておくとよいでしょう。以下のような感じでリクエストから画像の読み込み、モデル実行、結果を返すことができます。
STEP 2: ヘルスチェック用の関数を実装する
Vertex AI はカスタムコンテナの HTTP サーバーが正常に起動しているかを、指定のルートにリクエストを送ることで確認します。このルートが 200 を返さない限り、Vertex AI はこのコンテナに推論のリクエストをルーティングしません。モデルの読み込みや初期化に時間がかかる場合はここに判定を入れるとよいでしょう。今回は Flask を起動する前にモデルを読み込んでいるので、単純に 200 を返すコードのみを実装します。
@app.route("/healthcheck")
def healthcheck():
status_code = Response(status=200)
return status_codeなお、ヘルスチェックに失敗しても Vertex AI はコンテナを再起動せず、10秒毎にチェックして成功するのを待ちます。
STEP 3: コンテナをビルドし、コンテナリポジトリに登録する
次のステップで Vertex AI にモデルを読み込めるようにするため、Container Registry (gcr.io) に登録します。特にユニークな機構を入れる必要はないのですが、 ENTRYPOINT or CMD でHTTP サーバーを起動するのと、どの Port を開いているかは Dockerfile に記載する必要があります。必要最低限なところだけ書くと次のようになります。
FROM python:3.7-busterCOPY requirements.txt ./requirements.txt
RUN pip install -r requirements.txtCOPY app.py ./app.pyEXPOSE 8080ENTRYPOINT ["python3", "app.py"]
Port 番号は STEP4 でモデルを読み込むときに指定できます。ここではデフォルトの 8080 にしておきました。
ビルドと gcr.io への登録はこんな感じです。
$ docker build -t animegan .
$ docker tag animegan gcr.io/${PROJECT_ID}/animegan
$ docker push gcr.io/${PROJECT_ID}/animeganSTEP 4: Vertex AI で カスタムコンテナをモデルとして登録する
カスタムコンテナができたので、Vertex AI にモデルとして登録しましょう。Console を開いて、サイドメニューから「Vertex AI」→「Models」を開き、左上の IMPORT をクリックします。（初めて Vertex AI を開く場合は、先に「ENABLE VERTEX AI API」をクリックし、API を有効にしてください）
モデルの詳細を入力する画面が出てくるので、以下の必須項目を入力します。
- Model name — モデル名
- Container image — STEP3のgcr.ioを指定
- Prediction route —
/predictと記載 - Health route —
/healthcheckと記載
その他はデフォルト値でかまいません。終わったら、IMPORT ボタンを押すとモデルのインポートが開始します。数分かかるので、お茶でも入れて待ちましょう。
STEP 5: エンドポイントを作成する
モデルをインポートし終わったら、Vertex AI→Models のメニューに STEP4でインポートしたモデルがリストに表示されます。それをクリックし、「DEPLOY TO ENDPOINT」ボタンをクリックします。色々細かい設定ができるのですが、今回は使い方を覚えるだけなので以下の必須項目だけ入力し、「DEPLOY」ボタンをクリックします。
- Endpoint name — エンドポイント名
- Machine type — マシンタイプ。
n1-standard-2でとりあえずOK
エンドポイントの作成にも数分時間がかかります。これでサービング環境の出来上がりです！
リクエストを投げてみる
Vertex AI のエンドポイントは REST API として機能しているので、HTTP Post できるクライアントであればリクエストを投げることができます。認証も組み込まれているので、Vertex AI にアクセスできるサービスアカウントを利用するとよいでしょう。クライアントライブラリを使うと次のように簡単に認証も含めたリクエストを送れます。
エンドポイントでは以下のように秒間リクエスト数やレイテンシーをモニタリングできます。
おまけ：動画の場合は各フレームを変換するので時間がかかりますが、並列でリクエストを投げれば時間を短縮することもできますね。
後片付け
作成したリソースは削除しないと課金が続いてしまうので、使わない場合は削除しましょう。特にエンドポイントはノードが立ち上がっているので、早めに消した方が無駄になりません。他に使っているプロダクトがない場合はプロジェクトごと消すと簡単です。
まとめ
本記事では、 PyTorch 版 AnimeGAN をカスタムコンテナを使って Vertex AI でサービングする手順を紹介しました。元々のコードから、少し追加するだけで環境が立ち上がるのが体感できたのではないかと思います。
参照
本記事のソースコード
TorchServe をカスタムコンテナで使う手順https://cloud.google.com/blog/ja/topics/developers-practitioners/pytorch-google-cloud-how-deploy-pytorch-models-vertex-ai
PyTorch 版 AnimeGAN
TensorFlow 版 AnimeGAN
Disclaimer この記事は個人的なものです｡ここで述べられていることは私の個人的な意見に基づくものであり､私の雇用者には関係はありません｡
