新しくなった Amazon SageMaker Experiments で実験管理

## はじめに
Amazon SageMaker ユーザにはお馴染み（のはず）の Experiments ですが、これが [大幅にアップデート](https://aws.amazon.com/jp/blogs/machine-learning/next-generation-amazon-sagemaker-experiments-organize-track-and-compare-your-machine-learning-trainings-at-scale/) されました！
個人的に期待がつのるアップデートなので、早速どこが便利になったのかを確かめてみようと思います。

2022/12/20 追記：[公式のサンプルコード](https://github.com/aws/amazon-sagemaker-examples/tree/main/sagemaker-experiments) が公開されました！

## アップデートのポイント
まず一番大きいのが、記録した内容を表示するダッシュボード機能ができたことです！今までは Pandas の DataFrame としてノートブックで表示するくらいしかできませんでしたが、Amazon SageMaker Studio からこんなふうにみることができます。この例では、mnist-hand-written-digits-classification という Experiment の中に run1, run2... という Run が登録されています。今まで Trial といっていたものが Run になった感じかなと。（Experiment の中身を describe してみると Trial や TrialComponent として記録されているようです。）

![スクリーンショット 2022-12-18 9.35.00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/1c724595-5c6d-52b0-dc82-a0eacd00638b.png)

Run をクリックすると詳細画面を確認できます。Charts のところを見るとこんな感じでグラフを確認することができます。この例では、train と test の loss は自分で手動で追加しましたが、混同行列は学習スクリプトの中で `run.log_confusion_matrix` という API を呼ぶだけで勝手に SageMaker が作ってくれました。

![スクリーンショット 2022-12-18 9.37.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/818f0e44-9c2d-3895-cb34-f76754414ee0.png)

Run 一覧画面で複数の Run を選択して Analyze ボタンをクリックすると、以下のように Run を比較することができます。チャートは自分で追加しました。

![スクリーンショット 2022-12-18 9.42.24.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/0d0f4e8e-2cdf-e546-c068-435d830f082e.png)

## どうやって使うのか

基本的には学習スクリプトの中に Experiments の API を呼ぶコードを追加する形です。学習スクリプトを SageMaker 学習ジョブとして実行する場合は Estimator の作成方法も少し変える必要があります。

### 注意
はじめに、新しい Experiments を使うために気をつけておくポイントを整理しておきます。

#### 2.124.0 以降の SageMaker Python SDK を利用する

新しい機能なので SageMaker Python SDK のバージョンアップも必要です。学習スクリプトを実行する環境の SageMaker Python SDK が古い場合はエラーになりますので、2.124.0 以降のバージョンにアップグレードしてください。また、従来は sagemaker-experiments というモジュールをインストールする必要がありましたが、Experiments の機能が SageMaker Python SDK に統合されたことにより sagemaker-experiments のインストールは不要になりました。

#### Estimator 作成時に AWS_DEFAULT_REGION をセットする
学習ジョブの中で Experiments を使う場合、これを忘れると以下のエラーが出ます。
```console
ValueError: Must setup local AWS configuration with a region supported by SageMaker.
```

このエラーを回避するには、Estimator 作成時のパラメタ `environment` を以下のように設定を追加します。リージョン名は実際にお使いのものに書き換えてください。
```python
environment={
    "AWS_DEFAULT_REGION": "us-east-1"
},
```

#### Experiment の新規作成が不要に
これは注意というほどのものではないのですが、従来は `Experiment.create` API で Experiment を作成してから Trial を登録する流れでしたが、新しい Experiments では with Run(...) の部分で指定された名前の Experiment がない場合は自動的に作成してくれるようになりました。地味に便利！

## やってみる
~~まだサンプルコードが公開されていないようなので、[こちらのサンプルコード](https://github.com/aws-samples/aws-ml-jp/tree/main/sagemaker/hpo-pytorch-mnist) をベースに Experiments を試してみましょう。~~

2022/12/20 追記：[公式のサンプルコード](https://github.com/aws/amazon-sagemaker-examples/tree/main/sagemaker-experiments) が公開されました！

SageMaker Studio の Python3 (PyTorch 1.12 Python 3.8 CPU Optimized カーネルと ml.t3.medium の組み合わせで実行しました。

### ソースディレクトリを作成してファイルを移動

前述の通り 2.124.0 以降の SageMaker Python SDK が必要なのですが、今回使用するビルトインコンテナで使用している SDK のバージョンが古いので、それをアップグレードするために以下の requirements.txt を作成します。

```python:requirements.txt
sagemaker>=2.124.0
```

この requirements.txt を使うためには学習スクリプトと requirements.txt を同じディレクトリに入れておく必要があるので、その準備をします。root 直下に src というディレクトリを作成してその中に requirements.txt を配置し、root にある mnist.py も src 以下に移動します。最終的に、以下のようなディレクトリ構成になります。

```console
root/
├── src/
| ├── mnist.py
| └── requirements.txt
└── pytorch_mnist.ipynb
```

### mnist.py の変更

mnist.py に Experiments のコードを追加します。

モジュールを import しているあたりに以下のコードを追加します。
```python
from sagemaker.experiments import load_run
```

次に、一番下にある引数を受け取る部分に以下のコードを追加します。ここでは、Experiment 名と Run 名を取得しています。

```python
parser.add_argument('--experiment-name', type=str, default=None,
                    help='experiment name')
parser.add_argument('--run-name', type=str, default=None,
                    help='run name')
```

次に、以下のような `eval()` 関数を作成します。これは、モデルの学習が終わった後にテストデータを使って推論を実行し、`run.log_confusion_matrix` API で混同行列を作って Experiment に記録するための関数です。

```python
def eval(model, test_loader, device, run):
    model.eval()
    test_loss = 0
    correct = 0
    
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, size_average=False).item()  # sum up batch loss
            pred = output.max(1, keepdim=True)[1]  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()
            run.log_confusion_matrix(target, pred, "Confusion-Matrix-Test-Data")

    test_loss /= len(test_loader.dataset)
    test_accuracy = 100. * correct / len(test_loader.dataset)
    logger.info('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        test_accuracy))
```

次に、`train()` の中の `for epoch in range(1, args.epochs + 1):`（epoch のループが始まる部分）の手前に以下のコードを追加します。ここでは、`load_run` を使って、学習ジョブ実行時に指定された Experiment と Run の情報を取得して、学習データやハイパーパラメタなどのモデルに関連する情報を登録しています。なお、Estimator の hyperparameter に設定したパラメタは自動的に記録されるため、`run.log_parameter` などで再度登録する必要はありません。

```python
experiment_name = args.experiment_name
run_name = args.run_name

with load_run(experiment_name=experiment_name, run_name=run_name) as run:
    run.log_parameters({
        "num_train_samples": len(train_loader.sampler),
        "num_test_samples": len(test_loader.sampler)
    })

    # log the parameters of your model
    run.log_parameter("device", "cpu")
    run.log_parameters({
        "data_dir": args.data_dir,
        "optimizer": 'SGD',
        "epochs": args.epochs
    })
```

ループの中で `test()` を呼び出す部分がありますが、ここを以下のように書き換えます。ここでは、学習データで算出した loss の値を記録しています。

```python
with load_run(experiment_name=experiment_name, run_name=run_name) as run:
    test(model, test_loader, device, run, epoch)
    run.log_metric(name="train:loss", value=loss.item(), step=epoch)
```

さらに、`save_model()` を呼び出している部分の手前に 2行追加します。

```python
with load_run(experiment_name=experiment_name, run_name=run_name) as run:
    eval(model, test_loader, device, run)
save_model(model, args.model_dir)   // この行は元からある
```

最後に、`test()` を以下のように書き換えます。ここでは、テストデータで算出した loss と accuracy の値を登録します。

```python
def test(model, test_loader, device, run, epoch):
    model.eval()
    test_loss = 0
    correct = 0
    
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, size_average=False).item()  # sum up batch loss
            pred = output.max(1, keepdim=True)[1]  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    test_accuracy = 100. * correct / len(test_loader.dataset)
    logger.info('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        test_accuracy))

    run.log_metric(name="test:loss", value=test_loss, step=epoch)
    run.log_metric(name="test:accuracy", value=test_accuracy, step=epoch)
```

ここまでで、学習スクリプト mnist.py の変更は終了です。

### 最新の SageMaker Python SDK をインストール

次にノートブックを変更していきます。ノートブックの初めに以下のコードを記述したセルを追加して実行します。

```console
!pip install --upgrade sagemaker
```

### Estimator 作成部分の変更

「4.2.Estimatorの定義」と書かれた Estimator を作成するセルの上までセルを実行していきます。以下のコードが書かれたセルを追加および実行して、Run 名を作るのに使用するカウンタを初期化します。

```python
count = 1
```

次に、Estimator を作成しているセルを、以下のように変更してから実行します。一番下にあるように、**fit() を with Run() の中に入れるのを忘れないようにしてください。** fit() を入れ忘れても単体で Run の中身を確認することはできるのですが、複数の Run を比較することができません。

```python
from sagemaker.pytorch import PyTorch
from sagemaker.session import Session
from sagemaker.experiments.run import Run

run_name = 'run' + str(count)
count = count + 1

if on_studio:
    instance_type = 'ml.m5.xlarge'
else:
    instance_type = 'local'

experiment_name = "mnist-hand-written-digits-classification-2"

with Run(experiment_name=experiment_name, sagemaker_session=Session(), run_name=run_name) as run:
    print(run_name)
    estimator = PyTorch(entry_point="mnist.py",
                        source_dir='src',
                        role=role,
                        framework_version='1.12.0',
                        py_version='py38',
                        instance_count=1,
                        instance_type=instance_type,
                        hyperparameters={
                            'batch-size':128,
                            'lr': 0.01,
                            'epochs': 20,
                            'backend': 'gloo',
                            'experiment-name': experiment_name,
                            'run-name': run_name
                        },
                        environment={
                            "AWS_DEFAULT_REGION": "us-east-1"
                        },
                        keep_alive_period_in_seconds=1800
                       )
    estimator.fit({'training': inputs})
```
なお、今回は`keep_alive_period_in_seconds` というパラメタを設定することで、別の新機能である[学習ジョブの Warm pool](https://docs.aws.amazon.com/sagemaker/latest/dg/train-warm-pools.html) 機能を使っています。Warm pool 機能は、一旦学習ジョブを作成したらその際に使用したインスタンスを停止せず hold しておき、次に同じ設定（ハイパーパラメタは例外）の Estimator から学習ジョブが作成されたらそのインスタンスを再度使えるというものです。これにより、`keep_alive_period_in_seconds` で指定した時間内の試行錯誤の際にインスタンスの起動時間を待つ必要がなくいため、すばやく試行錯誤のサイクルを回すことができます。なお、Warm pool を使用するには、インスタンスタイプごとに [上限緩和申請](https://console.aws.amazon.com/servicequotas/home) が必要です。また、**インスタンスを hold している間はインスタンスと EBS ボリュームへの課金が続きます** のでご注意ください。

ハイパーパラメタを変えながら 2, 3 回学習ジョブを動かしてから結果を確認します。

場合によっては以下のエラーが出るかもしれません。[こちら](https://aws.amazon.com/jp/premiumsupport/knowledge-center/sagemaker-python-throttlingexception/) で解決方法が案内されています。ただしこのエラーが出た際は、意図せず高頻度で API を叩きまくっている可能性をまず疑うのが良いでしょう。
```console
botocore.exceptions.ClientError: An error occurred (ThrottlingException) when calling the AddAssociation operation
```

## Experiment の確認

SageMaker Studio の左側にあるアイコンの一番上にあるホームボタンをクリックすると、Experiments というメニューが表示されるので、さらにクリックします。
![スクリーンショット 2022-12-18 17.56.39.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/56554e04-0c21-31e5-8e41-91d518ab8599.png)

### Run の確認

作成した Experiment をクリックして Run 一覧画面に行き、さらにどれかの Run をクリックします。以下の図はメトリクス画面です。
![スクリーンショット 2022-12-18 19.18.55.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/a1d359c3-7ec0-ef95-d3a7-bb7af9496c99.png)

左側のメニューにある Charts をクリックすると混同行列が表示されます。右上にあるAdd Chart ボタンをクリックするとグラフを追加できます。
![スクリーンショット 2022-12-18 19.20.07.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/45949803-cff7-3d28-8483-2edeb12484ad.png)

Chart type を Line にして、Axis dimension を Stepに、Y-axis を test:loss にして Create ボタンをクリックします。
![スクリーンショット 2022-12-18 19.20.50.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/712aaf7c-cf99-6d2a-be5b-a1fbd50f9d64.png)

チャートが追加されました！
![スクリーンショット 2022-12-18 19.21.55.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/c51d1e4a-adb0-0749-5211-a5c9800cdda7.png)

Run 一覧画面で複数の Run の右側のチェックボックスにチェックを入れて右上の Analyze ボタンをクリックすると、複数 Run の比較ができます。グラフは Add Chart ボタンから自分で作ります。
![スクリーンショット 2022-12-18 19.36.53.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/c7e617c2-cbc0-fa76-9e5d-81550cb9e277.png)

S3 パスなどの長い文字列の全体を表示するにはカラムの幅を広げる必要があるのですが、カラム幅の合計の最大が画面の幅（横スクロールバーがない）なので、全ての情報が見られる状態にできないし、画面の幅より長い文字列を最後まで確認する方法がないのがちょっといけてないです。マウスカーソルを文字列上に持っていってもフルテキストがホバー表示されることもないです。
![スクリーンショット 2022-12-18 19.29.50.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/3ba9c00b-a0ec-04a1-9432-665f225a7530.png)

`log_precision_recall`, `log_roc_curve`, `log_confusion_matrix` で自動的に作成されるグラフはアーティファクトと呼ぶようですが、アーティファクトの数の上限は 6000 のようです。それ以上登録しようとすると以下のエラーが出ます。無駄に量産したアーティファクトを削除したいのですが、今のところアーティファクトを削除する方法がわかりません。

```console
An error occurred (ResourceLimitExceeded) when calling the CreateArtifact operation: The account-level service limit 'Total number of artifacts allowed, excluding those automatically created by SageMaker' is 6000 Artifacts, with current utilization of 6000 Artifacts and a request delta of 1 Artifacts. Please contact AWS support to request an increase for this limit.
```

## おわりに

今回のアップデートで、Experiments がその名の通り実験管理により使いやすくなったと思います。これでモデル開発時の試行錯誤がはかどるというものです。

学習データのパスも学習済みモデルの場所も Experiments さんは知ってるはずなので、ここから [SageMaker Clarify](https://aws.amazon.com/jp/sagemaker/clarify/?sagemaker-data-wrangler-whats-new.sort-by=item.additionalFields.postDateTime&sagemaker-data-wrangler-whats-new.sort-order=desc) を使ってデータの品質チェックやモデル説明性の確認までできると良いなぁと妄想が膨らみました。

## ドキュメント

[Experiments の API リファレンス](https://sagemaker.readthedocs.io/en/stable/experiments/sagemaker.experiments.html)
[開発者ドキュメント](https://docs.aws.amazon.com/sagemaker/latest/dg/experiments.html)
[AWS ブログ（英語）](https://aws.amazon.com/jp/blogs/machine-learning/next-generation-amazon-sagemaker-experiments-organize-track-and-compare-your-machine-learning-trainings-at-scale/)