SageMaker で学習ジョブを実行する手順をまとめます。
記事中の実行コード: https://github.com/nsakki55/code-for-blogpost/tree/main/sagemaker_training/built-in-algorithm

SagaMaker を用いたモデルの学習

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

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

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

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

出典：Amazon でモデルをトレーニングする SageMaker - アマゾン SageMaker

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

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

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

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

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

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

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

HyperparameterTuning
SagaMaker Model Monitoring
Clarify
Debugger
Endpoint

...

文量が多くなるため、組み込みアルゴリズム・独自スクリプト・独自コンテナを使用するケースで記事を分けます。

この記事では組み込みアルゴリズムを利用した学習の流れを取り扱います。

AWS SDK・SageMaker SDK

SageMaker を操作する際、AWS SDK (boto3)とSageMaker SDK (sagemaker)の二つの方法があります。それぞれのユースケースは以下のように公式で説明されています。

参考：Amazon SageMaker 紹介 & ハンズオン（2018/07/25 実施）

AWS SDK
- Document: SageMaker — Boto3 Docs 1.26.7 documentation
- Code: GitHub - boto/boto3: AWS SDK for Python
- 基盤エンジニア向け
- AWS のリソース全般への操作が可能
- プロダクション環境で、定期学習ジョブの実行や、エンドポイントへのモデルのデプロイ等を実施するためのもの
- API 例

create-endopoint
create-notebook-instance
create-training-job
delete-endopoint
delete-notebook-instance
describe-endpoint
describe-notebook-instance

SageMaker SDK
- Document: Amazon SageMaker Python SDK — sagemaker 2.116.0 documentation
- Code : GitHub - aws/sagemaker-python-sdk: A library for training and deploying machine learning models on Amazon SageMaker
- データサイエンティスト向け
- SageMakerリソースへの操作が可能. AWS SDKをscikit-learn ライクにラップしたインターフェイス
- Python 版と Spark 版がある
- Jupyter Notebook 上で機械学習モデルの高速な開発と学習を行うためのもの
- API 例

estimator = Tensorflow(...)
estimator.set_hypyerprameters(...)
estimator.fit(...)
predictor = estimator.deploy(...)
Predictor.predict(...)

データセット・モデル

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

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

Click-Through Rate Prediction | Kaggle

特徴量前処理はOneHotEncoding を行い、線形回帰モデルを使用します。 AWSのドキュメントに合わせ、本文以下では線形学習アルゴリムと呼ぶことにします。

組み込みアルゴリズム

組み込みアルゴリズムを利用する場合、解きたい問題に合わせてAWS が提供しているアルゴリズムの中から適切なものを選択する必要があります。

解きたい問題ごとのアルゴリズムの一覧は公式ドキュメント中にまとめられています。 Amazon を使った SageMaker 組み込みアルゴリズムまたは事前トレーニングモデル - アマゾン SageMaker

今回は教師あり学習アルゴリズムの二値分類問題を解くので、「線形学習アルゴリズム(Linear Lerner)」を利用します。
線形学習アルゴリズム - アマゾン SageMaker

組み込みアルゴリズムで学習を実行するために必要なことは以下の2点です

指定された形式のデータをS3に用意
AWSが提供する docker image を指定して学習を実行

アルゴリズムごとにサポートされているデータ形式が異なります

例えば線形学習モデル、XGBoostでは

線形学習モデル : recordIO-wrapped protobuf , CSV
XGBoost : CSV , libsvm

のデータ形式がサポートされています。

指定された形式のデータをS3に用意

AWS SDK, SageMaker SDKを利用する設定、データの読み込みを行います。

import yaml
import pandas as pd
import sagemaker
import boto3

SETTING_FILE_PATH = "../config/settings.yaml"
DATA_FOLDER_PATH = "avazu-ctr-prediction"

with open(SETTING_FILE_PATH) as file:
    aws_info = yaml.safe_load(file)
        
sess = sagemaker.Session()
role = aws_info['aws']['sagemaker']['role']
bucket = aws_info['aws']['sagemaker']['s3bucket']
region = aws_info['aws']['sagemaker']['region']

sm = boto3.client('sagemaker')
s3 = boto3.client('s3')

prefix = 'built-in-algorithm-training'

df_train = pd.read_csv(os.path.join(DATA_FOLDER_PATH, "train"), dtype="object")

データの前処理を行います。今回はデータを train ,validation, test に分割して、OneHotEncoding を行います。

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder

feature_columns = ['C1', 'banner_pos', 'site_category', 'app_category', 'device_type', 'device_conn_type', 'C15', 'C16', 'C18']

df_train, df_test = train_test_split(df_train, train_size=0.7, random_state=0, shuffle=True)
df_train, df_validation = train_test_split(df_train, train_size=0.8, random_state=0, shuffle=True)

