Amazon SageMaker を使用した分散型 TensorFlow トレーニングの実行

TensorFlow は、複数のホストで複数の GPU を使用した分散型トレーニングを必要とするヘビーウェイトなディープニューラルネットワーク (DNN) の開発に広く使用されるオープンソースの機械学習 (ML) ライブラリです。Amazon SageMaker は、アクティブラーニングを使用したデータのラベル付けから、モデルの分散型トレーニング、トレーニングの進捗状況のモニタリング、訓練されたモデルの自動スケーリング可能な RESTful サービスとしてのデプロイメント、並列的な ML 実験の一元管理までの ML ワークフローをシンプル化するマネージドサービスです。

この記事では、Amazon SageMaker を使用した分散型 TensorFlow トレーニングに注目していきます。

コンセプトの概要

この記事における分散型トレーニングコンセプトの多くは、さまざまなタイプの TensorFlow モデル全体に広く該当するものですが、この記事ではCommon Object in Context (COCO) 2017 データセットでの Mask R-CNN モデルの分散型 TensorFlow トレーニングに焦点を当てます。

モデル

Mask R-CNN モデルはオブジェクトインスタンスセグメンテーションのために使用されるものであることから、このモデルは画像内の各オブジェクトインスタンスを区分するために、ピクセルレベルのマスク (シグモイドバイナリ分類) と、オブジェクトカテゴリでアノテーションが付けられた (SoftMax 分類) 境界ボックス (Smooth L1 回帰) を生成します。Mask R-CNN の一般的なユースケースには、自動運転車における認知、表面欠陥検出、および地理空間画像の分析などがあります。

この記事に Mask R-CNN モデルを選択した主な理由は 3 つあります。

大型のデータセットにおける Mask R-CNN の分散データ並列トレーニングは、トレーニングパイプラインでの画像のスループットを向上させ、トレーニング時間を短縮する。
Mask R-CNN モデルには多数のオープンソース TensorFlow 実装が利用できる。この記事では、その主な例として Tensorpack Mask/Faster-RCNN 実装を使用しますが、高度に最適化された AWS Samples Mask-RCNN も推奨されます。
Mask R-CNN モデルは、MLPerf 結果の一部のヘビーウェイトなオブジェクト検出モデルとして送信されます。

以下の画像は、Mask R-CNN ディープニューラルネットワークアーキテクチャの概略図です。

分散型トレーニングにおける勾配の同期型 AllReduce

分散型 DNN トレーニングにおける中核的な課題は、複数のノード全体の複数の GPU でモデルウェイトを更新するために勾配を適用する前に、同期化されたステップで、複数の GPU 間でのバックプロパゲーション中に計算された勾配に AllReduce (平均化) が実行されなければならないことです。

同期型 AllReduce アルゴリズムは効率性に優れている必要があります。そうでなければ、同期型 AllReduce ステップの非効率性のために分散データ並列トレーニングから得たトレーニング迅速化が損なわれることになります。

同期型 AllReduce アルゴリズムの効率を高めるには 3 つの主な課題があります。

アルゴリズムを、分散型トレーニングクラスター内で増加するノードと GPU の数と共にスケールする必要がある。
アルゴリズムが、単一ノード内の高速な GPU 対 GPU 相互接続のトポロジを利用する必要がある。
アルゴリズムが、他の GPU との通信を効率的にバッチ化することによって、他の GPU との通信と GPU での計算を効率的にインターリーブする必要がある。

Uber のオープンソースライブラリである Horovod は、次のようにこれら 3 つの課題に対処しています。

Horovod は、増加する GPU とノードの数と共にスケールする効率性に優れた同期型 AllReduce アルゴリズムの選択肢を提供します。
Horovod ライブラリは、Nvidia GPU トポロジに対するアウェアネスを利用した Nvidia Collective Communications Library (NCCL) の通信プリミティブを使用します。
Horovod には Tensor Fusion が含まれており、これは AllReduce のためのデータ通信をバッチ化することによって、通信と計算を効率的にインターリーブします。

