Amazon SageMaker Examplesを読んでみた　〜Build and Train Models & MLOps編〜

この記事は[UMITRON Advent Calendar 2024](https://adventar.org/calendars/9995) 13日目の記事です

---

# まえがき
この記事は、[Amazon SageMaker Examplesを読んでみた 〜End To End ML Lifecycle & Prepare Data編〜](https://qiita.com/hosonot/items/fe497249e39f9b3b6c1a)の続きです。
この記事がどのようなものかについては、そちらをご参照ください。

先日は[Deploy and Monitorについてまとめました](https://qiita.com/hosonot/items/e089f0ecb4ee10737cea)。

当初、本日はBuild and Train Modelsについてのみまとめる予定でしたが、後述する理由により紹介するサンプルが少なくなってしまったので、MLOpsについてもまとめます。

# Build and Train Models

様々なモデルの学習方法が記載されたサンプル群

個々の組み込みアルゴリズムやフレームワークのコンテナの使い方に関するサンプルが多く、基本的な部分はこれまでのものと重複するので、学習過程の可視化と自作アルゴリズムコンテナの作成方法に関するサンプルのみについて記す

## [Visualize Training Jobs and Performance of Your Model Using TensorBoard on SageMaker](https://github.com/aws/amazon-sagemaker-examples/blob/default/%20%20%20%20%20%20build_and_train_models/sm-distributed_data_parallelism_pytorch/sm-distributed_data_parallelism_pytorch.ipynb)

### 概要

- [TensorBoard on Amazon SageMaker](https://docs.aws.amazon.com/sagemaker/latest/dg/tensorboard-on-sagemaker.html)を使って学習ログを可視化する
- PyTorchで分類器を学習し、その学習過程を可視化する

### 主要なAPI呼び出し

| カテゴリ | メソッド | 説明 |
| --- | --- | --- |
| 学習と推論 | sagemaker.pytorch.PyTorch.fit | PyTorchのコンテナを使って推論用モデルを学習。ノートブックでは引数tensorboard_output_configを設定することで、学習ログをリアルタイムでS3にアップロードするようにしている |
| 監視とトラッキング | .sagemaker.interactive_apps.tensorboard.TensorBoardApp.get_app_url | TensorBoradアプリケーションへのURLを生成する。SageMaker Studioで実行している場合はURLが生成され、それ以外の場合はSageMakerコンソールへのURLが生成される |

### 所感

- サポートされいるフレームワークは、TensorFlow、PyTorch、Hugging Face Transformersで、最近使われているものはカバーされている
- Cloud watchを使って学習過程を可視化するよりも簡単にできるかもしれないと感じた
    - 自作コンテナで学習する場合も、[sagemaker.estimator.Estimator](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html#sagemaker.estimator.Estimator)の初期化時にtensorboard_output_configが設定できるので、[torch.utils.tensorboard](https://pytorch.org/docs/stable/tensorboard.html)でログを適切なディレクトリに出力すればできそうだがサンプルは見つからなかった

## [Building your own container as Algorithm / Model Package](https://github.com/aws/amazon-sagemaker-examples/blob/default/%20%20%20%20%20%20build_and_train_models/sm-marketplace_building_your_own_container_as_package/sm-marketplace_building_your_own_container_as_package.ipynb)

### 概要

- SageMakerに独自アルゴリズムを実装する
    - 独自アルゴリズムを動作させるdocker imageをECRに登録
    - ローカル or ノートブック上で上記の動作を確認する
    - トレーニングジョブを実行し、バッチ変換でテスト
    - 推論用エンドポイントをデプロイ
- 独自アルゴリズムをパッケージ化して登録する
- 登録したアルゴリズムをファインチューニングして推論用エンドポイントをデプロイ
- AWSマーケットプレイスに出品する

### 主要なAPI呼び出し

| カテゴリ | メソッド | 説明 |
| --- | --- | --- |
| データ準備 | sagemaker.session.Session.upload_data | S3にデータをアップロードする。boto3でアップロードする場合と異なり、指定が無い場合デフォルトのバケットが選択される（無ければ作成される）。また、戻り値がS3 URIになっており、そのまま他のsagemaker関連の設定に使える。 |
| 学習と推論 | sagemaker.estimator.Estimator.fit | トレーニングジョブを実行。インフラ周りの設定はEstimatorインスタンス作成時に行う |
| 学習と推論 | sagemaker.estimator.Estimator.transformer.transform | バッチ変換ジョブを実行 |
| インフラ管理 | sagemaker.estimator.Estimator.create_model | 推論用のモデルを作成 |
| インフラ管理 | sagemaker.estimator.Estimator.deploy | 推論用エンドポイントを作成。
ノートブックではインスタンスタイプを指定することでリアルタイム推論でデプロイしている。 |
| 学習と推論 | sagemaker.predictor.Predictor.predict | エンドポイントに推論をリクエスト |
| 学習と推論 | SageMaker.Client.create_algorithm | アルゴリズムを登録。学習アルゴリズムと関連コンテナイメージが含まれる。TrainingSpecificationとInferenceSpecificationにAdditionalS3DataSourceを設定することで、学習の際の初期重みを登録できる（つまり、ファインチューニングできるようになる） |
| 学習と推論 | SageMaker.Client.create_model_package | モデルパッケージを登録。推論モデルと関連コンテナイメージ、モデルデータが含まれる（create_modelで作成されるデプロイ用のモデルとは別物） |
| 学習と推論 | SageMaker.Client.create_training_job | トレーニングジョブを作成 |
| 学習と推論 | SageMaker.Client.create_model | 推論用モデルを作成。ノートブックではモデルパッケージから作成している |
| インフラ管理 | SageMaker.Client.create_endpoint_config | 推論用のエンドポイント設定を作成 |
| インフラ管理 | SageMaker.Client.create_endpoint | 推論用のエンドポイントを作成 |
| 学習と推論 | SageMaker.Client.invoke_endpoint | エンドポイントに推論をリクエスト |

### 所感

- 独自アルゴリズムを実装する場合、大枠はこのノートブックのコピペでよさそう（マーケットプレイスへの出品は除く）
    - アルゴリズムのパッケージ化をする場合は、ファインチューニングとモデルの登録部分は別ファイルにしても良さそう
    - パッケージ化の必要性については、後に紹介するMLOpsのサンプルでは基本的にパッケージ化していることから、パッケージ化するのがデファクトなのかもしれない
        - パッケージ化することで、再現性が担保されやすくなったり、CI/CDパイプラインとの統合が容易になるメリットが大きいのかもしれない
- ローカルでテストをする場合、アーキテクチャの違いで環境構築に苦労することがあるので、SageMaker Notebookの方がスムーズに開発ができそう
    - ローカルでのテストを諦めると、テストのためにトレーニングジョブの実行やエンドポイント作成を毎回する必要があるのでテストに時間がかかる
    - ノートブックインスタンスの使用料金が別途発生する点に注意
- sagemaker.session.Session.upload_dataで作成されるデフォルトバケットは(`sagemaker-<region>-<account_id>`)になるので、アルゴリズム毎にバケットを分ける等、整理したい場合は指定したほうが良さそう
    - Sessionのインスタンスを作成する際にdefault_bucket引数に指定するのが良さそう

# MLOps

MLモデルのためのCI/CDに関する機能のサンプル群

Infrastructure-as-codeを実現するためのSageMaker Pipelinesや、MLflowとの連携によるモデル管理に関するサンプルがある

MLflowはそれをサービングするリソースが別途必要になることから、小規模プロジェクトには不向きと考え省略し、SageMaker Pipelinesの中でも基本的なもののみについて記す

また、1サンプルだけあったLinage Trackingについても記す

## [Model Lineage Tracking](https://github.com/aws/amazon-sagemaker-examples/blob/default/%20%20%20ml_ops/sm-ml_lineage_tracking_model_governance_graph/sm-ml_lineage_tracking_model_governance_graph.ipynb)

### 概要

- Model linageデータをグラフにして可視化する
    - Model linageデータ：データセット、トレーニングジョブ、デプロイ等、機械学習モデルの生成や運用に関わるプロセスを追跡・記録し、依存関係や履歴を明確にするための情報
        - データは以下4つのカテゴリで構成されている
            - Artifacts:： データ、モデル、メトリクス、ログ等の結果物や中間生成物
            - Trial components：データ準備、モデル学習、評価など、トレーニングの各ステージを構成する要素
            - Actions：各操作の実行履歴（例：ジョブ実行、モデルデプロイ、エンドポイント更新）
            - Contexts：特定のエンティティ（例：エンドポイント、モデルパッケージ）や論理的プロセス（例：モデル学習、評価）をグループ化したもの
- グラフの可視化は[Neptune ML](https://docs.aws.amazon.com/neptune/latest/userguide/machine-learning.html)で行う

### 主要なAPI呼び出し

| カテゴリ | メソッド | 説明 |
| --- | --- | --- |
| 監視とトラッキング | sagemaker.lineage.artifact.Artifact.list | Artifactsの一覧を取得 |
| 監視とトラッキング | sagemaker.lineage.trial_component.TrialComponent.list | Trial componentsの一覧を取得 |
| 監視とトラッキング | sagemaker.lineage.action.Action.list | Actionsの一覧を取得 |
| 監視とトラッキング | sagemaker.lineage.context.Context.list | Contextsの一覧を取得 |
| 監視とトラッキング | sagemaker.lineage.association.Association.list | 上記4つの構成要素間のエッジを取得 |

### 所感

- これらの関連性はWebコンソールからも確認できるものが多い気がするが、一望するのは難しいため、モデル開発の全体像を視覚的に表現するには良い機能かもしれない

## [Orchestrate Jobs to Train and Evaluate Models with Amazon SageMaker Pipelines](https://github.com/aws/amazon-sagemaker-examples/blob/default/%20%20%20ml_ops/sm-pipelines_preprocess_train_evaluate_batch_transform/sm-pipelines_preprocess_train_evaluate_batch_transform.ipynb)

### 概要

- SageMaker Pipelinesでデータ準備からエンドポイントのデプロイまでの一連の流れを定義する
    - 処理の流れは以下の画像の通り （ノートブックより引用）
    ![68747470733a2f2f7261772e67697468756275736572636f6e74656e742e636f6d2f6177732f616d617a6f6e2d736167656d616b65722d6578616d706c65732f6d61696e2f736167656d616b65722d706970656c696e65732f746162756c61722f6162616c6f6e655f6275696c645f747261696e5f646570.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2371662/d4988100-5705-ee54-7893-754480d90879.png)
    
    - SageMaker Pipelinesを使うことで、各ステップの依存関係を明示的に定義できる
- [UCI Machine Learning Abalone Dataset](https://archive.ics.uci.edu/ml/datasets/abalone) でアワビの年齢推定器をXGBoostで学習し、バッチ変換としてデプロイ
- 類似例 [SageMaker Pipelines](https://github.com/aws/amazon-sagemaker-examples/blob/default/%20%20%20ml_ops/sm-pipelines_train_model_registry_deploy.ipynb)：[California Housing dataset](https://lib.stat.cmu.edu/datasets/)で住宅価格予測器をニューラルネットで学習し、バッチ変換としてデプロイ

### 主要なAPI呼び出し

| カテゴリ | メソッド | 説明 |
| --- | --- | --- |
| パイプライン構築 | sagemaker.sklearn.processing.SKLearnProcessor.run | scikit-learn用のコンテナでSageMaker processing jobを実行。ノートブックでは、SKLearnProcessorインスタンス作成時にsagemaker.workflow.pipeline_context.PipelineSessionを渡すことで、この関数呼び出し時に処理は実行されず、ジョブに渡される引数を定義し、それをもとにsagemaker.workflow.steps.ProcessingStepを作成することでステップを定義している |
| パイプライン構築 | sagemaker.estimator.Estimator.fit | トレーニングジョブを実行。ノートブックでは、上記同様、Estimatorインスタンス作成時にPipelineSessionを渡すことで引数を定義し、それをもとにステップを定義している |
| パイプライン構築 | sagemaker.processing.ScriptProcessor.run | 処理ジョブを実行。ノートブックでは、上記同様、ScriptProcessorインスタンス作成時にPipelineSessionを渡すことで引数を定義し、それをもとにステップを定義している |
| パイプライン構築 | sagemaker.model.Model.deploy | 推論用モデルをデプロイ。ノートブックでは、上記同様、Modelインスタンス作成時にPipelineSessionを渡すことで引数を定義し、それをもとにステップを定義している |
| パイプライン構築 | sagemaker.workflow.pipeline.Pipeline.upsert | パイプラインの作成。インスタンス作成時に、各ステップの定義と依存関係を登録している |
| パイプライン構築 | sagemaker.workflow.pipeline.Pipeline.start | パイプラインの実行 |
| パイプライン構築 | sagemaker.workflow.pipeline._PipelineExecution.describe | パイプラインの状態確認。_PipelineExecutionインスタンスは上記upsert関数の戻り値として取得する |
| パイプライン構築 | sagemaker.workflow.pipeline._PipelineExecution.wait | パイプラインの完了を待つ |
| パイプライン構築 | sagemaker.workflow.pipeline._PipelineExecution.list_steps | パイプラインの各ステップの進行状況や結果を確認 |

### 所感

- 処理フローをパイプラインとして定義することで、ワークフロー全体をコード化して管理することが徹底されるのもメリットだと感じた
    - 手動プロセスとの混合が避けられ、手順全体をバージョン管理できる
- デプロイステップをLambdaで実行する例が[SageMaker Pipelines Lambda Step](https://github.com/aws/amazon-sagemaker-examples/blob/default/%20%20%20ml_ops/sm-pipelines_lambda_step/sm-pipelines_lambda_step.ipynb)にあるが、Lambdaにするメリットはこの例のみからでは理解できなかったが、以下のようなときはLambdaを使う必要があるかもしれない
    - 他のAWSサービスでPipeline外からもデプロイをトリガーする必要があるとき
    - AWS外のリソースに変更を加える必要があるとき

## [Basic Pipeline for Batch Inference using Low-code Experience for SageMaker Pipelines](https://github.com/aws/amazon-sagemaker-examples/blob/default/%20%20%20ml_ops/sm-pipelines_batch_inference_step_decorator/sm-pipelines_batch_inference_step_decorator.ipynb)

### 概要

- SageMaker Pipelinesの`@step`デコレータを使ってモデル構築のパイプラインを構築
    - `@step`デコレータ：SageMaker Pipelines内でステップを定義。Python関数をステップとして登録できる
        - ノートブックでは、以下３つのパラメータを設定している
            - name：ステップの名前
            - instance_type：使用するインスタンスタイプ
            - keep_alive_period_in_seconds：次のステップで利用されるまでインスタンスを生存させる秒数（これにより、以降のステップでインスタンスの起動を省略できる）
        - これにより、より簡素にパイプラインを構築できる
- 乳がんの診断データから良性か悪性かを予測分類器をXGBoostで学習し、バッチ変換としてデプロイ

### 主要なAPI呼び出し

| カテゴリ | メソッド | 説明 |
| --- | --- | --- |
| インフラ管理 | sagemaker.serve.builder.model_builder.ModelBuilder.build().register | モデルパッケージの作成と登録 |
| インフラ管理 | SageMaker.Client.describe_model_package | モデルパッケージの情報を取得。ノートブックではモデルアーティファクトのS3URIを取得している |
| パイプライン構築 | sagemaker.workflow.pipeline.Pipeline.upsert | パイプラインの作成。インスタンス作成時に、`@step`デコレータを付与して定義した関数とその依存関係を登録する |
| パイプライン構築 | sagemaker.workflow.pipeline.Pipeline.start | パイプラインの実行 |
| パイプライン構築 | sagemaker.workflow.pipeline._PipelineExecution.describe | パイプラインの状態確認。_PipelineExecutionインスタンスは上記upsert関数の戻り値として取得する |
| パイプライン構築 | sagemaker.workflow.pipeline._PipelineExecution.wait | パイプラインの完了を待つ |
| パイプライン構築 | sagemaker.workflow.pipeline._PipelineExecution.list_steps | パイプラインの各ステップの進行状況や結果を確認 |

### 所感

- デコレータを使うことで、各処理ステップを独立した関数として記述することが強制されるので、再利用性が高まるだけでなくコードの保守性も上がるように感じた
    - 個人的には、`@step`デコレータを使った例の方が各ステップの入出力の関係を追いやすかった
- [開発ガイド](https://docs.aws.amazon.com/sagemaker/latest/dg/pipelines-step-decorator.html)に記されているように、可能ならローカル（ノートブックインスタンス上）でテストし、次に`@remote`デコレータでリモート環境でユニットテスト、`@step`デコレータでパイプラインに組み込むという手順を踏めばスムーズに開発を進められそう

## [Use SageMaker Pipelines With Step Caching](https://github.com/aws/amazon-sagemaker-examples/blob/default/%20%20%20ml_ops/sm-pipelines_step_caching/sm-pipelines_step_caching.ipynb)

### 概要

- SageMaker Pipelinesのキャッシュを有効にする方法とそのテスト
- キャッシュを有効にすると各ステップの引数が追跡され、引数が一致する場合はキャッシュを使用する様になる。ただし、処理結果に影響の無い引数は追跡されない
- サンプルでは、インスタンスタイプを変更するとキャッシュされた結果が使われるが、前処理のスクリプトが変更されるとキャッシュが使われないことを試している
    - キャッシュの動作確認が目的のため、前処理ステップとトレーニングステップのみのパイプラインを定義している

### 主要なAPI呼び出し

| カテゴリ | メソッド | 説明 |
| --- | --- | --- |
| パイプライン構築 | sagemaker.sklearn.processing.SKLearnProcessor.run | scikit-learn用のコンテナでSageMaker processing jobを実行。ノートブックでは、SKLearnProcessorインスタンス作成時にsagemaker.workflow.pipeline_context.PipelineSessionを渡すことで、この関数呼び出し時に処理は実行されず、ジョブに渡される引数を定義し、それをもとにsagemaker.workflow.steps.ProcessingStepを作成することでステップを定義している |
| パイプライン構築 | sagemaker.estimator.Estimator.fit | トレーニングジョブを実行。ノートブックでは、上記同様、Estimatorインスタンス作成時にPipelineSessionを渡すことで引数を定義し、それをもとにステップを定義している |
| パイプライン構築 | sagemaker.workflow.pipeline.Pipeline.upsert | パイプラインの作成。インスタンス作成時に、各ステップの定義と依存関係を登録している |
| パイプライン構築 | sagemaker.workflow.pipeline.Pipeline.start | パイプラインの実行 |
| パイプライン構築 | sagemaker.workflow.pipeline._PipelineExecution.describe | パイプラインの状態確認。_PipelineExecutionインスタンスは上記upsert関数の戻り値として取得する |
| パイプライン構築 | sagemaker.workflow.pipeline._PipelineExecution.wait | パイプラインの完了を待つ |
| パイプライン構築 | sagemaker.workflow.pipeline._PipelineExecution.list_steps | パイプラインの各ステップの進行状況や結果を確認 |
| パイプライン構築 | sagemaker.workflow.pipeline.Pipeline.update | パイプラインの更新 |

## [SageMaker Pipeline - Local Mode](https://github.com/aws/amazon-sagemaker-examples/blob/default/%20%20%20ml_ops/sm-pipelines_local_mode/sm-pipelines_local_mode.ipynb)

### 概要

- SageMaker Pipelinesの各ジョブをローカルで実行する
    - SageMaker Studioではローカル実行ができない
- PipelineSessionクラスの代わりにLocalPipelineSessionを使用する
- これまでの例と異なり、FrameworkProcessorを使うことでジョブが実行されるコンテナに追加のPythonパッケージをインストールしている

### 主要なAPI呼び出し

| カテゴリ | メソッド | 説明 |
| --- | --- | --- |
| パイプライン構築 | sagemaker.processing.FrameworkProcessor.get_run_args | ProcessingStepで実行するジョブの入出力を取得。ノートブックでは、引数source_dirに実行するスクリプトとrequirements.txtを圧縮したファイルのS3パスを指定している |
| パイプライン構築 | sagemaker.workflow.pipeline.Pipeline.upsert | パイプラインの作成。インスタンス作成時に、各ステップの定義と依存関係を登録している |
| パイプライン構築 | sagemaker.workflow.pipeline.Pipeline.start | パイプラインの実行 |
| パイプライン構築 | sagemaker.workflow.pipeline._PipelineExecution.describe | パイプラインの実行 |
| パイプライン構築 | sagemaker.workflow.pipeline._PipelineExecution.list_steps | パイプラインの各ステップの進行状況や結果を確認 |

### 所感

- `@step`デコレータを使う場合はインスタンスの指定が必須なので、単に`@step`デコレータを外してローカルでテストすればよいのかもしれない（LocalPipelineSessionもあくまでクラウドで実行する前のテストという立ち位置）
    - この方法ならSageMaker Studioでも実行できるので、テストのしやすさという観点でも`@step`デコレータの方が便利かもしれない

---
ウミトロンは、「持続可能な水産養殖を地球に実装する」というミッション実現に向けて、日々プロダクト開発・展開にチーム一丸となって邁進しています。
ウミトロンのニュースや活動状況を各種SNSで配信していますので、ぜひチェックいただき、来年も応援よろしくお願いします！
Facebook https://www.facebook.com/umitronaqtech/
X https://x.com/umitron
Instagram https://www.instagram.com/umitron.aqtech/
Linkedin https://www.linkedin.com/company/umitron