G-gen の片岩です。当記事では Vertex AI Custom Training において カスタムコンテナ を使用し、標準では提供されていない LightGBM モデルの学習から 寄与度（SHAP）の出力 まで実行する方法を紹介します。

はじめに

ビルド済みコンテナとカスタムコンテナの使い分け

Vertex AI Custom Training には学習ジョブを実行するためのコンテナイメージとして、大きく 2 つの選択肢があります。

コンテナの種類	特徴	向いているケース
ビルド済みコンテナ	Google Cloud が用意したイメージ	XGBoost や TensorFlow など、標準的なフレームワークをすぐに使いたい時
カスタムコンテナ	自分で Dockerfile を書いて作成するイメージ	LightGBM など未提供のライブラリを使いたい時や、独自の処理を組み込みたい時

カスタムコンテナの利点

ビルド済みコンテナでもジョブ実行時の引数に requirements=["lightgbm", "shap"] のように指定することでライブラリを追加できます。ビルド済みコンテナについては以下の記事を参照してください。

blog.g-gen.co.jp

しかし実務の本番運用において、実行時にライブラリを動的にインストールすることは、以下のデメリットがあります。

1 点目は、環境の再現性が低下することです。

ジョブを実行するたびにインターネットから最新のパッケージを取得するため、依存ライブラリのバージョンが上がったために突然ジョブが落ちたり、学習結果が変わってしまうといった、本番運用で避けたいリスクを招きます。

2 点目は、実行のたびにオーバーヘッドが発生することです。

毎回ライブラリをダウンロードしてインストールする処理が走るため、余計な待ち時間が発生します。

カスタムコンテナを利用することにより、上記のデメリットを回避できます。

参考 : カスタムコンテナの概要

構成図

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

初期設定

はじめにライブラリのインストールと環境変数の設定を行います。今回は可視化や解釈のためのライブラリ（seaborn 、 shap）も追加します。

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

データの準備と分割

データは機械学習デモで使用されるダイヤモンドの価格データを使用します。このデータはカラットなどの数値データや、カットや色といったカテゴリ変数を含みます。

学習データと推論データに分割して Cloud Storage に保存します。

import seaborn as sns
from sklearn.model_selection import train_test_split
import pandas as pd
  
# データのロード (~54,000行)
df = sns.load_dataset('diamonds')

# 文字列カラムを 'category' 型に変換
cat_cols = ['cut', 'color', 'clarity']
for col in cat_cols:
    df[col] = df[col].astype('category')
  
# 学習データと推論データに 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 へアップロード
!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}")

カスタムコンテナの準備

ディレクトリとリポジトリの準備

Colab Enterprise 上に作業ディレクトリを用意し、Google Cloud 上に完成したコンテナの保存先となる Artifact Registry のリポジトリを作成します。

# 作業用ディレクトリの作成
!mkdir -p custom_container
  
# Artifact Registry にリポジトリを作成 (初回のみ)
!gcloud artifacts repositories create custom-training-repo \
    --repository-format=docker \
    --location={LOCATION} \
    --description="Custom Training Repository" || true

学習スクリプトの作成

コンテナ内で実行される task.py を作成します。今回はモデルの学習だけでなく、過学習を確認するための学習曲線と、予測の根拠を説明するための寄与度の画像を生成し、モデルと一緒に Cloud Storage へアップロードする処理を組み込みます。

%%writefile custom_container/task.py
import argparse
import os
import pandas as pd
import lightgbm as lgb
import shap
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from google.cloud import storage
from urllib.parse import urlparse
import warnings
warnings.filterwarnings('ignore')
  
parser = argparse.ArgumentParser()
parser.add_argument('--train-data-uri', dest='train_data_uri', type=str, required=True)
args = parser.parse_args()
  
# --- GCS ダウンロード / アップロード用の関数 ---
def download_from_gcs(gcs_uri, local_file):
    parsed_url = urlparse(gcs_uri)
    client = storage.Client()
    bucket = client.bucket(parsed_url.netloc)
    blob = bucket.blob(parsed_url.path.lstrip("/"))
    blob.download_to_filename(local_file)
  
def upload_to_gcs(local_file, gcs_dir):
    parsed_url = urlparse(gcs_dir)
    client = storage.Client()
    bucket = client.bucket(parsed_url.netloc)
    blob_path = f"{parsed_url.path.lstrip('/').rstrip('/')}/{local_file}"
    bucket.blob(blob_path).upload_from_filename(local_file)
  
# --- 1. データの準備 ---
print(f"Downloading data from {args.train_data_uri}...", flush=True)
local_train_file = "train.csv"
download_from_gcs(args.train_data_uri, local_train_file)
  
df = pd.read_csv(local_train_file)
cat_cols = ['cut', 'color', 'clarity']
for col in cat_cols:
    df[col] = df[col].astype('category')
  
X = df.drop(columns=["price"])
y = df["price"]
  
# スクリプト内で学習用と検証用に分割 (データリーク防止)
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.1, random_state=42)
  
# --- 2. モデルの学習 ---
print("Training LightGBM model...", flush=True)
model = lgb.LGBMRegressor(n_estimators=100, random_state=42)
  
# 学習過程を記録するために eval_set を渡す
model.fit(
    X_train, y_train,
    eval_set=[(X_train, y_train), (X_val, y_val)],
    eval_names=['train', 'valid']
)
  
# --- 3. 分析画像の生成と保存 ---
# ① 学習曲線の描画
lgb.plot_metric(model, metric='l2') 
plt.title('Learning Curve (MSE)')
plt.tight_layout()
plt.savefig("learning_curve.png")
plt.close()
  
