SageMaker で学習ジョブを実行する ~独自スクリプト~ - nsakki55 のアウトプットブログの続きです

記事中のコード github.com

SagaMaker を用いたモデルの学習

SageMaker では学習に必要なスクリプトやライブラリをコンテナベースで管理し、学習に必要なインフラ管理を自動で行なってくれます。

思想としては、データサイエンティストが面倒なインフラ管理や推論サービスの作業を行わず、MLモデルの開発に集中できることが SageMaker を使用するメリットとなっています。

SageMaker Training Jobが行なっていることは、大きく分けると以下の流れになります

S3から入力データを読み込み
学習用コンテナを実行
モデルアーティファクトをS3に書き戻す

引用：Amazon SageMaker でモデルをトレーニングする - Amazon SageMaker

データ入出力を利用するためのデータパスの対応や、コンテナ環境の準備など、SageMaker で学習を実行するためのルールがあります。

SageMakerで学習を実行する場合、3つパターンが存在します

組み込みアルゴリズム : AWSが提供するコンテナ・学習スクリプトを実行
独自スクリプト: AWSが提供するコンテナ上で、独自スクリプトを実行
独自コンテナ: 独自スクリプトを含めた独自コンテナを実行

コード量を少なくし手軽にモデルの学習を実行できるのは組み込みアルゴリズムを利用する場合ですが、実装の柔軟性は低くなります。一方、独自スクリプトや独自コンテナを利用する場合、コード量は増えますが実装の柔軟性は高くなり、独自のアルゴリズムでの学習を実行することができるようになります。

初めてSageMakerでモデルの学習を実行する際、実行パターンの多さ・守るルールの多さに面食らってしまう人が多いと思います。ドキュメントやサンプルコードは豊富に提供されていますが、どのパターンの学習方法で、最低限必要な要素は一体何なのかを把握するのが大変です。

今回は3つのパターンで独自の予測モデルの学習を行い、学習を実行する流れを整理ます。

SageMaker ではモデル学習のために以下のようなサービスを提供してくれていますが、本記事では扱いません。

HyperparameterTuning
SagaMaker Model Monitoring
Clarify
Debugger
Endpoint

...

本記事では独自コンテナを用いてSageMaker 学習ジョブで実行します SageMaker 学習ジョブを実行するには AWS SDK, SageMaker SDKを用いる2つの方法がありますが、今回は jupyter notebook から学習をジョブを実行するため、SageMaker SDKを使用します。

データセット・モデル

今回は広告をクリックする確率(CTR)を予測を予測する二値分類モデルの学習を行います。

データセットは kaggle の avazu-ctr-predictionのデータセットを使用します。

https://www.kaggle.com/c/avazu-ctr-prediction

Feature Engineering・モデルは

FeatureHashing・線形回帰モデル
OneHotEncoding・Factorization Machines

の2パターンを実装します

学習データの準備

学習ジョブに渡すデータをS3にアップロードしておきます。

今回はtrain , validation , test にデータを分割します

gist82e34280293abba2b1743f49e1cc7102

独自コンテナ

独自コンテナを用いて学習ジョブを実行する方法は、3パターンあります

AWS が提供しているコンテナイメージを拡張する方法
独自のコンテナイメージに SageMaker Training Toolkit をインストールする方法
スクラッチでコンテナイメージを作成する方法

引用：Amazon SageMaker におけるカスタムコンテナ実装パターン詳説〜学習編〜 | Amazon Web Services ブログ

公式サンプルコードやドキュメントの多くは、2つめのパターンの独自のコンテナイメージにSageMaker Training Toolkit で実装しているものが多いです。

自分がSageMaker で独自コンテナを利用しようと思い立った時、色々なやり方が錯綜していて困ったので、3つのパターンの学習ジョブを実行する流れを整理したいとおもいます。

独自コンテナで学習ジョブを実行する流れ

独自コンテナの学習ジョブを実行する流れは以下になります

学習スクリプトを用意
学習ライブラリ・スクリプトを設定したDockerfileの作成
ECRへdocker image をpush
Estimatorクラスにdocker image, 学習設定を渡して学習ジョブを実行

AWS が提供しているコンテナイメージを拡張する方法