Horovod は、TensorFlow を含む多くの ML フレームワークでサポートされています。TensorFlow の分散戦略も NCCL を使用しており、分散型 TensorFlow トレーニングを実行するための Horovod の使用に代わる手段を提供します。この記事では Horovod を使用します。

Mask R-CNN などのヘビーウェイト DNN のトレーニングでは、トレーニングパイプラインに 1 つ、または複数の高解像度画像を送り込むことができるように、GPU ごとに多くのメモリが必要になります。また、勾配の同期型 AllReduce を効率的に実行できるように、高速な GPU 対 GPU 相互接続と高速なネットワーキング相互接続マシンも必要になります。Amazon SageMaker の ml.p3.16xlarge と ml.p3dn.24xlarge の各インスタンスタイプは、これらすべての要件を満たします。詳細については、Amazon SageMaker ML インスタンスタイプを参照してください。8 台の Nvidia Tesla V100 GPU、128～256 GB の GPU メモリ、25～100 Gbps のネットワーキング相互接続、および高速な Nvidia NVLink GPU 対 GPU 相互接続を持つこれらのインスタンスは、Amazon SageMaker での分散型 TensorFlow トレーニングに最適です。

メッセージパッシングインターフェイス

分散型 TensorFlow トレーニングの次の課題は、複数ノード間におけるトレーニングアルゴリズムプロセスの適切な配置と、各アルゴリズムプロセスの一意のグローバルランクとの関連付けです。メッセージパッシングインターフェイス (MPI) は、並列コンピューティングのために幅広く使用されている集団通信プロトコルで、複数ノード全体でのトレーニングアルゴリズムのワーカープロセスグループの管理に役立ちます。

MPI は、トレーニングアルゴリズムプロセスを複数のノードに分散し、各アルゴリズムプロセスを一意のグローバルランクおよびローカルランクに関連付けるために使用されます。Horovod は、所定のノードにあるアルゴリズムプロセスを特定の GPU にピン留めするために使用されます。各アルゴリズムプロセスの特定 GPU への論理的なピン留めは、勾配の同期型 AllReduce に必要です。

この記事のために理解しておく MPI の重要なコンセプトは、MPI がマスターノードで mpirun コマンドを使用して複数ノードで並行プロセスを起動することです。MPI を使用すると、マスターホストが複数ノードで実行されている分散型トレーニングプロセスのライフサイクルを一元的に管理します。Amazon SageMaker を使用した分散型トレーニングの実行に MPI を使用するには、MPI を Amazon SageMaker のネイティブな分散型トレーニング機能と統合する必要があります。

MPI の Amazon SageMaker 分散型トレーニングとの統合

MPI を Amazon SageMaker の分散型トレーニングに統合する方法を理解するには、以下のコンセプトを理解する必要があります。

Amazon SageMaker は、Docker イメージ にパッケージ化されたトレーニングアルゴリズムとフレームワークを必要とします。
Docker イメージは Amazon SageMaker トレーニング用に有効化されなければなりません。この有効化は、Amazon SageMaker 対応の Docker イメージの作成に役立つライブラリである Amazon SageMaker コンテナを使用することで簡素化されます。
Amazon SageMaker トレーニングイメージに、Amazon SageMaker とアルゴリズムコード間の仲介役となるエントリポイントスクリプト (通常は Python スクリプト) を提供する必要があります。
所定のホストでトレーニングを開始するため、Amazon SageMaker がトレーニングイメージから Docker コンテナを実行し、ハイパーパラメータや入力データの場所などの情報を提供するエントリポイント環境変数でエントリポイントスクリプトを呼び出します。
エントリポイントスクリプトは、正しい args でアルゴリズムプログラムを起動するためにエントリポイント環境変数で渡された情報を使用し、実行中のアルゴリズムプロセスをポーリングします。
アルゴリズムプロセスが終了すると、エントリポイントスクリプトがアルゴリズムプロセスの終了コードを伴って終了します。Amazon SageMaker は、この終了コードを使用してトレーニングジョブが成功したか失敗したかを判断します。
エントリポイントスクリプトは、アルゴリズムプロセスの stdout と stderr の出力を独自の stdout にリダイレクトします。次に、Amazon SageMaker はエントリポイントスクリプトからの stdout をキャプチャして、それを Amazon CloudWatch Logs に送信します。Amazon SageMaker はトレーニングジョブで定義されたアルゴリズムメトリクスのために stdout 出力を解析し、メトリクスを Amazon CloudWatch メトリクスに送信します。
Amazon SageMaker が複数のトレーニングインスタンスをリクエストするトレーニングジョブを開始すると、一連のホストを作成し、各ホストを algo-k (k はホストのグローバルランク) として論理的に命名します。例えば、トレーニングジョブが 4 つのトレーニングインスタンスをリクエストする場合、Amazon SageMaker はホストを algo-1、algo-2、algo-3、および algo-4 と命名します。ホストは、これらのホスト名を使用してネットワークに接続できます。

