オープンソース Amazon SageMaker XGBoost アルゴリズムコンテナのご紹介

XGBoost は、表形式データセットでの回帰タスクと分類タスクの実行によく使用される効率的な機械学習 (ML) アルゴリズムです。このアルゴリズムはツリーに勾配ブースティングと呼ばれる技法を実装し、ML 大会で極めて優れた成績を残しています。

Amazon SageMaker は、提供開始時から XGBoost を組み込みのマネージドアルゴリズムとしてサポートしてきました。詳細については、XGBoost と Amazon SageMaker を使った機械学習の簡素化を参照してください。この記事を書いている時点で、より優れた柔軟性、スケーラビリティ、拡張性、およびマネージドスポットトレーニングを備えたオープンソース Amazon SageMaker XGBoost コンテナを活用することが可能です。詳細については、GitHub で Amazon SageMaker サンプルノートブックと sagemaker-xgboost-container を参照、または「Xboost アルゴリズム」をお読みください。

この記事では、オープンソース XGBoost アルゴリズムコンテナのメリットと、3 つのユースケースをご紹介します。

オープンソース SageMaker XGBoost コンテナのメリット

新しい XGBoost コンテナには、以下のメリットがあります。

柔軟性

新しいスクリプトモードによって、独自のトレーニングスクリプトをカスタマイズまたは使用することができるようになりました。TensorFlow、MXNet、PyTorch、および Chainer のユーザーも利用できるこの機能は、カスタマイズされたプリプロセスまたはポストプロセスロジックの追加、トレーニングプロセス後の追加ステップの実行、すべての XGBoost 機能 (交差検証のサポートなど) の活用を可能にします。他の Amazon SageMaker 組み込みアルゴリズムといった非スクリプトアルゴリズムモードも引き続き使用できます。このモードでは、データの場所とハイパーパラメータの指定のみが必要です。

スケーラビリティ

オープンソースコンテナでは、分散型トレーニングがより効率的に実装されているため、より多くのインスタンスにスケールアウトし、メモリ不足エラーを削減することができます。

拡張性

コンテナはオープンソースであることから、スクリプトモードの使用以外にも、ニーズに合わせてアルゴリズムを拡張、分割、または変更できます。これには、追加のライブラリのインストール、および基盤となる XGBoost バージョンの変更が含まれます。

マネージドスポットトレーニング

マネージドスポットトレーニングのサポートにより、Amazon SageMaker XGBoost トレーニングジョブでの最大 90% の節約が可能になります。この完全マネージド型オプションでは、AWS クラウドにある未使用のコンピューティング性能を活用することができます。スポットインスタンスは Amazon SageMaker がユーザーに代わって管理するので、キャパシティーのポーリングに煩わされる必要はありません。XGBoost の新バージョンは、ジョブが信頼性に優れた形で完了することを確実にするために、チェックポイントを自動で管理します。詳細については、「Amazon SageMaker のマネージドスポットトレーニング」および「Amazon SageMaker でのチェックポイントの使用」を参照してください。

追加の入力形式

XGBoost は、Parquet および Recordio-protobuf の各入力形式をサポートするようになりました。Parquet は、データ分析システムでの使用のための、標準化されたオープンソースの自己記述型列指向ストレージ形式です。Recordio-protobuf は、さまざまなアルゴリズムのために Amazon SageMaker 全体で使用される一般的なバイナリデータ形式で、今後は XGBoost でもトレーニングと推論のためにサポートされます。詳細については、「トレーニングの共通データ形式」を参照してください。このコンテナはさらに、これらのデータ形式について Pipe モードトレーニングもサポートします。詳細については、Amazon SageMaker アルゴリズムでパイプ入力モードを使用するを参照してください。

組み込みアルゴリズムとしての最新 XGBoost コンテナの使用

既存の Amazon SageMaker XGBoost ユーザーとして、トレーニングジョブの作成時にバージョンを指定することで新機能と向上したパフォーマンスを活用することができます。XGBoost の使用開始、または最新バージョンの使用に関する詳細については、GitHub リポジトリを参照してください。

新しいコンテナには、フレームワークバージョン (1.0-1) を指定することでアップグレードできます。このバージョンは、アップストリーム XGBoost フレームワークバージョン (1.0) と追加の Amazon SageMaker バージョン (1) を指定します。レガシー 0.72 コンテナに基づいた既存の XGBoost ワークフローがある場合、このコンテナで同じワークフローを機能させる上で必要になるのはこの変更のみです。コンテナは、0.90-1 のバージョンを使用することによって XGBoost 0.90 もサポートします。

以下のコードを参照してください。

