はじめに

LightGBMは、Gradient Boosting Decision Tree (GBDT) アルゴリズムをベースとした人気の高い機械学習アルゴリズムです。
そのため、SageMaker でもLightGBM を使いたいというユースケースは非常に多いと思います。
SageMaker では、LightGBM を使う方法がいくつかあるため、ここではそれらの方法について整理したいと思います。

SageMakerでLightGBMを使ってモデル構築する方法

SageMakerでLightGBMを使ってモデル構築する方法は、現時点（2023/3）で以下の4つの方法があります。

1. SageMaker built-in LightGBM を利用
2. SageMaker でカスタムスクリプトを利用
3. SageMaker Autopilot を利用
4. SageMaker JumpStart を利用

上から順番に一般的な利用方法かと思います。それぞれの使いどころと利用方法についてまとめていきます。

1. SageMaker built-in LightGBM を利用

SageMaker では、built-in アルゴリズム（組み込みアルゴリズム）といって、入力データ、ハイパーパラメータ、インスタンスタイプなど実行環境の指定をするだけで、すぐに機械学習アルゴリズムを利用できる機能があります。
built-in アルゴリズムには、LightGBM も用意されており、ローコードでLightGBM を利用することが可能になっています。

回帰、分類ともに対応しており、主要なハイパーパラメータの自動チューニングやDask を利用した分散学習にも対応しており、大半のユースケースはこれでカバーできると思います。

実装サンプル

まずは、実行環境となるコンテナイメージと提供されているトレーニングスクリプト、およびトレーニング済みのモデルのURIを取得します。
SageMaker Python SDK で用意されている、各retrieveを利用して簡単に取得することができます。

built-in LightGBM は、スクリプトモードで実行されますが、トレーニングスクリプトが提供されているため、置き換える必要はありません。SageMaker をよく利用される方はお気づきかもですが、このトレーニングスクリプトを独自のLightGBM トレーニングスクリプトに指定する方法が「SageMaker でカスタムスクリプトを利用」にあたります。

また、トレーニング済みのモデルについては、必ずしも取得する必要はありません。Fine-tuning したい場合のみ指定すればOK です。（今回やろうとしている分類問題には対応していますが、対応しているユースケースも限られているので、利用される場合はご注意ください）

from sagemaker import image_uris, model_uris, script_uris

train_model_id, train_model_version, train_scope = "lightgbm-classification-model", "*", "training"
training_instance_type = "ml.m5.xlarge"

# 実行環境となるコンテナイメージを取得します。
train_image_uri = image_uris.retrieve(
    region=None,
    framework=None,
    model_id=train_model_id,
    model_version=train_model_version,
    image_scope=train_scope,
    instance_type=training_instance_type
)

# SageMaker が提供するトレーニングスクリプトを取得します。
train_source_uri = script_uris.retrieve(
    model_id=train_model_id, model_version=train_model_version, script_scope=train_scope
)

# トレーニング済みのモデルを取得します。
train_model_uri = model_uris.retrieve(
    model_id=train_model_id, model_version=train_model_version, model_scope=train_scope
)

※余談：上記で取得したコンテナイメージは763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-trainingでPytorchがベースイメージになっており、デフォルトだとLightGBM のパッケージは入ってません。train_source_uriの中にLightGBM のパッケージが含まれており、実行時にインストールされるように構成されているようです。

次に、built-in LightGBM に指定する、IAMロールや入出力データの設定を定義します。

import sagemaker, boto3, json
from sagemaker import get_execution_role

aws_role = get_execution_role()
aws_region = boto3.Session().region_name
sess = sagemaker.Session()

training_data_bucket = f"jumpstart-cache-prod-{aws_region}"
training_data_prefix = "training-datasets/tabular_multiclass"

training_dataset_s3_path = f"s3://{training_data_bucket}/{training_data_prefix}/train" 
validation_dataset_s3_path = f"s3://{training_data_bucket}/{training_data_prefix}/validation" 

output_bucket = sess.default_bucket()
output_prefix = "jumpstart-example-tabular-training"

s3_output_location = f"s3://{output_bucket}/{output_prefix}/output"

built-in LightGBM のデフォルトのハイパーパラメータを取得します。

from sagemaker import hyperparameters