MPI を使用した分散型トレーニングの場合は、複数のノード (algo-1 から algo-n のノードで、n は Amazon SageMaker トレーニングジョブでリクエストされたトレーニングインスタンスの数) 全体に分散されたすべてのアルゴリズムプロセスのライフサイクルを制御するマスターノード (ホスト) で実行される単一の mpirun コマンドが必要です。ただし、Amazon SageMaker は MPI、または複数ノードにアルゴリズムプロセスを分散させるために使用するその他の並列処理フレームワークを認識しません。Amazon SageMaker は各ノードで実行されている Docker コンテナでエントリポイントスクリプトを呼び出します。これは、エントリポイントスクリプトがそのノードのグローバルランクを認識し、それがマスターノード、または非マスターノードの 1 つで呼び出されるかどうかに応じて異なるロジックを実行する必要があることを意味します。

具体的には、MPI の場合、マスターノードで呼び出されたエントリポイントスクリプトが mpirun コマンドを実行して、現行の Amazon SageMaker トレーニングジョブのホストセット内のすべてのノードでアルゴリズムプロセスを開始する必要があります。非マスターノードのいずれかで同じエントリポイントスクリプトが Amazon SageMaker によって呼び出されると、スクリプトがその非マスターノード (mpirun コマンドがマスターノードからリモートで管理するもの) でアルゴリズムプロセスが引き続き実行されているかどうかを定期的にチェックし、実行されていなければ終了します。

MPI のマスターノードは論理的なコンセプトなので、現行のトレーニングジョブのホストセットにあるすべてのホストの中からマスターノードとして設定するホストを指定するのはエントリポイントスクリプトになります。この指定は、分散的な方法で行われる必要があります。algo-1 をマスターノードとして指定し、他のすべてのノードを非マスターノードとするのがシンプルなアプローチです。Amazon SageMaker はエントリポイント環境変数で各ノードにその論理ホスト名を提供するため、ノードがマスターノードか非マスターノードかを判断するのは簡単です。

付随する GitHub リポジトリに含まれ、Tensorpack Mask/Faster-RCNN アルゴリズムの Docker イメージにパッケージ化されている train.py は、このセクションで説明したロジックに従います。

これらのコンセプトを理解したところで、Amazon SageMaker を使用した Mask R-CNN のための分散型 TensorFlow トレーニングを実行する方法におけるステップバイステップのチュートリアルに進む準備が整いました。

ソリューションの概要

このチュートリアルには、次の主なステップがあります。

AWS CloudFormation のオートメーションスクリプトを使用してプライベート Amazon VPC を作成し、このプライベート VPC にアタッチされた Amazon SageMaker ノートブックインスタンスネットワークを作成する。
Amazon SageMaker ノートブックインスタンスから、プライベート VPC にアタッチされた Amazon SageMaker 管理の Amazon VPC ネットワークで分散型トレーニングジョブを起動する。トレーニングデータパイプラインのデータソースとしては、Amazon S3、Amazon EFS、および Amazon FSx を使用できます。

前提条件

以下は、必要となる前提条件です。