from sagemaker.amazon.amazon_estimator import get_image_uri
container = get_image_uri(region, 'xgboost', '1.0-1')

estimator = sagemaker.estimator.Estimator(container,
                                          role,
                                          hyperparameters=hyperparameters,
                                          train_instance_count=1,
                                          train_instance_type='ml.m5.2xlarge',
                                          )

estimator.fit(training_data)

マネージドスポットインスタンスの使用

Estimator で train_use_spot_instances フラグを有効化することで、マネージドスポットインスタンスのサポートも活用できます。詳細については、GitHub リポジトリを参照してください。

マネージドスポットインスタンスを使ってトレーニングするときは、トレーニングジョブが中断される場合があり、開始から終了までの時間が長くなる原因となります。トレーニングジョブが中断された場合は、チェックポイントとしたスナップショットを使って以前に保存されたポイントから再開することができ、トレーニング時間 (およびコスト) を削減することができます。また、トレーニングジョブがスナップショットを保存する場所である checkpoint_s3_uri を使って、スポットインスタンスの中断時にシームレスに再開することもできます。以下のコードを参照してください。

estimator = sagemaker.estimator.Estimator(container,
                                          role,
                                          hyperparameters=hyperparameters,
                                          train_instance_count=1,
                                          train_instance_type='ml.m5.2xlarge',
                                          output_path=output_path,
                                          sagemaker_session=sagemaker.Session(),
                                          train_use_spot_instances=train_use_spot_instances,
                                          train_max_run=train_max_run,
                                          train_max_wait=train_max_wait,
                                          checkpoint_s3_uri=checkpoint_s3_uri
                                         )
                                         
estimator.fit({'train': train_input})

このジョブが終了する頃に、以下の 2 行の出力が表示されます。

Training seconds: X – トレーニングジョブの実際のコンピューティング時間
Billable seconds: Y – スポット割引の適用後に課金される時間

train_use_spot_instances を有効化した場合は X と Y の間に顕著な違いが見られるはずで、この違いはマネージドスポットトレーニングの使用からのコスト節約を示しています。これは、以下のコードに反映されています。

Managed Spot Training savings: (1-Y/X)*100 %

スクリプトモードの使用

スクリプトモードは、オープンソース Amazon SageMaker XGBoost コンテナの新機能です。独自のトレーニングまたはホスティングスクリプトを使用して、XGBoost トレーニングまたは推論ワークフローを完全にカスタマイズできます。以下のコード例は、スクリプトモードでのカスタマイズされたトレーニングスクリプトの使用を、順を追って説明するものです。詳細については、GitHub リポジトリを参照してください。

エントリポイントスクリプトの準備

一般的なトレーニングスクリプトは、入力チャネルからデータをロード、ハイパーパラメータを使ってトレーニングを設定、モデルをトレーニング、それからモデルを model_dir に保存して、後ほどホストできるようにします。ハイパーパラメータは引数としてスクリプトに渡され、argparse.ArgumentParser インスタンスで取得できます。

トレーニングジョブを作成するときは、メインガードを初めとして、パーサーを使って Amazon SageMaker estimator に渡されたハイパーパラメータを読み込みます。これらのハイパーパラメータは、入力スクリプトが引数として使用できるようにされます。また、Amazon SageMaker 固有の環境変数もいくつか解析し、入力データの場所、モデルの保存先などのトレーニング環境に関する情報を取得します。以下のコードを参照してください。

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    # Hyperparameters are described here
    parser.add_argument('--num_round', type=int)
    parser.add_argument('--max_depth', type=int, default=5)
    parser.add_argument('--eta', type=float, default=0.2)
    parser.add_argument('--objective', type=str, default='reg:squarederror')
    
    # Sagemaker specific arguments.Defaults are set in the environment variables.
    parser.add_argument('--train', type=str, default=os.environ['SM_CHANNEL_TRAIN'])
    parser.add_argument('--validation', type=str, default=os.environ['SM_CHANNEL_VALIDATION'])
    
    args = parser.parse_args()
    
    train_hp = {
        'max_depth': args.max_depth,
        'eta': args.eta,
        'gamma': args.gamma,
        'min_child_weight': args.min_child_weight,
        'subsample': args.subsample,
        'silent': args.silent,
        'objective': args.objective
    }
    
    dtrain = xgb.DMatrix(args.train)
    dval = xgb.DMatrix(args.validation)
    watchlist = [(dtrain, 'train'), (dval, 'validation')] if dval is not None else [(dtrain, 'train')]

    callbacks = []
    prev_checkpoint, n_iterations_prev_run = add_checkpointing(callbacks)
    # If checkpoint is found then we reduce num_boost_round by previously run number of iterations
    
    bst = xgb.train(
        params=train_hp,
        dtrain=dtrain,
        evals=watchlist,
        num_boost_round=(args.num_round - n_iterations_prev_run),
        xgb_model=prev_checkpoint,
        callbacks=callbacks
    )
    
    model_location = args.model_dir + '/xgboost-model'
    pkl.dump(bst, open(model_location, 'wb'))
    logging.info("Stored trained model at {}".format(model_location))

