はじめに
Amazon SageMaker ユーザにはお馴染み（のはず）の Experiments ですが、これが 大幅にアップデート されました！
個人的に期待がつのるアップデートなので、早速どこが便利になったのかを確かめてみようと思います。
2022/12/20 追記：公式のサンプルコード が公開されました！
アップデートのポイント
まず一番大きいのが、記録した内容を表示するダッシュボード機能ができたことです！今までは Pandas の DataFrame としてノートブックで表示するくらいしかできませんでしたが、Amazon SageMaker Studio からこんなふうにみることができます。この例では、mnist-hand-written-digits-classification という Experiment の中に run1, run2... という Run が登録されています。今まで Trial といっていたものが Run になった感じかなと。（Experiment の中身を describe してみると Trial や TrialComponent として記録されているようです。）
Run をクリックすると詳細画面を確認できます。Charts のところを見るとこんな感じでグラフを確認することができます。この例では、train と test の loss は自分で手動で追加しましたが、混同行列は学習スクリプトの中で run.log_confusion_matrix という API を呼ぶだけで勝手に SageMaker が作ってくれました。
Run 一覧画面で複数の Run を選択して Analyze ボタンをクリックすると、以下のように Run を比較することができます。チャートは自分で追加しました。
どうやって使うのか
基本的には学習スクリプトの中に 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 を使う場合、これを忘れると以下のエラーが出ます。
ValueError: Must setup local AWS configuration with a region supported by SageMaker.
このエラーを回避するには、Estimator 作成時のパラメタ environment を以下のように設定を追加します。リージョン名は実際にお使いのものに書き換えてください。
environment={
"AWS_DEFAULT_REGION": "us-east-1"
},
Experiment の新規作成が不要に
これは注意というほどのものではないのですが、従来は Experiment.create API で Experiment を作成してから Trial を登録する流れでしたが、新しい Experiments では with Run(...) の部分で指定された名前の Experiment がない場合は自動的に作成してくれるようになりました。地味に便利！
やってみる
まだサンプルコードが公開されていないようなので、こちらのサンプルコード をベースに Experiments を試してみましょう。
2022/12/20 追記：公式のサンプルコード が公開されました！
SageMaker Studio の Python3 (PyTorch 1.12 Python 3.8 CPU Optimized カーネルと ml.t3.medium の組み合わせで実行しました。
ソースディレクトリを作成してファイルを移動
前述の通り 2.124.0 以降の SageMaker Python SDK が必要なのですが、今回使用するビルトインコンテナで使用している SDK のバージョンが古いので、それをアップグレードするために以下の requirements.txt を作成します。
sagemaker>=2.124.0
この requirements.txt を使うためには学習スクリプトと requirements.txt を同じディレクトリに入れておく必要があるので、その準備をします。root 直下に src というディレクトリを作成してその中に requirements.txt を配置し、root にある mnist.py も src 以下に移動します。最終的に、以下のようなディレクトリ構成になります。
root/
├── src/
| ├── mnist.py
| └── requirements.txt
└── pytorch_mnist.ipynb
mnist.py の変更
mnist.py に Experiments のコードを追加します。
モジュールを import しているあたりに以下のコードを追加します。
from sagemaker.experiments import load_run
次に、一番下にある引数を受け取る部分に以下のコードを追加します。ここでは、Experiment 名と Run 名を取得しています。
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 に記録するための関数です。
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 などで再度登録する必要はありません。
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 の値を記録しています。
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行追加します。
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 の値を登録します。
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 をインストール
次にノートブックを変更していきます。ノートブックの初めに以下のコードを記述したセルを追加して実行します。
!pip install --upgrade sagemaker
Estimator 作成部分の変更
「4.2.Estimatorの定義」と書かれた Estimator を作成するセルの上までセルを実行していきます。以下のコードが書かれたセルを追加および実行して、Run 名を作るのに使用するカウンタを初期化します。
count = 1
次に、Estimator を作成しているセルを、以下のように変更してから実行します。一番下にあるように、fit() を with Run() の中に入れるのを忘れないようにしてください。 fit() を入れ忘れても単体で Run の中身を確認することはできるのですが、複数の Run を比較することができません。
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 機能を使っています。Warm pool 機能は、一旦学習ジョブを作成したらその際に使用したインスタンスを停止せず hold しておき、次に同じ設定（ハイパーパラメタは例外）の Estimator から学習ジョブが作成されたらそのインスタンスを再度使えるというものです。これにより、keep_alive_period_in_seconds で指定した時間内の試行錯誤の際にインスタンスの起動時間を待つ必要がなくいため、すばやく試行錯誤のサイクルを回すことができます。なお、Warm pool を使用するには、インスタンスタイプごとに 上限緩和申請 が必要です。また、インスタンスを hold している間はインスタンスと EBS ボリュームへの課金が続きます のでご注意ください。
ハイパーパラメタを変えながら 2, 3 回学習ジョブを動かしてから結果を確認します。
場合によっては以下のエラーが出るかもしれません。こちら で解決方法が案内されています。ただしこのエラーが出た際は、意図せず高頻度で API を叩きまくっている可能性をまず疑うのが良いでしょう。
botocore.exceptions.ClientError: An error occurred (ThrottlingException) when calling the AddAssociation operation
Experiment の確認
SageMaker Studio の左側にあるアイコンの一番上にあるホームボタンをクリックすると、Experiments というメニューが表示されるので、さらにクリックします。
Run の確認
作成した Experiment をクリックして Run 一覧画面に行き、さらにどれかの Run をクリックします。以下の図はメトリクス画面です。
左側のメニューにある Charts をクリックすると混同行列が表示されます。右上にあるAdd Chart ボタンをクリックするとグラフを追加できます。
Chart type を Line にして、Axis dimension を Stepに、Y-axis を test:loss にして Create ボタンをクリックします。
チャートが追加されました！
Run 一覧画面で複数の Run の右側のチェックボックスにチェックを入れて右上の Analyze ボタンをクリックすると、複数 Run の比較ができます。グラフは Add Chart ボタンから自分で作ります。
S3 パスなどの長い文字列の全体を表示するにはカラムの幅を広げる必要があるのですが、カラム幅の合計の最大が画面の幅（横スクロールバーがない）なので、全ての情報が見られる状態にできないし、画面の幅より長い文字列を最後まで確認する方法がないのがちょっといけてないです。マウスカーソルを文字列上に持っていってもフルテキストがホバー表示されることもないです。
log_precision_recall, log_roc_curve, log_confusion_matrix で自動的に作成されるグラフはアーティファクトと呼ぶようですが、アーティファクトの数の上限は 6000 のようです。それ以上登録しようとすると以下のエラーが出ます。無駄に量産したアーティファクトを削除したいのですが、今のところアーティファクトを削除する方法がわかりません。
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 を使ってデータの品質チェックやモデル説明性の確認までできると良いなぁと妄想が膨らみました。