独自スクリプトを利用する場合も、AWSが提供するコンテナ環境に独自実装のライブラリやthird party のライブラリをrequirements.txt を配置しておくことで利用することができます。

requriements.txt による公式コンテナ環境の拡張方法はこちらの記事でまとめています SageMaker で学習ジョブを実行する ~独自スクリプト~ - nsakki55 のアウトプットブログ

AWSが提供するコンテナイメージをベースとして、新しく作成したコンテナイメージを利用する方法を説明します。

まずベースとなるAWS提供のコンテナイメージのimage uriを取得します。

Scikit-Learn 環境のimage: https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/pre-built-docker-containers-scikit-learn-spark.html

利用したいライブラリのEstimator クラスのtraining_image_uriメソッドでimage uri を取得することができます。

from sagemaker.sklearn.estimator import SKLearn

estimator = SKLearn(entry_point="",
                    framework_version="0.23-1",
                    py_version="py3",
                    role=role,
                    instance_type='local')

print(estimator.training_image_uri())
# 354813040037.dkr.ecr.ap-northeast-1.amazonaws.com/sagemaker-scikit-learn:0.23-1-cpu-py3

third party library のoptuna を利用してハイパーパーパラメータチューニングを行う学習スクリプトを用意します。
AWS提供のimage にはsagemaker training toolkit がインストールされているため、環境変数を用いて学習データ・ハイパーパラメータを学習スクリプトに受け渡すことができます。
sagemaker training toolkit を用いた学習スクリプトの記述についてはSageMaker で学習ジョブを実行する ~独自スクリプト~ - nsakki55 のアウトプットブログの記事で取り上げています。

gist39c43d40dd3232b9c41433c485d3e0d5

SageMaker の公式コンテナイメージをベースイメージにした Dockerfile を作成します。

拡張ライブラリ・学習スクリプトをコピーします。

SageMaker の公式コンテナイメージにには、sagemaker-training-toolkit がインストールされているため、学習ジョブを開始する際の実行スクリプトを環境変数 SAGEMAKER_PROGRAM で設定します。

gistf7d8917f9dc5fdf54ae29e71fad71f86

Dockerfileを元に、ECRへイメージをpushするスクリプト

gist47fa76080719b17210c1327a2f17d1a2

docker image をECRへpushします

./buiid_push_image.sh ctr-prediction-extend-sagemaker-container extend_sagemaker_container

作成したコンテナイメージで学習ジョブを実行するには、Estimator クラスの image_uri にECRに登録した image のURIを設定します。

以下のディレクトリ構成で実行します

src
 ├── extend_sagemaker_container_trainer.ipynb
 │
 ├── build_push_image.sh
 │
 └── extend_sagemaker_container
       │
       ├── Dockerfile
       │
       └── extend_sagemaker_container_trainig_script.py

gistbfabcc8d05f9cc82ebfa04571749dc86

管理画面から学習ジョブが実行されてるのを確認できます

学習ジョブを実行すると、Dockerfile中に環境変数 SAGEMAKER_PROGRAM に設定したスクリプトが実行されます。ログを確認すると、extend_sagemaker_container_trainig_script.py が実行されていることを確認できます。

参考: 公式コンテナイメージの拡張を行う場合の公式サンプルコードamazon-sagemaker-examples/advanced_functionality/pytorch_extending_our_containers at main · aws/amazon-sagemaker-examples · GitHub

独自のコンテナイメージに SageMaker Training Toolkit をインストールする方法

AWSが提供するコンテナイメージを拡張する場合、sagemaker-training-toolkit を自分でインストールせずに利用することができます。

完全に新しく自作のコンテナイメージを作成し、AWSが提供するコンテナイメージと同じように実行する場合、sagemaker-training-toolkit をインストールする必要があります。

ここでは、fastFMをインストールした独自コンテナイメージで学習ジョブを実行してみます。独自のトレーニングコンテナを適応させる - Amazon SageMaker

fastFMライブラリを利用する学習スクリプトを用意します。

AWS提供のコンテナイメージを利用する場合と同様、sagemaker-training-toolkit の環境変数を用いてデータ、ハイパーパラメータの受け渡しを行います。