hyperparameters = hyperparameters.retrieve_default(
    model_id=train_model_id, model_version=train_model_version
)

# 取得したパラメータは自由に変更でき、チューニングすることが可能です。
hyperparameters[
    "num_boost_round"
] = "500"
print(hyperparameters)

SageMaker Estimator のインスタンス（実行環境定義）を作成します。
ここまで、記述してきた各種定義をEstimator の引数に指定します。

from sagemaker.estimator import Estimator
from sagemaker.utils import name_from_base

training_job_name = name_from_base(f"built-in-algo-{train_model_id}-training")

# SageMaker Estimator のインスタンス（実行環境定義）を作成します。
tabular_estimator = Estimator(
    role=aws_role,
    image_uri=train_image_uri,
    source_dir=train_source_uri,
    model_uri=train_model_uri,
    entry_point="transfer_learning.py",
    instance_count=1, # Dask を利用した分散学習をしたい場合は、これを2 以上に指定します。
    instance_type=training_instance_type,
    max_run=360000,
    hyperparameters=hyperparameters,
    output_path=s3_output_location
)

最後に、入力データとなるS3 パスを指定して、トレーニングジョブを実行します。

tabular_estimator.fit(
    {
        "train": training_dataset_s3_path,
        "validation": validation_dataset_s3_path,
    }, logs=True, job_name=training_job_name
)

トレーニングジョブが完了すると、ほかのbuilt-in アルゴリズム同様に、簡単にLightGBM モデルをエンドポイントにデプロイし、推論を実行することができます。
推論用のコンテナイメージや推論スクリプトもSageMaker がすべて提供してくれます。

inference_instance_type = "ml.m5.large"

# 実行環境となるコンテナイメージを取得します。
deploy_image_uri = image_uris.retrieve(
    region=None,
    framework=None,
    image_scope="inference",
    model_id=train_model_id,
    model_version=train_model_version,
    instance_type=inference_instance_type,
)

# SageMaker が提供する推論スクリプトを取得します。
deploy_source_uri = script_uris.retrieve(
    model_id=train_model_id, model_version=train_model_version, script_scope="inference"
)

endpoint_name = name_from_base(f"built-in-algo-{train_model_id}-")

# SageMaker endpoint にモデルをデプロイします。
predictor = tabular_estimator.deploy(
    initial_instance_count=1,
    instance_type=inference_instance_type,
    entry_point="inference.py",
    image_uri=deploy_image_uri,
    source_dir=deploy_source_uri,
    endpoint_name=endpoint_name,
)

データさえ用意すれば、上記のコードのみで、モデル構築・推論まで非常に簡単に実行できます。また、大規模データでモデル構築する場合は、分散学習を選択したり、自動モデルチューニングでより精度の高いモデルを構築できたりもし、非常に便利な機能です。

2. SageMaker でカスタムスクリプトを利用

自前のLightGBM スクリプトを利用したい場合の方法を紹介します。built-in LightGBM では対応していないCross Validation を実行したり、対応していないパラメータをチューニングしたりする場合に利用することが想定されます。
※Cross Validation は、SageMaker TrainingJob を複数実行させ、うまくPipeline 化してやれば、built-in LightGBM を利用しても実現可能かと思いますが、SageMaker の機能をよく理解していないと難しいと思われるので、利用方法としてあげました。（機会があったら、試してみたいと思います）

実装サンプル

まずは、自前のLightGBM スクリプト（train.py）を用意します。このスクリプトを動作させることだけが目的のため、非常にシンプルな実装のみしておきます。（これだけだと、built-in LightGBM で十分笑）

train.py

from sklearn.metrics import accuracy_score, confusion_matrix
import lightgbm as lgb
import pandas as pd


df_train = pd.read_csv("/opt/ml/input/data/train/train.csv")
df_validation = pd.read_csv("/opt/ml/input/data/validation/validation.csv")


clf = lgb.LGBMClassifier()

clf.fit(
    df_train.drop("y_yes", axis=1), 
    df_train["y_yes"],
    eval_set = [(df_train.drop("y_yes", axis=1), df_train["y_yes"]),(df_validation.drop("y_yes", axis=1), df_validation["y_yes"])],
    early_stopping_rounds=100)

