🌅

SageMaker Notebook 上で PyTorch モデルの学習を回す

2022/09/18に公開

AWS

いままでEC2でゴリゴリ学習させていたものをSageMaker Notebookに移植することになり、それなりに苦労したので自分用のテンプレート作成も兼ねてハマりどころをまとめます。

今回作成した雛型

そもそも SageMaker Notebook とは

SageMaker Notebook自体はAWSのサーバー上で動くJupyter Notebookです。ただし、ここから学習時にTraining Jobというものを作成して、それを別のインスタンス上のコンテナで実行させることができます。基本的には ml.t3.medium あたりの通常のインスタンス上でNotebookを動かし、Training Jobを p3 系なので強力なインスタンスで実行することが多いと思います。

概要

SageMaker Notebookでの学習手順は次のようになります。

Training Jobの作成（Training用コンテナの作成）。
コンテナにS3から学習に必要なデータをダウンロード。
コンテナで学習を実行。
学習結果をS3に保存。
コンテナの終了。

手動での操作が必要なのは、上記の1番 Estimator.fit() のみで、それ以降の処理は自動で実行されます。

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    entry_point="entry.py",
    source_dir="src",
    role=role,
    py_version="py38",
    framework_version="1.11.0",
    instance_count=1,
    instance_type="ml.p3.8xlarge",
    output_path=f"s3://{bucket}/{prefix}",
    hyperparameters={
        "worker_params": worker_params
    },
)

上記はPyTorchの例ですが、Tensorflowなども同様にEstimatorを作成し、その際にコンテナを作成する学習用プログラムのパスやインスタンスタイプなどを設定します。

estimator.fit({'train': 's3://sagemaker-ap-northeast-1-xxxxxxxxxxxx/mydata'})

その後、データセットのS3 URIなどを渡してTraining Jobを作成します。あとは学習が完了するのを待つだけです。学習完了後は自動でコンテナが終了するので、EC2のようにインスタンスの停止忘れの心配などがありません。

既存プログラムからの移植

基本的にはこれまで学習用に作成していたプログラムの80-90%ぐらいはそのまま流用できます。綺麗に実装していれば既存プログラムにはほぼ手を入れずに移植できると思います。残りはSageMakerから既存スクリプトを呼び出すための玄関部分と、データ入出力用のS3とのやり取りスクリプトさえ用意してやれば移植は完了です。

今回用意したサンプルプログラムだと以下のような感じです。

src/ 以下のプログラムはほぼそのまま以前のものと一緒。
main.ipynb は新たに作成したSageMaker Notebook用のプログラム。

ハマりどころ

謎の環境変数群

SageMakerから呼び出されるentry_pointのスクリプトには、モデルの保存場所や入力データの場所など各種パラメタを渡すことになります。

その際に謎の環境変数が登場するので、公式ドキュメントを頼りに一つずつ確認するしかありません。ネット上には古い情報が転がっているので、一次情報に当たらずに盲信すると動かないこともあります（一敗）

ライフサイクル設定

事前に設定しておくと Notebook インスタンスの作成・起動時に自動でスクリプトを実行することができます。例えば

idle 状態が 1 時間続くと自動でインスタンスを停止
インスタンス起動時に追加のライブラリを自動でインストール

といったことが可能です。

※ただし、Windowsユーザーはライフサイクルスクリプトの設定時に改行コードに注意しないとエラーを吐かれてインスタンスが起動しません （二敗）

/bin/bash^M: bad interpreter: No such file or directory

参考: https://github.com/aws-samples/amazon-sagemaker-notebook-instance-lifecycle-config-samples/issues/8

換装した感想

ネット中の大海から情報をかき集めようとせず始めから公式のexampleを読もう！

実際に使ってみた感想ですが、インスタンスの停止忘れが避けられるのと、スポットインスタンスが気軽に使えるのは非常にありがたいですね。費用面でもかなりEC2時代から抑えられている印象です。