one_hot_encoder = OneHotEncoder(handle_unknown='ignore')

y_train = df_train['click'].to_numpy()
X_train = one_hot_encoder.fit_transform(df_train[feature_columns]).toarray()

y_validation = df_validation['click'].to_numpy()
X_validation = one_hot_encoder.transform(df_validation[feature_columns]).toarray()

y_test = df_test['click'].to_numpy()
X_test = one_hot_encoder.fit_transform(df_test[feature_columns]).toarray()

線形学習アルゴリズムがサポートしているrecordIO-wrapped protobuf 形式に変換し、S3へ保存します

import numpy as np 
import io
import sagemaker.amazon.common as smac

def upload_protobuf_to_s3(data_type: str, X: np.ndarray, y: np.ndarray) -> None:
    file_name = f'{data_type}.data'
    f = io.BytesIO()
    smac.write_numpy_to_dense_tensor(f, X.astype("float32"), y.astype("float32"))
    f.seek(0)
    
    boto3.Session().resource("s3").Bucket(bucket).Object(
        os.path.join(prefix, data_type, file_name)
    ).upload_fileobj(f)
    
    
upload_protobuf_to_s3('train', X_train, y_train)
upload_protobuf_to_s3('validation', X_validation, y_validation)
upload_protobuf_to_s3('test', X_test, y_test)

S3にデータが保存されていることを確認します。

学習ジョブ名と、train, validation, test データのS3パス、学習時の生成ファイルの出力先パスの設定を保持しておきます

job_name = "built-in-linear-learner-ctr-prediction-" + strftime("%Y%m%d-%H-%M-%S", gmtime())

output_location = f"s3://{bucket}/{prefix}/output"
s3_train_data = f"s3://{bucket}/{prefix}/train/{train_file}"
s3_validation_data = f"s3://{bucket}/{prefix}/validation/{validation_file}"
s3_test_data = f"s3://{bucket}/{prefix}/test/{test_file}"

組み込みアルゴリズムの学習の実行方法は3つあります

SageMaker SDKのEstimator クラスに、AWSが提供する線形学習アルゴリズム用のdocker imageを渡して実行する方法
AWS SDK で、AWSが提供する線形学習アルゴリズム用のdocker imageを指定して学習ジョブを作成する方法
SageMaker SDKのLinearLearnerクラスを使用する方法

SageMaker SDKのEstimator クラスを利用する方法

SageMaker SDKのimage_uris.retrieve を用いて組み込みアルゴリズムのdocker imageのURIを取得できます。

AWSが提供しているdocker imageの一覧はこちらにまとまっています

アジアパシフィック (東京) (ap-northeast-1) の Docker レジストリパスとサンプルコード - アマゾン SageMaker

取得したdocker image uriをEsimatorクラスに渡すことで、組み込みアルゴリズムを利用できます。

AWS提供の線形学習アルゴリズムの内部では以下の処理が行われます

入力データの正規化
validation データを用いたハイパーパラメータチューニング
test データに対して評価
学習済みモデルをS3にアップロード

線形学習の仕組み - アマゾン SageMaker

set_hyperparameters で学習ハイパーパラメータの値を設定します。

fit メソッドにS3のデータパスを渡すことで学習ジョブを実行できます。train, ,validation, test の3つのデータを渡すことができます。

from sagemaker import image_uris
import sagemaker

# AWSが提供するdocker imageのURIを取得
container = image_uris.retrieve(region=region, framework="linear-learner")
print(container) # 351501993468.dkr.ecr.ap-northeast-1.amazonaws.com/linear-learner:1

# Estimatorクラスにdocker imageを渡す
linear = sagemaker.estimator.Estimator(
    container,
    role,
    instance_count=1,
    instance_type="ml.m5.large",
    output_path=output_location,
    sagemaker_session=sess,
)

# 学習ジョブの実行
linear.set_hyperparameters(predictor_type="binary_classifier", mini_batch_size=200)
linear.fit({"train": s3_train_data, "validation": s3_validation_data, "test": s3_test_data}, job_name=job_name)

SageMaker のAWSコンソール画面で、学習ジョブが実行されていることを確認できます。

CloudWatchで学習ログを確認すると、データが読み込まれモデルの学習が進んでいることを確認できます。

学習ジョブ作成時に設定した出力先S3パスに学習済みモデルが保存されていることを確認できます

参考：SageMaker SDKを用いて組み込みアルゴリズムを実行するサンプルノートブック An Introduction to Linear Learner with MNIST — Amazon SageMaker Examples 1.0.0 documentation

AWS SDKを利用する方法

AWS SDK を利用する場合も、SaeMaker SDK で学習ジョブを実行する場合の流れと一緒です

指定された形式のデータをS3に用意
AWSが提供する docker image を指定して学習を実行