AWS アカウントを作成してアクティブ化する、または既存の AWS アカウントを使用する。
Amazon SageMaker インスタンス制限を管理する。最小で 2 つの ml.p3dn.24xlarge または 2 つの ml.p3.16xlargeインスタンスが必要です。サービスの制限をそれぞれ 4 にすることが推奨されます。サービスの制限は各 AWS リージョンに固有であることに注意してください。この記事では us-west-2 を使用します。
この記事の GitHub リポジトリをクローンして、記事にある手順を完了します。記事内のすべてのパスは、GitHub リポジトリのルートに相対します。
Amazon SageMaker、EFS、および Amazon FSx をサポートする任意の AWS リージョンを使用してください。この記事では us-west-2 を使用します。
新しい S3 バケットを作成、または既存のバケットを選択します。

VPC にアタッチされた Amazon SageMaker ノートブックインスタンスの作成

AWS CloudFormation のオートメーションスクリプトを実行するための最初のステップは、プライベート VPC にアタッチされた Amazon SageMaker ノートブックインスタンスの作成です。このスクリプトを実行するには、ネットワーク管理者機能と一致する IAM ユーザーのアクセス許可が必要です。そのようなアクセス権がない場合は、ネットワーク管理者に助けを求めて、このチュートリアルに含まれる AWS CloudFormation のオートメーションスクリプトを実行してもらう必要があるかもしれません。詳細については、職務機能の AWS 管理ポリシーを参照してください。

AWS CloudFormation テンプレート cfn-sm.yaml を使用して、プライベート VPC にアタッチされた Amazon SageMaker ノートブックインスタンスを作成する AWS CloudFormation スタックを作成します。これには、AWS CloudFormation サービスコンソールで cfn-sm.yaml を使用して AWS CloudFormation スタックを作成するか、stack-sm.sh スクリプトの変数をカスタマイズして、AWS CLI がインストールされている場所でスクリプトを実行することができます。

AWS CLI アプローチを使用するには、以下の手順を完了してください。

AWS CLI をインストールして設定します。
stack-sm.sh で AWS_REGION をお使いの AWS リージョン、S3_バケット をお使いの S3 バケットに設定します。これら 2 つの変数は必須です。
オプションとして、既存の EFS ファイルシステムを使用したい場合は EFS_ID 変数を設定します。EFS_ID を空のままにしておくと、新しい EFS ファイルシステムが作成されます。既存の EFS ファイルシステムの使用を選択する場合は、既存のファイルシステムに既存のマウントターゲットがないことを確認してください。詳細については、Amazon EFS ファイルシステムの管理を参照してください。
オプションとして、GIT_URL を指定して GitHub リポジトリを Amazon SageMaker ノートブックインスタンスに追加します。GitHub リポジトリがプライベートの場合は、GIT_USER および GIT_TOKEN 変数を指定できます。
AWS CLI を使用してカスタマイズされた stack-sm.sh スクリプトを実行し、AWS CloudFormation スタックを作成します。

後ほど使用するために、AWS CloudFormation スクリプトのサマリー出力を保存します。この出力は、AWS マネジメントコンソールにある AWS CloudFormation スタックの [出力] タブで表示することもできます。

Amazon SageMaker トレーニングジョブの起動

Amazon SageMaker コンソールで、作成したノートブックインスタンスを開きます。このノートブックインスタンスには、Mask R-CNN のトレーニングのために利用できる 3 つの Jupyter ノートブックがあります。

S3 バケットをデータソースとして使用する Mask R-CNN ノートブック、mask-rcnn-s3.ipynb。
EFS ファイルシステムをデータソースとして使用する Mask R-CNN ノートブック、mask-rcnn-efs.ipynb。
Amazon FSx Lustre ファイルシステムをデータソースとして使用する Mask R-CNN ノートブック、mask-rcnn-fsx.ipynb。

3 つすべてのデータソースオプションのトレーニング時間におけるパフォーマンスは、この記事で選択した Mask R-CNN モデルおよび COCO 2017 データセットとほとんど違いません (が、同一ではありません)。コストプロファイルはデータソースごとに異なります。以下は、トレーニングデータパイプラインをセットアップするためにかかる時間における違いです。