# モデルを使って訓練データ・検証データを予測し、評価
y_tr_pred = clf.predict(df_train.drop("y_yes", axis=1))
y_va_pred = clf.predict(df_validation.drop("y_yes", axis=1))

y_tr_pred = y_tr_pred.round(0)
y_va_pred = y_va_pred.round(0)

accuracy_tr = accuracy_score(y_tr_pred,df_train["y_yes"])
accuracy_va = accuracy_score(y_va_pred,df_validation["y_yes"])
print(f'[accuracy]train_data: {accuracy_tr:.5f} , validation_data: {accuracy_va:.5f}')

次に、LightGBM のパッケージファイルを用意します。
SageMaker Trainingjob では、Estimator のsource_dir でディレクトリを指定し、そのディレクトリの直下にrequirements.txt を配置して必要なライブラリを記載しておくと、実行コンテナにパッケージが自動でインストールされます。

LightGBMパッケージファイルを取得します。（pip installでインストールするのではなく、オフラインインストールなどで利用することが多い、ファイルを取得するコマンドを実行します。）

pip download -d train_dir/lib lightgbm===3.3.3

自前のLightGBM スクリプト（train.py）と、上記で取得したLightGBMパッケージファイル（依存関係含む）を以下のように配置します。

source_dir="train_dir"

train_dir
├── train.py
├── lib
│   ├── joblib-1.2.0-py3-none-any.whl
│   ├── lightgbm-3.3.3-py3-none-manylinux1_x86_64.whl
│   ├── numpy-1.24.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
│   ├── scikit_learn-1.2.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
│   ├── scipy-1.10.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
│   ├── threadpoolctl-3.1.0-py3-none-any.whl
│   └── wheel-0.40.0-py3-none-any.whl
├── requirements.txt

SageMaker Trainingjob の実行コンテナにパッケージを持ち込む方法は、AWS のSA さんが分かりやすく解説してくれているので、詳細は以下ご参照ください。

ここからは、built-in LightGBM 同様に、実行環境の定義をしていきます。
まずは、コンテナイメージを取得します。

from sagemaker import image_uris, model_uris, script_uris

train_model_id, train_model_version, train_scope = "lightgbm-classification-model", "*", "training"
training_instance_type = "ml.m5.xlarge"

# 実行環境となるコンテナイメージを取得します。
train_image_uri = image_uris.retrieve(
    region=None,
    framework=None,
    model_id=train_model_id,
    model_version=train_model_version,
    image_scope=train_scope,
    instance_type=training_instance_type
)

次に、built-in LightGBM に指定する、IAMロールや入出力データの設定を定義します。

import sagemaker, boto3, json
from sagemaker import get_execution_role

aws_role = get_execution_role()
aws_region = boto3.Session().region_name
sess = sagemaker.Session()

training_data_bucket = f"jumpstart-cache-prod-{aws_region}"
training_data_prefix = "training-datasets/tabular_multiclass"

training_dataset_s3_path = f"s3://{training_data_bucket}/{training_data_prefix}/train" 
validation_dataset_s3_path = f"s3://{training_data_bucket}/{training_data_prefix}/validation" 

output_bucket = sess.default_bucket()
output_prefix = "jumpstart-example-tabular-training"

s3_output_location = f"s3://{output_bucket}/{output_prefix}/output"

SageMaker Estimator のインスタンス（実行環境定義）を作成します。
ここまで、記述してきた各種定義をEstimator の引数に指定します。built-in LightGBM との大きな違いは、source_dir 引数が上記で用意した自前のLightGBM スクリプトとパッケージファイルが格納されているディレクトリを指定している点です。

from sagemaker.estimator import Estimator
from sagemaker.utils import name_from_base

training_job_name = name_from_base(f"{train_model_id}-training-custom")

# Create SageMaker Estimator instance
tabular_estimator = Estimator(
    role=role,
    image_uri=train_image_uri,
    source_dir="train_dir", # 
    entry_point="train.py",
    instance_count=1,
    instance_type=training_instance_type,
    max_run=360000,
    output_path=s3_output_location
)

最後に、入力データとなるS3 パスを指定して、トレーニングジョブを実行します。

tabular_estimator.fit(
    {
        "train": f"s3://{bucket}/{prefix}/train/train.csv",
        "validation": f"s3://{bucket}/{prefix}/validation/validation.csv",
    }, logs=True, job_name=training_job_name
)

