G-gen の片岩です。当記事では Vertex AI Custom Training を使用して、機械学習モデルのトレーニングから推論まで実行する方法を紹介します。Vertex AI Custom Training を使うことで、クラウド上のフルマネージド環境で機械学習モデルをトレーニングできます。

はじめに

機械学習モデルのトレーニング手法

Google Cloud のマネージドサービスを使って機械学習モデルをトレーニングする方法は、大きく分けて3つあります。

手法	向いているケース
Vertex AI AutoML	データを用意するだけで、少ない手順でモデルをトレーニングしたい時
BigQuery ML	BigQuery のデータを使用して SQL でモデルをトレーニングしたい時
Vertex AI Custom Training	モデルの開発環境や工程を細かく制御したい時

当記事では、最も柔軟性の高い Vertex AI Custom Training に焦点を当て、モデルのトレーニングおよび推論を実行する方法をご紹介します。

Vertex AI についての解説、Vertex AI AutoML や BigQuery ML についての詳細は以下の記事を参照してください。

blog.g-gen.co.jp

Vertex AI Custom Training とは

Vertex AI Custom Training は、PyTorch、TensorFlow、Scikit-learn や XGBoost といったフレームワークを実行できるフルマネージドサービスです。

ローカルで開発した Python コードを Docker コンテナとして実行するため、サーバーの起動や停止といったインフラ管理を意識せず、コードの自由度とクラウドのスケーラビリティを両立できます。

AutoML では対応しきれない細かいチューニングや、独自のアルゴリズムを実装したいエンジニアにとって、有用な選択肢です。

参考 : カスタムトレーニングの初心者向けガイド

構成図

当記事で紹介する手順に関する構成図は、以下のとおりです。環境構築の負荷を軽減するため、ソースコードの作成や Python 実行環境に Colab Enterprise を使用します。

Colab Enterprise については以下の記事を参照してください。

blog.g-gen.co.jp

初期設定

はじめに、ライブラリのインストールと環境変数の設定を行います。

今後、別のモデルを作成する時にバケットを使い回せるようにするため、Cloud Storage はバケットの中にディレクトリを作成します。

# 必要なライブラリのインストール
!pip install google-cloud-aiplatform xgboost scikit-learn pandas -q
  
# プロジェクトとリージョンの設定
# ※ ご自身の環境に合わせて書き換えてください
PROJECT_ID = "your-project-id"
LOCATION = "asia-northeast1"
  
# バケットとフォルダの定義
ROOT_BUCKET = "gs://your-bucket"
EXPERIMENT_NAME = "california-housing-xgb-v1"
WORK_DIR = f"{ROOT_BUCKET}/{EXPERIMENT_NAME}"
  
# Vertex AI SDK の初期化
from google.cloud import aiplatform
# staging_bucket を指定すると、自動生成されるファイルもこのフォルダに整理されます
aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=WORK_DIR)
  
# バケットが存在しない場合のみ作成
!gsutil mb -l {LOCATION} {ROOT_BUCKET}

データの準備と分割

データは機械学習デモにしばしば使用される California Housing（カリフォルニアの住宅価格）を使用します。このデータセットは、scikit-learn に含まれています。

モデルの作成に使用する学習データと推論に使用する推論データに分割し、Cloud Storage に保存します。

from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
import pandas as pd
  
# データのロード
data = fetch_california_housing(as_frame=True)
df = data.frame
  
# 学習データと推論データに 90:10 の割合で分割
train_full_df, test_df = train_test_split(df, test_size=0.1, random_state=42)
  
# 学習データの保存
train_filename = "train.csv"
train_full_df.to_csv(train_filename, index=False)
  
# 推論データの保存
test_filename = "test.csv"
test_df.to_csv(test_filename, index=False)
  
# GCS へアップロード (フォルダ /data 配下へ)
!gsutil cp {train_filename} {WORK_DIR}/data/{train_filename}
!gsutil cp {test_filename} {WORK_DIR}/data/{test_filename}
  
print(f"学習データ: {WORK_DIR}/data/{train_filename}")
print(f"推論データ: {WORK_DIR}/data/{test_filename}")

学習スクリプトの作成

Vertex AI Custom Training として実行する Python スクリプトを作成します。今回は機械学習アルゴリズムに XGBoost を採用します。

環境変数 AIP_MODEL_DIR で与えられた Cloud Storage に作成したモデルを保存すると、自動的に Vertex AI Model Registry にアップロードされます。

%%writefile task.py
import argparse
import pandas as pd
import xgboost as xgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
import os
from google.cloud import storage
from urllib.parse import urlparse
  
# 引数の受け取り
parser = argparse.ArgumentParser()
parser.add_argument('--train-data-uri', dest='train_data_uri', type=str, required=True)
args = parser.parse_args()
  
# GCS から学習用データセットをダウンロード
parsed_url = urlparse(args.train_data_uri)
bucket_name = parsed_url.netloc
blob_path = parsed_url.path.lstrip("/")
local_filename = "downloaded_train.csv"
  
client = storage.Client()
bucket = client.bucket(bucket_name)
blob = bucket.blob(blob_path)
print(f"Downloading training data from: {args.train_data_uri}", flush=True)
blob.download_to_filename(local_filename)
  
# 学習用データセットを読み込む
df = pd.read_csv(args.train_data_uri)
  
# 特徴量とターゲットに分離し、Numpy Array に変換
target_col = "MedHouseVal"
X = df.drop(columns=[target_col]).values
y = df[target_col].values
  
