🌊

初手BQMLのメリデメ

2024/04/26に公開

機械学習

概要

スタートアップで初めて機械学習を用いた予測システムをデプロイするときにBigQueryMLを採用したので、やったこととメリデメを整理する。

背景・課題

機械学習を用いて需要予測を行うシステムをデプロイしたい
予測システムの導入は初めてなので、現在簡単にデプロイできるような体制はない
dbtでデータ基盤を管理している
予測モデルがワークするか確かめられるまでは、リッチなシステム構築は避けたい

やったこと

BigQueryMLを用いて、学習・モデルデプロイ・予測を行うシステムを組んだ
https://cloud.google.com/bigquery/docs/bqml-introduction?hl=ja
あらかじめデータ基盤やジョブスケジューリング機能を作れていたので、実装自体は調査含めて4日くらいでできた

システム構成

・データ処理はdbtで管理する。これにより異常データなどが入ったときに予測を実行せずに処理を止めることができる
・github actionsでスケジューリングを実行する
・BQMLで学習を実行すると自動でvertex aiにデプロイされ、モデルのバージョニングや指標の監視などができるようになる

BQMLの導入

dbt_mlの利用

dbt上でBQMLを利用するために便利なパッケージ「dbt_ml」を利用した。

dbt_mlについては他に色々記事がありそうだったので、そちらを参照していただきたい。
直近更新があまりないのと、使用実績が少なそうなのがネックだったが、最悪なくなっても自前で実装できそうなくらいのボリュームだったので一旦良しとした。

学習の実装

train.sql

{{
    config({
        "materialized": 'model'
        , "tags": ["monthly","model"] -- 月次で再学習するためmonthlyタグを指定
        , "post-hook": "{{ dbt_ml.model_audit() }}"
        , "enabled": true
        , "full_refresh": false
        , "ml_config": {
            'model_registry': "vertex_ai",
            'vertex_ai_model_id': "{{MDOEL_ID}}", -- vertex ai上で登録するID
            'model_type': 'BOOSTED_TREE_CLASSIFIER',
            'booster_type': 'GBTREE',
            -- 任意のパラメータを指定
        }
    }
    )
}}

{% set features = [
  "feature1"
    , "feature2"
    -- 任意の特徴量を追加
] %}

select {% for feature in features %}{{ feature }},{% endfor %}
from {{ ref('feature_table') }}

model_registryでvertex aiを指定することで、学習完了時にvertex aiにデプロイが実行されるようになる。
デプロイされたモデルは、モデルレジストリから確認できる。

予測の実装

predict.sql

-- depends_on: {{ ref('train') }}
{{
    config(
        materialized='incremental' -- 予測のたびに上書きされるようにincrementalにした
        , unique_key=['feature1','feature2']
        , tags=["weekly"] -- 予測は週次で行うためweeklyタグを指定
        , on_schema_change="append_new_columns"
    )
}}
{% set threshold_query %}
select
  array_agg(threshold order by total_abs_error asc)[offset(0)]
from
  {{ref("_get_threshold")}}
{% endset %}

{% set result = run_query(threshold_query) %}

{% if execute %}
{% set threshold = result.columns[0].values() %}
{% else %}
{% set threshold = [] %}
{% endif %}

{{print("optimized threshold for prediction: " ~ threshold[0])}}

with
import_features as (
  select
    *
  from
    {{ ref('feature_table') }}
  where
    week = xxx -- 予測したい週に限定
)

,
predicted as (
  select
    *
  from
    {{ dbt_ml.predict(ref('train'), "import_features", threshold[0]) }}
)

select * from predicted

今回は分類を行うモデルであるが、分類に使う閾値はpredict関数に指定する必要がある。
自社の要件に応じた評価指標を最適化する閾値を、予測の実行時に自動的に決定してセットするような実装をした。

良かったこと・微妙だったこと

良かったこと

とにかく実装が速い
- 既存のdbtスタックの上にそのまま乗せられたのがかなり大きい
データをBigQuery外に出すことなく、学習・予測が完了した
学習結果がvertex aiで自動的に表示されるため、性能指標のモニタリングについてゼロから開発する必要がない
BigQueryMLはONNXなどにも対応しているため、今後より複雑なモデルを扱いたい場合は、学習だけを別のシステムで行ってvertex aiにアップロードし、予測はモデルファイルをロードする形に書き直せば対応できそう

微妙だったこと

モデルの実験はLightGBMで行ったが、BigQueryML上だとLightGBMがそのまま使えない？
n_iterationが上限1000回までしか指定できないので、学習が収束するまでに終わってしまう

まとめ

とりあえずあまり時間をかけずにサクッとモデルデプロイを行いたい場合は、BigQueryMLが良さそう。