gista8f890bad85b6fedce7656850956be99

必要なライブラリを含めたDockerfile を用意します。

この時、sagemaker-training をインストールすることで、sagemaker-training-toolkitの環境変数によるデータ・ハイパーパラメータの受け渡しができるようになります。

環境変数 SAGEMAKE_PROGRAM で実行スクリプトを指定できます。

gistd615eda7343d433c88b6fcdec8fd5fe5

環境変数 SAGEMAKER_PROGRAM を指定しない場合、以下のコマンドでコンテナイメージが実行されます。 Amazon SageMaker がトレーニングイメージを実行する方法 - Amazon SageMaker

docker run image名 train

この場合、学習スクリプトをキックするtrainというファイル名のシェルスクリプトを用意します

#! /bin/sh
python /opt/ml/code/custom_toolkit_container_training_script.py

Dockerfile にtrain ファイルのコピーを追加します

gist218d3893ec4fe3ec6b28770db2d277a2

ECRへ imageをpushします

./buiid_push_image.sh ctr-prediction-custom-toolkit-container custom_toolkit_container

Estimator クラスのimage_uri にECRにpushした独自コンテナを指定します。

ハイパーパラメータをAWS提供のコンテナイメージを利用する場合と同様に、Estimatorクラスに渡します。

src
 ├── custom_toolkit_container_trainer.ipynb
 │
 ├── build_push_image.sh
 │
 └── custom_toolkit_container
       │
       ├── Dockerfile
       │
       └── custom_toolkit_container_trainig_script.py

gist4082bd52bf45468acb96395d0f41c3d7

独自コンテナを指定した学習ジョブが実行されていることを確認できます。

スクラッチでコンテナイメージを作成する方法

sagemaker-training-toolkit を利用せず、独自コンテナをSageMaker の学習上で実行する場合、以下の機能を実装する必要があります

コンテナ中の学習スクリプトを ENTRYPOINTとして指定できる機能
ハイパーパラメータを学習スクリプトの実行時引数として渡される機能
学習ジョブの正常・異常終了をサービスに通知する機能

sagemaker-training-toolkit をインストールする場合、上記の機能が自動で追加されます。

公式サンプルコードを参考に実装していきます amazon-sagemaker-examples/advanced_functionality/custom-training-containers at main · aws/amazon-sagemaker-examples · GitHub

SageMaker Training Jobではハイパーパラメータ・データの情報がそれぞれ /opt/ml/input/config/hyperparameters.json , /opt/ml/input/config/inputdataconfig.json に書き込まれます。

Amazon SageMaker がトレーニング情報を提供する方法 - Amazon SageMaker

データは/opt/ml/input/data/チャンネル名以下に自動的に配置されます。

学習が正常終了した際は sys.exit(0) , 異常終了した際は sys.exit(1) を返すようにします。

上記の機能を追加した学習スクリプトを作成します

gista6729d6984de03fd7473f3cbeda15712

Dockerfile に必要なライブラリのインストールと、ENTRYPOINT に学習スクリプトを実行するコマンド設定を記述します

gist082eeba3782b1a5da3690f914f88c0f4

image を ECR へ pushします

./buiid_push_image.sh ctr-prediction-scratch-container scratch_container

以下のフォルダ構成で学習ジョブを実行します

src
 ├── scratch_container_trainer.ipynb
 │
 ├── build_push_image.sh
 │
 └── scratch_container
       │
       ├── Dockerfile
       │
       └── scratch_container_trainig_script.py

学習ジョブの実行は、sagemaker-training-toolkitを利用する場合と同様の方法で行います。

gist61eac4ca470901c1569042878978677b

学習ジョブが実行されていることを確認できます。

実行ログを見ると、ハイパーパラメータや学習データが学習ジョブに渡せていることを確認できます。

まとめ

独自コンテナでSageMaker 学習ジョブを実行する3つの方法を試してみました。データ・ハイパーパラメータの設定・学習ジョブを実行するためのentrypointの設定など、利用するために覚えることが多いなと思います。一度ルールを覚えてしまえば、インフラ管理を気にする必要なくなるため実験を沢山実行したいDS・MLエンジニアに便利なサービスだと思います。