学習データをS3に用意する手順はSageMaker SDKを利用する場合と共通なので、省略します。

学習ジョブの実行はcreate_training_job で行い、SageMaker SDKのEstimatorを利用する場合と同様、学習データのS3パス・AWSが提供する線形学習アルゴリズムのdocker image を設定する必要があります。
SageMaker — Boto3 Docs 1.26.7 documentation

job_name = "built-in-linear-learner-ctr-prediction-aws-sdk" + strftime("%Y%m%d-%H-%M-%S", gmtime())
container = image_uris.retrieve(region=region, framework="linear-learner")

linear_training_params = {
    "RoleArn": role,
    "TrainingJobName": job_name,
    "AlgorithmSpecification": {"TrainingImage": container, "TrainingInputMode": "File"},
    "ResourceConfig": {"InstanceCount": 1, "InstanceType": "ml.m5.large", "VolumeSizeInGB": 10},
    "InputDataConfig": [
        {
            "ChannelName": "train",
            "DataSource": {
                "S3DataSource": {
                    "S3DataType": "S3Prefix",
                    "S3Uri": s3_train_data,
                    "S3DataDistributionType": "ShardedByS3Key",
                }
            },
            "CompressionType": "None",
            "RecordWrapperType": "None",
        },
        {
            "ChannelName": "validation",
            "DataSource": {
                "S3DataSource": {
                    "S3DataType": "S3Prefix",
                    "S3Uri": s3_validation_data,
                    "S3DataDistributionType": "FullyReplicated",
                }
            },
            "CompressionType": "None",
            "RecordWrapperType": "None",
        },
        {
            "ChannelName": "test",
            "DataSource": {
                "S3DataSource": {
                    "S3DataType": "S3Prefix",
                    "S3Uri": s3_test_data,
                    "S3DataDistributionType": "FullyReplicated",
                }
            },
            "CompressionType": "None",
            "RecordWrapperType": "None",
        },
    ],
    "OutputDataConfig": {"S3OutputPath": output_location},
    "HyperParameters": {
        "mini_batch_size": "300",
        "predictor_type": "binary_classifier",
        "epochs": "5",
        "num_models": "1",
        "loss": "absolute_loss",
    },
    "StoppingCondition": {"MaxRuntimeInSeconds": 60 * 60},
}

sm = boto3.client('sagemaker')
sm.create_training_job(**linear_training_params)

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

AWS SDKを用いて組み込みアルゴリズムを実行するサンプルノートブック Breast Cancer Prediction — Amazon SageMaker Examples 1.0.0 documentation

SageMaker SDKのLinearLerner クラスを利用する方法

Estimator クラスを利用する際は、AWS提供のdocker image URIを指定する必要がありましたが、線形学習アルゴリズム用のLinearLernerクラスを使用すれば、docker image の指定を省略することができます。
https://sagemaker.readthedocs.io/en/stable/algorithms/linear_learner.html

内部実装を見ると、線形学習アルゴリズムのdocker image を取得しているので、実行される処理はEstimator を利用する場合と同じです。
sagemaker-python-sdk/linear_learner.py at 255a339ae985041ef47e3a80da91b9f54bca17b9 · aws/sagemaker-python-sdk · GitHub

LinearLernerのように、特定のアルゴリズムに特化したEstimator クラスは限られているので注意してください。

LinearLernerクラスで学習ジョブを実行する場合、線形学習アルゴリズムのdocker imageの取得を内部で行なってくれるため、学習データの準備だけ行います。

LinearLearner を使用する場合、入力データを Record 形式にする必要があります。

https://sagemaker.readthedocs.io/en/stable/algorithms/linear_learner.html#sagemaker.LinearLearner.record_set

学習ジョブ実行時に Record はprotocol buf 形式にシリアライズされ、S3へ自動アップロードされます。

train = linear.record_set(X_train.astype('float32'), labels=y_train.astype('float32'), channel='train')
validation = linear.record_set(X_validation.astype('float32'), labels=y_validation.astype('float32'), channel='validation')
test = linear.record_set(X_test.astype('float32'), labels=y_test.astype('float32'), channel='test')

Estimatorクラスを利用する場合と同様、LinearLernerクラスを作成し、 fit メソッドで学習ジョブを登録できます。 fit メソッドにデータの Record クラスのリストを渡すことで、学習データを渡すことができます。

job_name = "built-in-linear-learner-ctr-prediction" + strftime("%Y%m%d-%H-%M-%S", gmtime())

linear = sagemaker.LinearLearner(
    role=role,
    train_instance_count=1,
    train_instance_type="ml.m5.large",
    output_path=output_location,
    predictor_type="binary_classifier",
    sagemaker_session=sess
)

linear.fit([train, validation, test], mini_batch_size=200, wait=False, job_name=job_name)

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