# ② SHAP値（寄与度）の描画
print("Calculating SHAP values...", flush=True)
explainer = shap.TreeExplainer(model)
shap_values = explainer(X_val.sample(min(1000, len(X_val)), random_state=42))
  
plt.figure()
shap.plots.beeswarm(shap_values, show=False)
plt.title("SHAP Feature Importance")
plt.tight_layout()
plt.savefig("shap_importance.png")
plt.close()
  
# --- 4. 成果物のアップロード ---
aip_model_dir = os.getenv("AIP_MODEL_DIR")
if aip_model_dir:
    print(f"Uploading artifacts to: {aip_model_dir}", flush=True)
    model.booster_.save_model("model.txt")
    upload_to_gcs("model.txt", aip_model_dir)
    upload_to_gcs("learning_curve.png", aip_model_dir)
    upload_to_gcs("shap_importance.png", aip_model_dir)
    print("Upload completed.", flush=True)

Dockerfile の作成

Dockerfile を記述します。ベースイメージには Python 3.12 を指定し、LightGBM に必要な libgomp1 をインストールします。

%%writefile custom_container/Dockerfile
FROM python:3.12-slim
  
# LightGBM に必須の OS ライブラリをインストール
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*
  
# 必要な Python ライブラリのインストール
RUN pip install --no-cache-dir \
    pandas scikit-learn lightgbm shap matplotlib google-cloud-storage
  
WORKDIR /app
COPY task.py /app/task.py
  
ENTRYPOINT ["python", "task.py"]

コンテナのビルドとプッシュ

Cloud Build を使用してコンテナをビルドし、プッシュします。

# Cloud Build でビルドとプッシュを実行
REPO_NAME = "custom-training-repo"
IMAGE_URI = f"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{REPO_NAME}/lgbm-shap-trainer:latest"
  
!gcloud builds submit --tag {IMAGE_URI} ./custom_container

学習ジョブの実行

作成した自作コンテナ (IMAGE_URI) を指定して、学習ジョブを送信します。引数 base_output_dir を指定することで、指定した Cloud Storage のパス配下にモデルや画像を保存できます。

# ジョブの定義
job = aiplatform.CustomContainerTrainingJob(
    display_name="diamonds-lgbm-shap-job",
    container_uri=IMAGE_URI,
)
  
# ジョブの実行
print("ジョブを送信しました。完了までお待ちください...")
job.run(
    machine_type="n1-standard-4",
    replica_count=1,
    args=[
        f"--train-data-uri={WORK_DIR}/data/train.csv"
    ],
    # 成果物の保存先フォルダを指定
    base_output_dir=f"{WORK_DIR}/model_output"
)

推論と評価指標の確認

ジョブ完了後、Cloud Storage から学習済みモデルをダウンロードし、Colab Enterprise 上でテストデータに対する精度評価を行います。

import numpy as np
import lightgbm as lgb
from sklearn.metrics import r2_score, mean_squared_error
import pandas as pd
  
# 1. 学習の成果物のダウンロード
MODEL_DIR = f"{WORK_DIR}/model_output/model"
  
print("学習済みモデルと分析画像をダウンロードします...")
!gsutil cp {MODEL_DIR}/model.txt .
!gsutil cp {MODEL_DIR}/learning_curve.png .
!gsutil cp {MODEL_DIR}/shap_importance.png .
  
# 2. テストデータの読み込み
df_test = pd.read_csv(f"{WORK_DIR}/data/test.csv")
cat_cols = ['cut', 'color', 'clarity']
for col in cat_cols:
    df_test[col] = df_test[col].astype('category')
  
X_test = df_test.drop(columns=["price"])
y_true = df_test["price"]
  
# 3. ローカル推論の実行
local_model = lgb.Booster(model_file="model.txt")
predictions = local_model.predict(X_test)
  
# 4. 評価指標の計算と表示
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)

以下は筆者の環境における実行結果です。R2スコアが 0.98 を超える精度の高いモデルが作成できました。

------------------------------
評価結果 (データ数: 5394件)
R2 Score (決定係数): 0.9817
RMSE (誤差の大きさ): 543.6218
------------------------------

分析レポートの解釈

単に予測精度を出すだけでなく、AI がなぜその予測をしたのかを解釈することは実務において重要です。コンテナ内で生成した学習曲線の画像と SHAP を用いた個別データの分析結果を確認します。

import shap
from IPython.display import Image, display
  
print("=== 学習曲線 (過学習の確認) ===")
display(Image("learning_curve.png"))
  
print("\n=== 全体の寄与度 (SHAP Beeswarm) ===")
display(Image("shap_importance.png"))
  
# --- 個別データに対するSHAP（表形式）---
print("\n=== 特定のデータ（1件目）の予測の根拠 ===")
  
explainer = shap.TreeExplainer(local_model)
single_instance = X_test.iloc[[0]]
shap_values_single = explainer(single_instance)
  
shap_df = pd.DataFrame({
    "特徴量 (Feature)": single_instance.columns,
    "実際の値 (Value)": single_instance.values[0],
    "価格への影響 (SHAP値)": shap_values_single.values[0]
})
  
shap_df = shap_df.reindex(shap_df["価格への影響 (SHAP値)"].abs().sort_values(ascending=False).index)
  
base_value = explainer.expected_value
predicted_price = predictions[0]
  
print(f"【ベースライン価格 (平均)】: {base_value:.2f}")
display(shap_df.style.format({"価格への影響 (SHAP値)": "{:+.2f}"}).hide(axis="index"))
print(f"【最終予測価格】: {predicted_price:.2f}")