# トレーニングに必要な学習データと評価データに分割する
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.1, random_state=42)
  
# XGBoost モデルの構築と学習
model = xgb.XGBRegressor(
    n_estimators=100, 
    learning_rate=0.1, 
    max_depth=5, 
    random_state=42,
    objective='reg:squarederror'
)
  
model.fit(
    X_train, y_train,
    eval_set=[(X_val, y_val)],
    verbose=True
)
  
# 評価スコアを Cloud Logging に出力する
score = model.score(X_val, y_val)
mse = mean_squared_error(y_val, model.predict(X_val))
print(f"Validation Score (R^2): {score:.4f}", flush=True)
print(f"Validation MSE: {mse:.4f}", flush=True)
  
# モデルの保存
model_filename = "model.bst"
model.save_model(model_filename)
  
aip_model_dir = os.getenv("AIP_MODEL_DIR")
if aip_model_dir:
    print(f"Uploading model to: {aip_model_dir}", flush=True)
    parsed_url = urlparse(aip_model_dir)
    bucket_name = parsed_url.netloc
    blob_path = parsed_url.path.lstrip("/").rstrip("/") + "/" + model_filename
    
    # Cloud Storage にアップロード
    bucket = client.bucket(bucket_name)
    blob = bucket.blob(blob_path)
    blob.upload_from_filename(model_filename)
    
    print("Model upload completed.", flush=True)

学習ジョブの実行

作成したスクリプトをコンテナで実行します。今回は Google Cloud が提供するビルド済みコンテナを使用するため、Dockerfile の作成は不要です。

また、軽量なモデルのため n1-standard-4 マシンタイプを使用しますが、巨大なモデルをトレーニングしたい場合は GPU を使用することで高速に処理できます。

# ジョブの定義
job = aiplatform.CustomTrainingJob(
    display_name="california-housing-xgb-job",
    script_path="task.py",
    container_uri="asia-docker.pkg.dev/vertex-ai/training/xgboost-cpu.2-1:latest",
    requirements=["pandas", "google-cloud-storage"],
    staging_bucket=WORK_DIR,
    model_serving_container_image_uri="asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.2-1:latest"
)
  
# ジョブの実行
print("ジョブを送信しました。完了まで5分ほどお待ちください...")
model = job.run(
    machine_type="n1-standard-4",
    replica_count=1,
    model_display_name="california-housing-xgb-model",
    # GCS上の学習データパスを引数で渡す
    args=[f"--train-data-uri={WORK_DIR}/data/train.csv"]
)

ジョブ終了後に Cloud Storage にモデルが保存されていることを確認できます。

Google Cloud コンソールのモデルレジストリの画面でモデルが作成されていることを確認できます。

推論の実行

本番運用では Vertex AI に備わっているオンライン推論やバッチ推論を使用して大量のリクエストを捌きますが、ここでは Colab Enterprise の実行環境上で推論を実行します。

# ---------------------------------------
import xgboost as xgb
import pandas as pd
import matplotlib.pyplot as plt
  
# --- モデルの場所 ---
model_gcs_uri = model.uri
# GCSのパスを直接指定する場合
# model_gcs_uri = WORK_DIR + "/aiplatform-custom-training-xxxxx/model"
print(f"参照モデル: {model_gcs_uri}")
  
# 推論データのダウンロード
!gsutil cp {model_gcs_uri}/model.bst .
!gsutil cp {WORK_DIR}/data/test.csv .
  
print("ダウンロード完了。推論を開始します...")
  
# モデルの読み込み
local_model = xgb.XGBRegressor()
local_model.load_model("model.bst")
  
# データの読み込み
df_test = pd.read_csv("test.csv")
  
# 推論用に「正解列」を削除したデータフレームを作ります
target_col = "MedHouseVal"
X_test = df_test.drop(columns=[target_col])
  
# 推論の実行
# .values をつけて数値行列として渡します
predictions = local_model.predict(X_test.values)
  
# --- 結果の表示 ---
print("\n=== 推論結果 (最初の10件) ===")
for i, pred in enumerate(predictions[:10]):
    print(f"データ{i}: 予測価格 {pred:.4f}")

推論結果の確認

最後に予測結果と正解データを照合して評価指標を計算し散布図を描画します。

from sklearn.metrics import r2_score, mean_squared_error
import numpy as np
  
# 正解データの抽出
y_true = df_test["MedHouseVal"]
  
# 精度評価 (答え合わせ)
r2 = r2_score(y_true, predictions)
rmse = np.sqrt(mean_squared_error(y_true, predictions))
  
print("-" * 30)
print(f"評価結果 (データ数: {len(y_true)}件)")
print(f"R2 Score (決定係数): {r2:.4f}")
print(f"RMSE (誤差の大きさ): {rmse:.4f}")
print("-" * 30)
  
# 結果の可視化 (散布図)
plt.figure(figsize=(8, 8))
plt.scatter(y_true, predictions, alpha=0.5, color='blue', label='Predictions')
# 理想線 (y=x)
min_val = min(y_true.min(), predictions.min())
max_val = max(y_true.max(), predictions.max())
plt.plot([min_val, max_val], [min_val, max_val], 'r--', label='Ideal Fit')
  
plt.title(f"Actual vs Predicted (R2: {r2:.3f})")
plt.xlabel("Actual Price")
plt.ylabel("Predicted Price")
plt.legend()
plt.grid(True)
plt.show()