BigQueryデータでXGBoostモデルを訓練するスクリプト

表形式データの二値分類モデルの選定にあたって、高速な非ディープラーニングモデルである勾配ブースティング決定木(GBDT)をサクッと試してみようと思い、Jupyterノートブック用のスクリプトを作りました。BigQueryをデータソースとして、XGBoostアルゴリズムを用いる方法を記述します。

スクリプト全文

解説はいらないからスクリプトだけ試したいという方はこちらをコピペしてつかってください。
project_name、dataset_name、table_nameはご自身のプロジェクトに合わせて書き換えてください。

import xgboost as xgb
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import precision_score, recall_score, f1_score, log_loss
from google.cloud import bigquery

# BigQueryからデータを読み込む関数
def load_data_from_bigquery(query):
    client = bigquery.Client()
    query_job = client.query(query)
    df = query_job.to_dataframe()
    return df

def compute_metrics(labels, predictions, prediction_probs):
    precision = precision_score(labels, predictions, average='macro')
    recall = recall_score(labels, predictions, average='macro')
    f1 = f1_score(labels, predictions, average='macro')
    loss = log_loss(labels, prediction_probs)
    return {
        'precision': precision,
        'recall': recall,
        'f1': f1,
        'loss': loss
    }

# BigQueryでのクエリ
query = """
SELECT *
FROM `<project_name>.<dataset_name>.<table_name>`
"""

# データの読み込み
df = load_data_from_bigquery(query)

# 正解データ
y = df["reaction"]

# 入力データ
X = df.drop(columns=["reaction"], axis=1)

# データをトレーニングセットと検証セットに分割
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=1)

# XGBoostのモデルを訓練
model = xgb.XGBClassifier(eval_metric='logloss')

# パラメータグリッドの設定
param_grid = {
    'max_depth': [3, 4, 5],
    'learning_rate': [0.01, 0.1, 0.2],
    'n_estimators': [100, 200, 300],
    'subsample': [0.8, 0.9, 1.0]
}

# GridSearchCVの初期化
grid_search = GridSearchCV(estimator=model, param_grid=param_grid, cv=3, scoring='accuracy', verbose=1, n_jobs=-1)

# グリッドサーチの実行
grid_search.fit(X_train, y_train)

# 最適なパラメータの表示
print("Best parameters:", grid_search.best_params_)

# 最適なパラメータでのモデル
best_model = grid_search.best_estimator_

# 検証データでの予測
val_predictions = best_model.predict(X_val)
val_prediction_probs = best_model.predict_proba(X_val)

# トレーニングデータでの予測
train_predictions = best_model.predict(X_train)
train_prediction_probs = best_model.predict_proba(X_train)

# モデルの評価（検証データ）
val_metrics = compute_metrics(y_val, val_predictions, val_prediction_probs)
print("Optimized Validation Metrics:", val_metrics)

# モデルの評価（トレーニングデータ）
train_metrics = compute_metrics(y_train, train_predictions, train_prediction_probs)
print("Optimized Training Metrics:", train_metrics)

解説

BigQueryからのデータ読み込み

以前はCloud StorageにCSVファイルを保存していましたが、データ読み込みの遅さが学習の効率を低下させてしまっていたので、より高速にデータを取り扱えるBigQueryへ移行しました。

BigQueryクライアントの設定

from google.cloud import bigquery
client = bigquery.Client()

このコードは、Google Cloudの認証情報を用いてBigQueryクライアントを初期化します。認証情報は環境変数やGoogle Cloud SDKを通じて設定できます。

データのクエリと読み込み

def load_data_from_bigquery(query):
    query_job = client.query(query)
    df = query_job.to_dataframe()
    return df

この関数はSQLクエリを実行し、結果をPandasのDataFrame形式で返します。この方法で、大量のデータを効率的に処理できます。

XGBoostによるモデル訓練

XGBoostは、勾配ブースティングを応用した高性能な機械学習アルゴリズムで、分類や回帰問題に広く使用されています。

モデルの初期設定

import xgboost as xgb
model = xgb.XGBClassifier(eval_metric='logloss')

XGBClassifierはXGBoostの分類器で、ここでは損失関数としてログロスを使用しています。

データの分割

from sklearn.model_selection import train_test_split
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=1)

train_test_split関数を使用して、データを訓練セットと検証セットに分割します。

パラメータの最適化

from sklearn.model_selection import GridSearchCV
param_grid = {
    'max_depth': [3, 4, 5],
    'learning_rate': [0.01, 0.1, 0.2],
    'n_estimators': [100, 200, 300],
    'subsample': [0.8, 0.9, 1.0]
}
grid_search = GridSearchCV(estimator=model, param_grid=param_grid, cv=3, scoring='accuracy', verbose=1, n_jobs=-1)
grid_search.fit(X_train, y_train)

GridSearchCVは、指定されたパラメータの中で最適な組み合わせを見つけるために交差検証を行います。

モデルの評価

モデルのパフォーマンスは、検証データセット上での精度、再現率、F1スコア、ログロスを計算することで評価します。

def compute_metrics(labels, predictions, prediction_probs):
    from sklearn.metrics import precision_score, recall_score, f1_score, log_loss
    return {
        'precision': precision_score(labels, predictions, average='macro'),
        'recall': recall_score(labels, predictions, average='macro'),
        'f1': f1_score(labels, predictions, average='macro'),
        'loss': log_loss(labels, prediction_probs)
    }
val_metrics = compute_metrics(y_val, val_predictions, val_prediction_probs)
print("Optimized Validation Metrics:", val_metrics)

出力結果

ノートブックを実行すると、以下のように最適なパラメータとその時のモデル評価指標が表示されます。

Best parameters: {'learning_rate': 0.2, 'max_depth': 5, 'n_estimators': 300, 'subsample': 0.9}
Optimized Validation Metrics: {'precision': 0.8919952583956949, 'recall': 0.753797304483842, 'f1': 0.8078981867164722, 'loss': 0.014006406471894417}
Optimized Training Metrics: {'precision': 0.8969556573175115, 'recall': 0.7681976753444204, 'f1': 0.8199353049298048, 'loss': 0.012475375680566196}

補足

データソースをCloud Storageにする場合

場合によっては、データをBigQueryではなく、Google Cloud Storageから読み込む方が適切かもしれません。以下の関数は、Cloud StorageからCSVファイルを読み込み、PandasのDataFrameとして返します。load_data_from_bigquery関数と同様に使用できます。

from google.cloud import storage

def load_data_from_gcs(bucket_name, file_path):
    client = storage.Client()
    bucket = client.get_bucket(bucket_name)
    blob = bucket.blob(file_path)
    data = blob.download_as_text()
    df = pd.read_csv(io.StringIO(data), encoding='utf-8')
    return df

以下が使用例です。

bucket_name = '<backet-name>'
file_path = '<file-path>'

df = load_data_from_gcs(bucket_name, file_path)

LightGBMによるモデル訓練を行う場合

XGBoostと同様の手順で、LightGBMを使用することも可能です。以下のコードでは、XGBoostのXGBClassifierをLightGBMのLGBMClassifierに置き換えることで、同様の設定を適用しています。

import lightgbm as lgb
model = lgb.LGBMClassifier()

おわりに

BigQuery ML(BQML)による学習についてはいずれ書きます。