S3 データソースでは、トレーニングジョブが起動されるたびに、S3 バケットから各トレーニングインスタンスにアタッチされたストレージボリュームに COCO 2017 データセットをレプリケートするための時間が約 20 分かかります。
EFS データソースでは、S3 バケットから EFS ファイルシステムに COCO 2017 データセットをコピーするために約 46 分かかります。このデータは一度コピーするだけで済みます。トレーニング中、データはネットワークインターフェイスを通じてすべてのトレーニングインスタンスにマウントされた共有 EFS ファイルシステムから入力されます。
Amazon FSx では、新しい Amazon FSx Lustre ファイルシステムを作成して、COCO 2017 データセットを S3 バケットから新しい Amazon FSx Lustre ファイルシステムにインポートするために約 10 分かかります。これを実行するのは一度だけです。トレーニング中、データはネットワークインターフェイスを通じてすべてのトレーニングインスタンスにマウントされた共有 Amazon FSx Lustre ファイルシステムから入力されます。

どのデータソースが最適なのかがよくわからない場合は、S3 から始めて、各トレーニングジョブの開始時におけるトレーニングデータのダウウンロード時間が許容範囲外である場合に EFS または Amazon FSx を検討してください。どのデータソースについても、トレーニング時間のパフォーマンスについて臆測しないようにしてください。トレーニング時間のパフォーマンスは多くの要因に応じて変化するため、実験して測定することが最も望ましい方法です。

これら 3 つのケースのすべてにおいて、トレーニング中のログとモデルチェックポイントの出力が各トレーニングインスタンスにアタッチされたストレージボリュームに書き込まれ、トレーニングが終了すると S3 バケットにアップロードされます。ログは、トレーニングの進捗状況として Amazon CloudWatch にもフィードされ、これらはトレーニング中に確認することが可能です。システムおよびアルゴリズムトレーニングのメトリクスは、トレーニング中に Amazon CloudWatch メトリクスにフィードされ、これらは Amazon SageMaker サービスコンソールで視覚化できます。

トレーニング結果

次のグラフは、COCO 2017 データセットでの 24 個のエポックのトレーニング後における 2 つのアルゴリズムの結果例です。

以下では、TensorPack Mask/Faster-RCNN アルゴリズムの結果例を確認することができます。以下のグラフは 3 つのバケットに分けることができます。

Intersection over Union (IoU) のさまざまな値、および小、中、大のオブジェクトサイズに対する境界ボックス (bbox) 予測の Mean Average Precision (mAP) グラフ
Intersection over Union (IoU) のさまざまな値、および小、中、大のオブジェクトサイズに対するオブジェクトインスタンスセグメンテーション (segm) 予測の Mean Average Precision (mAP) グラフ
トレーニングロスまたはラベルの正確性に関連するその他のメトリクス

以下では、最適化された AWS Samples Mask R-CNN アルゴリズムの結果例を確認することができます。以下のグラフに示されるコンバージされた mAP メトリクスは、前のアルゴリズムとほぼ同一ですが、コンバージェンスの進行が異なります。

まとめ

Amazon SageMaker は、Docker ベースのシンプル化された分散型 TensorFlow トレーニングプラットフォームを提供します。これによって、ユーザーがその ML アルゴリズムに集中し、インフラストラクチャの可用性とスケーラビリティの仕組みや並列実験の管理といった付随的な懸念に気を取られることがなくなります。モデルが訓練されたら、Amazon SageMaker の統合されたモデルデプロイメント機能を使用して、モデルのために自動スケーリングが可能な RESTful サービスエンドポイントを作成し、モデルのテストを開始できます。詳細については、Amazon SageMaker ホスティングサービスでモデルをデプロイするを参照してください。モデルの準備が整ったら、モデルの RESTful サービスを本番環境にシームレスにデプロイできます。

著者について

Ajay Vohra は、自動運転車開発のための認知機械学習を専門とするプリンシパルソリューションアーキテクトです。Amazon に入社する前、Ajay は金融リスクのモデル化のための超並列グリッドコンピューティング分野、およびオンプレミスデータセンターでのアプリケーションプラットフォームエンジニアリングの自動化に携わっていました。

Amazon Web Services ブログ