ノートブックインスタンスのトレーニングに Amazon SageMaker のローカルモードを使用する

Amazon SageMaker が最近、事前構築型の TensorFlow と MXNet コンテナを使用したローカルトレーニングのサポートを開始しました。  Amazon SageMaker はフレキシブルな機械学習プラットフォームで、機械学習モデルを本番環境に、より効果的に構築、トレーニング、およびデプロイできるようにします。  Amazon SageMaker トレーニング環境は管理されています。つまり、即時のスピンアップ、アルゴリズムコンテナのロード、Amazon S3 からのデータ取得、コードの実行、Amazon S3 への結果の出力、クラスターの解体などの処理について、ユーザーが考える必要がなくなります。  トレーニングを別々のマルチノード GPU クラスターにオフロードできる機能は大きな利点です。毎回新しいハードウェアをスピンアップすることは、反復性とセキュリティの面で良いことですが、アルゴリズムコードのテストやデバッグの際には無駄に回転させることになります。

Amazon SageMaker の深層学習コンテナを使用することで、普段通りに TensorFlow または MXNet のスクリプトを書けるようになります。しかし、今後は、これらをトレーニングとホスティングの両方の目的で、管理されている本番グレードの環境にプレビルトコンテナとしてデプロイすることになります。  これまで、こうしたコンテナは Amazon SageMaker 固有の環境でのみ使用できました。  最近これらのコンテナがオープンソース化されました。これにより、ユーザーは自身のアルゴリズムをローカルでテストするためにわずか 1 行のコードを書き換えるだけで、自分の作業環境にコンテナを取り込み、Amazon SageMaker Python SDK でカスタムコードビルトを使用できるようになるのです。  つまり、毎回新しいトレーニング用またはホスティング用のクラスターが構築されるのを待つことなく、自分の作業を反復したり、テストしたりできることを意味します。  小さなデータセットのサンプルをローカルで反復し、フルデータセットを分散してトレーニングするためにスケーリングするのが機械学習の一般的な手法です。  多くの場合、プロセス全体を再度書き込むことになるため、バグが生じないのを願うばかりです。  Amazon SageMaker のローカルモードでは、コードを 1 行書き換えるだけで、ローカルと分散されたマネージドトレーニング間をシームレスに切り替えられます。その他の動作はすべて同じです。

Amazon SageMaker Python SDK のローカルモードは、TensorFlow または MXNet エスティメーターの 1 つの引数を変えることで CPU (シングルおよびマルチインスタンス) と GPU (シングルインスタンス) SageMaker トレーニングジョブのエミュレートを実行できます。  これを実行するために、Docker 構成と NVIDIA Docker を使用します。  また、Amazon ECS から Amazon SageMaker TensorFlow または MXNet コンテナを引き出すため、自身のローカル環境からパブリックの Amazon ECR リポジトリにアクセスできる必要があります。  ローカル環境として SageMaker ノートブックインスタンスを使用すると、このスクリプトは必要な要件をインストールします。  そうでない場合、ユーザー自身がそれらをインストールし、次を使って SageMaker Python SDK の最新バージョンにアップグレードされるようにすることも可能です:  pip install -U sagemaker。

ユースケースの例

TensorFlow と MXNet のそれぞれに、MNIST でローカルモードを使用する方法を示したサンプルノートブックを 2 種類用意しました。  しかし、Amazon SageMaker でもまた、より大きなノートブックインスタンスタイプが利用できることになったと発表されたばかりなので、4 GPU で 5 万色のカラーイメージのイメージセパレータをトレーニングする、より大きなサンプルを試してみましょう。  CIFAR-10 データセットで MXNet Gluon を使用して ResNet ネットワークをトレーニングし、1 つの ml.p3.8xlarge ノートブックインスタンスで全体を実行します。  もし、複数のマシンでこれを実行する、または、ハードウェアを管理することなく、反復ベースでトレーニングを行う場合、ここで構築する同じコードを簡単に Amazon SageMaker マネージドトレーニング環境へ移行できます。

新しいノートブックインスタンスを作成するところから始めましょう。AWS マネジメントコンソールにログインし、Amazon SageMaker サービスを選択します。そして、Amazon SageMaker コンソールで [Create notebook instance] を選択し、次のページを開きます。

ノートブックのインスタンスが開始したら、新しい Jupyter ノートブックを作成し、セットアップを開始します。  または、ここで事前定義したノートブックを使用することでもできます。  ローカルモードに集中するため、いくつかの背景情報は飛ばします。  深層学習が初めての方には、このブログ記事が入門として役立つかもしれません。  または他の SageMaker サンプルノートブックをご覧ください。このサンプルは、CIFAR-10 の ResNet に該当しますが、Amazon SageMaker マネージドトレーニング環境でトレーニングを受けます。

