2024.08.16

【SageMaker Model Registry】AWS上での機械学習モデル管理入門

Champ

どうも、Champ です 🙌

先日、髪の毛を青に染めたんですが、思った様に色が入らず、黒っぽい感じになっちゃいました 🥺
単に青色と言っても、色の濃さや薬剤につける時間など人に応じた調整が必要らしく、今回はその調整が微妙だったようです。。。

最近、AWS の機械学習サービスにハマっていて、毎日のようにいじくり回しているのですが、機械学習モデルもカラーリングと同じで、パラメータをいじって実験しているとどれが一番良かったか分からなくなっちゃいますよね😅

そこで今回は、Amazon SageMaker Model Registry について掘り下げてみたいと思います。こちらは、機械学習モデルの管理を効率化してくれるツールです！

Contents

SageMaker Model Registry って何？
なぜモデルバージョン管理が必要なの？
SageMaker Model Registry でできること
モデル登録の手順
疑問：SKLearn クラスの fit メソッドで output_path を指定していなくても実行できる理由は？
疑問：モデルレジストリの詳細にコンテナ？
まとめ

SageMaker Model Registry って何？

簡単に言うと、機械学習モデルのライフサイクル管理をサポートする AWS のサービスです。主な機能としては：

モデルのバージョン管理
メタデータの関連付け
承認ステータスの管理
本番環境へのデプロイ支援

モデルグループを作って、その中に複数のモデルバージョンを登録できます。
これにより、モデルの変更履歴をしっかり追跡できるわけです。

なぜモデルバージョン管理が必要なの？

ここで、ちょっと考えてみましょう。モデルのバージョン管理をしないと、どんな「つらさ」が待っているでしょうか？

本番環境のモデルのバージョンが分からなくなる
性能が良いモデルのトレーニングコードが見つからない
同じトレーニングコードを使用しても、環境が違うため失敗する

機械学習モデルは複数作成し、性能の良いモデルを選ぶことは多々あるので適当な管理をしていたら大変なことになりそうですね。。。

SageMaker Model Registry でできること

それでは、具体的に SageMaker Model Registry で何ができるのか？
以下の５つのことが可能です。

モデルのバージョン管理：変更履歴をしっかり追跡
メタデータの関連付け：学習データやハイパーパラメータなどの情報を記録
承認ステータスの管理：モデルの品質や本番環境への準備状況を管理
コレクションによる整理：関連するモデルを階層的にグループ化
クロスアカウントでのモデル共有：AWS RAM を使って他のアカウントとも共有可能

モデルのバージョン管理だけでなく、クロスアカウントでの共有や承認ステータスの管理などもあり、一通りのことはできそうですね。

モデル登録の手順

それでは、実際に SageMaker Model Registry にモデルを登録する手順を紹介します。
今回は簡単な scikit-learn モデルを作成して登録してみます。

1.モデル作成に使用するトレーニングコード（train.py）を作成

詳細は省略しますが、デモ用に簡単な分類問題を解くトレーニングコードを生成しています。

import argparse
import os
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import joblib
import json

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])
    args, _ = parser.parse_known_args()

    # テストデータの生成
    np.random.seed(42)
    X = np.random.rand(1000, 4)
    y = (X[:, 0] + X[:, 1] > 1).astype(int)

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # モデルの学習
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)

    # モデルの評価
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print(f'Accuracy: {accuracy}')

    # モデルの保存
    joblib.dump(model, os.path.join(args.model_dir, 'model.joblib'))

    # メトリクスの保存
    with open(os.path.join(args.model_dir, 'metrics.json'), 'w') as f:
        json.dump({"accuracy": accuracy}, f)

2. SageMaker Studio を開いて新しいノートブックを作成

詳細な手順は割愛しますが、Sagemaker Studio Domainを作成した後、任意のユーザでログインし、下記画面からノートブックを作成します。

以降の手順ではノートブックに記述するコードを記述します。

3. 必要なライブラリをインポート