エントリポイントスクリプト内では、オプションとして Amazon SageMaker のホスティングまたはバッチ変換を使用するときの推論エクスペリエンスをカスタマイズすることができます。以下をカスタマイズできます。

input_fn() – 入力の処理方法
predict_fn() – XGBoost モデルの呼び出し方法
output_fn() – レスポンスを返す方法

このユースケースではデフォルトを使えるので、これらを定義する必要はありません。

Amazon SageMaker XGBoost estimator を使ったトレーニング

トレーニングデータとスクリプトを準備したら、Amazon SageMaker Python SDK の XGBoost estimator クラスによって、Amazon SageMaker が管理するトレーニングインフラストラクチャでそのスクリプトをトレーニングジョブとして実行できます。また、IAM ロール、使用したいインスタンスのタイプ、およびスクリプトに渡したいハイパーパラメータのディクショナリも estimator に渡します。以下のコードを参照してください。

from sagemaker.session import s3_input
from sagemaker.xgboost.estimator import XGBoost

xgb_script_mode_estimator = XGBoost(
    entry_point="abalone.py",
    hyperparameters=hyperparameters,
    image_name=container,
    role=role,
    train_instance_count=1,
    train_instance_type="ml.m5.2xlarge",
    framework_version="1.0-1",
    output_path="s3://{}/{}/{}/output".format(bucket, prefix, "xgboost-script-mode"),
    train_use_spot_instances=train_use_spot_instances,
    train_max_run=train_max_run,
    train_max_wait=train_max_wait,
    checkpoint_s3_uri=checkpoint_s3_uri
)

xgb_script_mode_estimator.fit({"train": train_input})

カスタム XGBoost モデルのデプロイメント

モデルをトレーニングしたら、Amazon SageMaker エンドポイント (つまり、推論を実行するために使用できるホストされるマネージド型予測サービス) を作成するために estimator を使用することができます。以下のコードを参照してください。

predictor = xgb_script_mode_estimator.deploy(initial_instance_count=1, instance_type="ml.m5.xlarge")
test_data = xgboost.DMatrix('/path/to/data')
predictor.predict(test_data)

Parquet 入力でのトレーニング

これで、Amazon SageMaker がサポートするオープンソースの ML-IO ライブラリを使用して、Parquet 形式のファイルまたはストリームで最新の XGBoost アルゴリズムを直接トレーニングすることができるようになりました。ML-IO は、複数のデータ形式をサポートする ML フレームワークのための高機能データアクセスライブラリで、最新の XGBoost コンテナにデフォルトでインストールされています。Parquet ファイルのインポートと、それを使ったトレーニングの詳細については、GitHub リポジトリを参照してください。

まとめ

Amazon SageMaker 向けのオープンソース XGBoost コンテナは、完全マネージド型のエクスペリエンスと、トレーニングにおけるコストを節約する追加のメリットを提供し、さらなる柔軟性を可能にします。

著者について

Rahul Iyer は AWS AI のソフトウェア開発マネージャーです。フレームワークアルゴリズムチームを率いる Rahul は、XGBoost および Scikit-learn といった機械学習フレームワークの構築と最適化を行っています。仕事以外では、ネーチャーフォトを楽しみ、家族との時間を大切にしています。

Rocky Zhang は AWS SageMaker のシニアプロダクトマネージャーです。Rocky は、お客様が機械学習を使って実世界のビジネス課題を解決するために役立つ製品を構築しています。仕事以外では、サッカーの観戦、プレー、およびコーチ業にほとんどの時間を費やしています。

Eric Kim は、Amazon AI のアルゴリズム & プラットフォームグループのエンジニアです。AWS のサービスである SageMaker のサポートを支援する Eric は、機械学習の研究、開発、および応用での経験を持っています。仕事以外では、音楽愛好家であり、あらゆるタイプの犬が大好きです。

Laurence Rouesnel は Amazon AI のシニアマネージャーです。SageMaker AutoPilot およびアルゴリズムといった深層学習および機械学習の研究と製品に携わるエンジニアと科学者のチームを率いています。仕事以外では、旅行、テーブルトーク RPG、およびランニングの熱烈なファンです。

Amazon Web Services ブログ