前提条件をインストールし、ライブラリーがロードされ、データセットがダウンロードされたら、データセットを Amazon S3 バケットへロードします。  私たちはローカルでトレーニングを実行していますが、SageMaker トレーニングとの整合性を保つため、Amazon S3 のデータにアクセスする点に注意してください。

inputs = sagemaker_session.upload_data(path='data', key_prefix='data/DEMO-gluon-cifar10')

次に、MXNet エスティメーターを定義します。  エスティメーターは cifar10.py スクリプトをポイントします。このスクリプトには、ネットワークの仕様と train()関数が含まれます。  また、ハイパーパラメーターや IAM ロールなどのジョブに関する情報を提供します。  しかし、最も重要な点としては train_instance_type を  'local_gpu' に設定する点です。  SageMaker マネージドトレーニング環境から、ローカルノートブックインスタンスにトレーニング全体を切り替えるのに必要な手順はこれだけです。

m = MXNet('cifar10.py',
          role=role,
          train_instance_count=1,
          train_instance_type='local_gpu',
          hyperparameters={'batch_size': 1024,
                           'epochs': 50,
                           'learning_rate': 0.1,
                           'momentum': 0.9})

m.fit(inputs)

初めてエスティメーターが稼動するときは、Amazon ECR リポジトリからコンテナイメージをダウンロードする必要がありますが、トレーニングはすぐに開始できます。  別のトレーニングクラスターがプロビジョンされるまで待機する必要はありません。  また、反復またはテスト時に再実行が必要になる可能性があります。MXNet または TensorFlow スクリプトへの変更点は、瞬時に開始します。

また、ローカルモードでトレーニングすることにより、私たちの使用しているハードウェアを活用するため、コードが適切に書き込まれるよう GPU の消費量などのメトリクスを簡単に監視できます。  この場合、 nvidia-smiを、ResNet モデルのごく簡単なトレーニングのため、ml.p3.8xlarge の GPU のすべてを使用するターミナルで実行することで、確認できます。

エスティメーターのトレーニングが終了したあとは、エンドポイントをローカルで作成しテストできます。  再び 'local_gpu' を instance_type に設定します。

predictor = m.deploy(initial_instance_count=1, instance_type='local_gpu')

これで推論コードが機能することを確認するために、いくつかの予測を立てることができます。  これは本番環境のエンドポイントにデプロイする前に行うのが望ましいですが、1 回限りのモデル精度の評価のために、いくつかの予測を立てることもできます。

from cifar10_utils import read_images

filenames = ['images/airplane1.png',
             'images/automobile1.png',
             'images/bird1.png',
             'images/cat1.png',
             'images/deer1.png',
             'images/dog1.png',
             'images/frog1.png',
             'images/horse1.png',
             'images/ship1.png',
             'images/truck1.png']

image_data = read_images(filenames)

for i, img in enumerate(image_data):
    response = predictor.predict(img)
    print('image {}: class: {}'.format(i, int(response)))

結果は次のようになるはずです。

image 0: class: 0
image 1: class: 9
image 2: class: 2
image 3: class: 3
image 4: class: 4
image 5: class: 5
image 6: class: 6
image 7: class: 7
image 8: class: 8
image 9: class: 9

これで、トレーニングと推論スクリプトを評価できたので、大きなスケールのトレーニング、または反復ベースで SageMaker マネージド環境にデプロイし、リアルタイムのホストエンドポイントから予測を立てることもできます。

ただし、最初にローカルエンドポイントをクリーンアップしましょう。これはローカルで 1 度に実行できるエンドポイントは 1 つのみであるためです。

m.delete_endpoint()

You can shut down your notebook instance from the [Notebook] ページにナビゲートし、[Stop] を選択することで、Amazon SageMaker コンソールからノートブックインスタンスをシャットダウンできます。  このようにすることで、ユーザーがバックアップを開始するまで、コンピューティングの料金が発生するのを防ぐことができます。  または、[Actions] と [Delete] を選択することで、ノートブックインスタンスを削除することもできます。

まとめ

このブログ記事では、大規模なイメージ分類モデルを素早くテストするために、最近リリースされたマルチ GPU ノートブックインスタンスで Amazon SageMaker Python SDK ローカルモードを使用する方法を紹介しています。  今日からテストとデバッグサイクルを加速するために、ローカルモードのトレーニングを使用できます。  SageMaker Python SDK の最新バージョンがお手元にあることを確認し、他のツールをインストールして、コードを 1 行変更するだけです。

今回のブログ投稿者について

David Arpin は AWS の AI Platforms Selection Leader で、データサイエンスチームと製品管理を統率してきた経歴があります。