import sagemaker
from sagemaker import get_execution_role
from sagemaker.sklearn.estimator import SKLearn

4. SageMaker セッションを作成し、実行ロールを取得

sagemaker_session = sagemaker.Session() role = get_execution_role()

下記の様なレスポンスが得られればOKです。

5. SageMaker の Estimator を設定してモデルをトレーニング

role = get_execution_role()
sklearn_estimator = SKLearn(
    entry_point='train.py',
    role=role,
    instance_type='ml.m5.large',
    framework_version='0.23-1',
    py_version='py3',
    instance_count=1
)
sklearn_estimator.fit()

上記のコードでは SKLearn エスティメータを設定しています。これは、scikit-learn を使用したモデルのトレーニングジョブを定義します。

entry_point: トレーニングコードのファイル名
role: 使用する IAM ロール
instance_type: トレーニングに使用するインスタンスタイプ
framework_version: 使用する scikit-learn のバージョン
py_version: Python のバージョン
instance_count: 使用するインスタンス数
output_path: トレーニング結果の保存先 S3 パス ※省略可

6. 設定したエスティメータを使用してモデルのトレーニングを開始

sklearn_estimator.fit()

7. モデルメトリクス（評価指標）のオブジェクトを作成

model_metrics = ModelMetrics(
    model_statistics=MetricsSource(
        s3_uri=f'{sklearn_estimator.output_path}/{sklearn_estimator.latest_training_job.job_name}/output/metrics.json',
        content_type='application/json'
    )
)

上記のコードでは、トレーニング中に生成されたメトリクスファイルの S3 パスを指定しています。

8. Model オブジェクトを作成

model = Model(
    image_uri=sklearn_estimator.image_uri,
    model_data=sklearn_estimator.model_data,
    role=role,
    sagemaker_session=sagemaker_session
)

9. モデルレジストリにモデルを登録

model.register(
    content_types=['text/csv'],
    response_types=['text/csv'],
    inference_instances=['ml.t2.medium', 'ml.m5.large'],
    transform_instances=['ml.m5.large'],
    model_package_group_name='SKLearnModelPackageGroup',
    model_metrics=model_metrics
)

上記のコードは、トレーニングされたモデルを SageMaker Model Registry に登録するためのものです。
各パラメータの説明は以下の通りです：

content_types: モデルが受け入れる入力データの形式
response_types: モデルが返す出力データの形式
inference_instances: リアルタイム推論に使用可能なインスタンスタイプのリスト
transform_instances: バッチ変換ジョブに使用可能なインスタンスタイプのリスト
model_package_group_name: このモデルを追加する Model Package Group の名前
model_metrics: モデルの性能メトリクスを含むディクショナリ

これらのパラメータを設定することで、モデルの使用方法や制約、性能などの情報を Model Registry に記録し、後でモデルを簡単に検索、比較、デプロイできるようになります。

Sagemaker StudioのModelsに以下の様に登録されていればOKです。

手順は以上です。

疑問：SKLearn クラスの fit メソッドで output_path を指定していなくても実行できる理由は？

SKLearn クラスの fit メソッドで output_path を指定していなくても実行できる理由は、SageMaker が自動的にデフォルト S3 バケットを設定するためです。
以下に詳細を説明します：

S3 バケットの自動生成：
明示的にバケットを指定しておらず、Sagemaker 用デフォルト S3 バケットが存在しない場合、デフォルトの S3 バケットが自動的に生成されます。
このバケット名は通常、sagemaker-{リージョン}-{アカウントID}の形式になり、明示的に指定しない場合はこの S3 バケットにモデルアーティファクトを保存します。
出力プレフィックスの自動生成：
トレーニングジョブごとに一意のプレフィックスが自動的に生成されます。このプレフィックスはトレーニングジョブ名に基づいています。

この機能により、ユーザーは出力パスを明示的に指定しなくても、SageMaker でモデルのトレーニングと保存を行うことが可能です。