Vertex AIのバッチ予測で入出力の対象をBigQueryにする
おはこんハロチャオ～！何者(なにもん)なんじゃ？じょんすみすです。
Vertex AIではModel Registryに登録されたモデルを使ってバッチ推論を実行できます。 バッチ推論を行うために、必要な設定やCloud Storageを利用する際の入力に関しては以前紹介しました。
今回は入力や出力をBigQueryとして動かしてみます。
Model Registryとエンドポイント, バッチ推論
まずは、簡単に機能をおさらいしておきましょう。
Vertex AIではModel Registryに以下の2つの情報を登録しておきます。
- モデルのCloud Storage上のパス
- 推論処理を実行するためのコンテナのパス
- コンテナ実行時のコマンド
これらの情報に基づいて、エンドポイントを立てる際にはコンテナを起動し、その中から処理を実施する仕組みになっています。
この際、コンテナ内には推論をするために処理が実装されたコードを含めますが、HTTPでリクエストを受け付ける仕組みになっている必要があります。
環境変数 AIP_HEALTH_ROUTE, AIP_PREDICT_ROUTE で渡されるパスへのそれぞれGETとPOSTでリクエストを受け付ける仕組みを実装しておけばいいわけです。
推論に利用するモデルのパスも環境変数 AIP_STORAGE_URI として渡されるため、 AIP_PREDICT_ROUTE へのリクエストがあった際に推論結果を返すコードを実装します。
リクエストはJSON形式で、 instances というキーに対して配列で複数の値が渡されます。これらが推論対象となるデータです。
推論結果も同様にJSON形式で返します。 predictions というキーの中なに入力で与えられたそれぞれに対する推論結果を配列で返します。
これらの配列の中身自体は、どのような形式で与えるかは仕様としては決まっていませんので、JSONで記述可能な任意の形式で与えることが可能です。 よくある形式としては、入力は特徴を列挙した配列や特徴量名と値のkey-value形式となるでしょう。 出力に関しても、推論結果だけを返す場合や、それぞれの予測ラベルに対する確信度を返すなどが考えられます。
バッチ推論でも同様にModel Registryに登録されているモデルを使って推論を実施します。 バッチ推論の入出力はCloud StorageやBigQueryが入出力となります。 これらは一致している必要はなく、入力はBigQueryで出力はCloud Storageとしたりその逆にすることも可能です。
バッチ推論でも同じモデルを利用することから、内部ではHTTPによるリクエストで処理を行いますが、利用者が意識しなくて以下の部分を実施してくれます。
- 入力元からデータを取得
- エンドポイントにリクエスト投げて推論結果を取得
- 取得した結果を出力先に書き込む
これらをVertex AI側でやってくれているため、同じモデルでエンドポイントの作成とバッチ推論の両方を行えるわけです。
ただし、自動的にやってくれているため、利用者側で前後の処理を挟むことはできません。 そのため、特に入力に関してはバッチ推論の仕様に合わせて内部の処理が実行できるように実装しておく必要があります。
BigQueryを入出力として利用する
BigQueryを入出力として利用する方法はそれほど難しくありません。
コンソール上で行う時は、以下のようにそれぞれでBigQueryを選択して対象のテーブルを指定します。
また、Pyhton SDKから呼び出す際も引数でBigQuery用のものを指定するのみとなります。
# aiplatform.Model("<モデルID>")で既存のモデルを取得するか
# 学習ジョブのrunの結果を受け取る
model = ...
# 入力はbigquery_source、出力はbigquery_destination_prefixで指定
# gcs_source, gcs_distination_prefixは指定しない
model.batch_predict(
job_display_name="titanic-predict",
bigquery_source="bq://<project id>.<dataset>.titanic_test",
bigquery_destination_prefix="bq://<project id>.<dataset>.titanic_batch_predict",
machine_type="n1-standard-4"
)これだけで利用が可能ですが、注意点としては入出力の仕様となります。
入力は1レコードごとにカラム名をkey、その値をvalueとした形式で、instancesの中身がその配列になっているような形式となります。 一方、出力はpredictionの中身が1レコードごとに順に付与されます。
以下のように返ってきたJSONが付与されたテーブルが出力されます。 (Titanicデータセットに対して画面表示用にいくつかの特徴量に絞ったもので推論を実施しています)
上述の記載の通り、BigQueryからの入力はテーブル名を指定する形式になっています。 そのため、SELECT文を使って何らかの処理を行ったものを渡すということができません。 公式に提供されているサンプルでは、リクエスト時に対象とするカラムや除外するカラムを設定する例が記載されていますが、なんらかの特徴量エンジニアリングを行うなどはここではできません。
これを踏まえ、他での実用性の観点からも見ると推論時の処理として学習時と同じ特徴量エンジニアリングを行う部分を内部のコードで実装しておくといいでしょう。
実際の使いどころは？
どのような利用シーンが考えられるでしょうか？ ここまでで、入出力双方でBigQueryを利用する場合は、BigQuery MLでもいいようにも見えます。
様々なケースが考えられますが、ここでは2つほど例を挙げてみたいと思います。
BigQuery MLでは利用できない手法を使いたい
BigQyery MLには豊富なアルゴリズムやAutoMLが実装されており、それらが簡単に利用可能です。 また、TensorflowやXGBoostの学習済みモデルをインポートして推論に利用することも可能です。
一件、これだけあればだいたいのことに使えそうですが、 例えばOSSで公開されているライブラリを利用しており、それがPyTorchで実装されているなどのインポートできないモデルを利用したい場合があります。
Vertex AIであれば、コンテナ上で任意のライブラリを実行できるのでこういったケースにも対応できます。
入力か出力のどちらかがCloud Storageである必要がある
推論は定期的なバッチ処理で問題ないが、出力結果を他と連携させて使うのにCloud Storageに置いておく必要がある。 入力データは構造化データではないのでCloud Storageに保存するようにしているが、それに対する推論結果はBigQueryで管理しておきたい。
などの要望があるケースもあるでしょう。 そういった場合に、追加で連携させる仕組みを作らなくてよくなるのは便利ですね。
おわりに
今回は、Vertex AIのバッチ推論においてBigQueryを入出力に利用する方法を紹介しました。
Model Registryに登録済みのモデルがあればお手軽に利用できるので、ぜひ試してみてください。
9