この後は、自前の推論スクリプトを用意し、エンドポイントにデプロイして推論できるようにするのもいいですし、TrainingJob内で推論まで記載してしまうのもいいですし、作成したモデルを自由に扱うことができます。

SageMaker TrainingJob が提供しているカスタムスクリプト実行機能をうまく使うことで、非常に簡単に自前のLightGBM スクリプトを実行することが可能です。

3. SageMaker Autopilot を利用

上２つの方法とは少し毛色が頃なりますが、SageMaker Autopilot を利用して、LightGBM のモデルを構築することも可能です。Autopilot とは、データに基づいて最適なモデルを自動で構築してくれるSageMaker のAutoML 機能です。
Autopilot のアンサンブルモード（AutoGluon を利用したアンサンブル学習）を選択すると、LightGBM を含む複数のアルゴリズムを自動でチューニングし、高い精度のモデルを簡単に作ってくれます。
以下の記事でも紹介したのですが、Autopilot で利用するアルゴリズムを選択できるようになり、LightGBM だけに絞ってモデルのチューニングをできるようになっています。詳しい利用方法は以下記事をご覧ください。

LightGBM のチューニング方法に詳しくなくても利用可能なため、データサイエンティストが不足しているケースや、試行錯誤のサイクルを高速に回すことができるので、短期間でモデルを構築しないといけないケースで有用な方法になると思います。

Autopilot で自動構築したモデルをインプットにして、built-in LightGBM で細かいチューニングをして、より精度を高めていくなど段階的に各方法を使っていくのも良いと思います。

4. SageMaker JumpStart を利用

最後は、SageMaker JumpStart を利用する方法です。
SageMaker JumpStart は、さまざまな種類の問題に対応するトレーニング済みモデルを提供し、それらを簡単にデプロイ・チューニングできるテンプレートを提供してくれています。

LightGBM についても、トレーニング済みモデルおよび、チューニングするためのソースコードを提供してくれています。利用方法は以下ご参照ください。
https://docs.aws.amazon.com/sagemaker/latest/dg/jumpstart-content-sharing.html

また、LightGBM のトレーニング済みモデルをチューニングするソースコードは以下となります。
https://github.com/aws/amazon-sagemaker-examples/blob/main/introduction_to_amazon_algorithms/lightgbm_catboost_tabular/Amazon_Tabular_Classification_LightGBM_CatBoost.ipynb

トレーニング済みモデルをそのまま利用できたり、ちょっとのチューニングのみで済む場合は、この方法がもっとも労力が少ないと思います。
また、上記のノートブックファイルは、中身はbuilt-in LightGBM の実装コードになってます。そのため、とりあえず、この方法を検討して、ダメそうなら上記のノートブックファイルを参考にbuilt-in LightGBM を実装してみるという手もいいと思います。

使い分けまとめ

Well-Architected フレームワークの機械学習版である、機械学習レンズでも言及されていますが、より「マネージド」なサービス、機能を使って、ビジネス価値につながらないものは「なるべく作らない」を考え、各実装方法を選択していくのが良いと思います。

MLCOST-02: Use managed services to reduce total cost of ownership (TCO)

そのため、まずは最小限の労力でモデルを利用できる 「4. SageMaker JumpStart を利用」 を検討するのがいいでしょう。トレーニング済みのモデルで事足りるなら、モデル構築にかかるコスト（工数や従量課金）も最小限で済みますし、推論の実装に注力することが可能です。

JumpStart で提供されているモデルがマッチしない場合、 「3. SageMaker Autopilot を利用」 が次に選択肢になるかと思います。
自動化により、手作業による実験や比較を排除しながら、モデル構築・検証のプロセスをスピードアップすることができます。

MLCOST-15: Use automated machine learning

また、Autopilot の章でも言及しましたが、データサイエンティストが不足していたり、短期間で結果を出す必要があるケースが昨今多いと思うので、そのようなケースの解決策として有効になるのではないでしょうか。

さらに、細かいチューニングをしたい場合、 「1. SageMaker built-in LightGBM を利用」 が選択肢になります。そして、built-in LightGBM で対応していないようなパラメータチューニングやモデル構築をしたい場合は、 「2. SageMaker でカスタムスクリプトを利用」 が最後の選択肢になると思います。