AWS Batch上で深層学習

同僚のKiuk ChungがAWS Batchを使って深層学習をするという素晴らしい記事を書いてくれました。

GPUインスタンスは当然のように深層学習とペアになりますが、それはそのニューラルネットワークのアルゴリズムがGPUインスタンスの超並列処理能力を活かすことができるからです。AWSではg2やp2といったGPUインスタンスを提供しており、お客様はスケーラブルなGPUワークロードを実行することができます。AWS Batchを使うことでそのスケーラビリティをもっと効率よく使うことができます。(訳注: 丁度GTC 2017のKeynoteにて次期NVIDIA GPUであるV100に関する情報も発表されましたのでご参考頂ければ幸いです: AWS and NVIDIA Expand Deep Learning Partnership at GTC 2017)

AWS Batchは皆さんの代わりに下回りの計算リソースを管理してくれるので、リソース管理のオーバーヘッド無しにモデリングすることに集中できます。AWS Batchにおける計算環境 (すなわちクラスタ)とは、皆さんのアカウント内のインスタンスのプールであり、AWS Batchはジョブの数に応じてインスタンスを起動したり削除したりしながらそれを動的にスケールしてくれます。これによって無駄なインスタンスを最小化でき、コストを最適化することができます。

さらに、AWS Batchは登録されたジョブが適切なインスタンスに配置されるように確実にスケジュールしてくれるので、ジョブのライフサイクルが管理されます。お客様独自のAMI利用の機能追加によって、AWS Batchの利用者はGPUが必要とされるジョブのためにもこの弾力性や利便性を活用することができるようになりました。

この記事ではGPUベースの深層学習ワークロードをAWS Batch上でどのように動かせばよいかをお見せします。Apache MXNetを使ってMNISTデータセットから手書きの数字を認識するための、畳み込みニューラルネットワーク(LeNetアーキテクチャ)の学習を例として使います。

MXNetのジョブをAWS Batchで実行する

Apache MXNetは機能が豊富で、柔軟にプログラムが書け、高いスケーラビリティをもった深層学習フレームワークで、畳み込みニューラルネットワーク (CNNs)やlong short-term memory networks (LSTMs)を含む最新の深層モデルをサポートしています。

AWS Batchでジョブを実行するには3つのステップがあります:

• カスタムAMIを作成
• AWS Batchのリソースを作成
• 学習ジョブを登録

カスタムAMIを作成

まず、NVIDIAドライバとAmazon ECSエージェントを含むAMIを作成するところから始めます。AWS Batchでは、計算環境を作成する時にimage_idを指定することで特定のAMIからインスタンスを起動させることができます。GPUが必要なジョブを実行しようとしているので、NVIDIAドライバが含まれたAMIが必要となります。

Launch Stackを選択して、あなたのアカウント上でus-east-1にCloudFromationテンプレートを起動します: 

下にある様に、CloudFormationスタックのOutputsタブの中にあるAMIという値をメモしておきます。次のセクションで計算環境を作成する時にこれをimage_idの値として使います。

または、AWS BatchのドキュメンテーションのGPU有効なAMIを作成するに従っても良いです。

AWS Batchのリソースを作成

AMIを作成したら、以下のリソースを作成します:

• 計算環境
• ジョブキュー
• ジョブ定義

計算環境は、同じまたは異なるインスタンスタイプのインスタンス(計算リソース)の集合です。ここでは、p2.xlargeのインスタンスを使ったマネージド計算環境を作成します。imageIdには、前のセクションで作成したAMIを指定します。

そうしたら、ジョブキューを作成します。AWS Batchでは、ジョブは計算環境の順序付きリストに紐付いたジョブキューに登録されます。順位が優先の計算環境が埋まってしまったら、ジョブは次の順位の計算環境へと漏れていきます。今回の例ではジョブキューには1つの計算環境のみ紐付けます。

最後に、ジョブの仕様のテンプレートとなるジョブ定義を作成します。Amazon ECSに馴染みのある方には、これはタスク定義と似たようなものですと言えばお分かりでしょう。ホスト上のNVIDIAドライバが含まれるディレクトリをコンテナの/usr/local/nvidia上にマウントします。またコンテナのプロパティでprivilegedフラグを設定する必要もあります。

以下のコードは前述のAWS Batchのリソースを作成してくれます。より詳しい情報は、AWS Batchユーザガイドをご覧ください。

git clone https://github.com/awslabs/aws-batch-helpers
cd aws-batch-helpers/gpu-example

python create-batch-entities.py\
 --subnets <subnet1,subnet2,…>\
 --security-groups <sg1,sg2,…>\
 --key-pair <ec2-key-pair>\
 --instance-role <instance-role>\
 --image-id <custom-AMI-image-id>\
 --service-role <service-role-arn> 

学習ジョブの登録

ここまできたら、手書き数字の認識のための畳み込みニューラルネットワークモデルを学習するジョブを登録します。Amazon ECSのタスクの様に、AWS BatchではジョブはDockerコンテナ内のコマンドとして実行されます。MXNetを深層学習ライブラリとして使うためには、MXNetが含まれたDockerイメージが必要になります。ここではmxnet/python:gpuを利用します。

submit-job.pyスクリプトはジョブを登録し、CloudWatch Logsから出力をtailしてくれます。

# cd aws-batch-helpers/gpu-example
python submit-job.py --wait

下のような出力を見ることができると思います:

Submitted job [train_imagenet - e1bccebc-76d9-4cd1-885b-667ef93eb1f5] to the job queue [gpu_queue]
Job [train_imagenet - e1bccebc-76d9-4cd1-885b-667ef93eb1f5] is RUNNING.
Output [train_imagenet/e1bccebc-76d9-4cd1-885b-667ef93eb1f5/12030dd3-0734-42bf-a3d1-d99118b401eb]:
 ================================================================================

[2017-04-25T19:02:57.076Z] INFO:root:Epoch[0] Batch [100]	Speed: 15554.63 samples/sec Train-accuracy=0.861077
[2017-04-25T19:02:57.428Z] INFO:root:Epoch[0] Batch [200]	Speed: 18224.89 samples/sec Train-accuracy=0.954688
[2017-04-25T19:02:57.755Z] INFO:root:Epoch[0] Batch [300]	Speed: 19551.42 samples/sec Train-accuracy=0.965313
[2017-04-25T19:02:58.080Z] INFO:root:Epoch[0] Batch [400]	Speed: 19697.65 samples/sec Train-accuracy=0.969531
[2017-04-25T19:02:58.405Z] INFO:root:Epoch[0] Batch [500]	Speed: 19705.82 samples/sec Train-accuracy=0.968281
[2017-04-25T19:02:58.734Z] INFO:root:Epoch[0] Batch [600]	Speed: 19486.54 samples/sec Train-accuracy=0.971719
[2017-04-25T19:02:59.058Z] INFO:root:Epoch[0] Batch [700]	Speed: 19735.59 samples/sec Train-accuracy=0.973281
[2017-04-25T19:02:59.384Z] INFO:root:Epoch[0] Batch [800]	Speed: 19631.17 samples/sec Train-accuracy=0.976562
[2017-04-25T19:02:59.713Z] INFO:root:Epoch[0] Batch [900]	Speed: 19490.74 samples/sec Train-accuracy=0.979062
[2017-04-25T19:02:59.834Z] INFO:root:Epoch[0] Train-accuracy=0.976774
[2017-04-25T19:02:59.834Z] INFO:root:Epoch[0] Time cost=3.190
[2017-04-25T19:02:59.850Z] INFO:root:Saved checkpoint to "/mnt/model/mnist-0001.params"
[2017-04-25T19:03:00.079Z] INFO:root:Epoch[0] Validation-accuracy=0.969148

================================================================================
Job [train_imagenet - e1bccebc-76d9-4cd1-885b-667ef93eb1f5] SUCCEEDED

実際には、学習されたモデル情報をAmazon S3に保存することで後続の予測ジョブが学習したモデルを使って予測を生成できるように、ジョブのコマンドを修正したくなると思います。Amazon S3のオブジェクトをジョブからどのように参照するかについては、Creating a Simple “Fetch & Runb” AWS Batch Jobの記事をご覧ください。

まとめ

この記事では、MXNetを深層学習のライブラリとして使いながらAWS Batch上でGPU有効なジョブを実行する例をお見せしました。AWS Batchはご自身のワークロードのために最も最適なアルゴリズムを実装することに集中できるだけの基礎的な部分を提供してくれます。登録したジョブのライフサイクルを管理し、指定した範囲内でジョブの要求に合わせて動的にインフラを最適化することが可能です。コスト最適なやり方でAWSが提供する計算用インスタンスの水平スケーラビリティの利点を簡単に活かすことができます。

MXNetは、一方で、自身の深層学習アルゴリズムを実装し始めるために必要な、高度に最適化されスケーラブルなビルディングブロックを豊富に提供しています。これによって巨大なニューラルネットワークモデルを必要とする問題を解決するだけでなく、Amazon EC2の無制限の計算リソースをつなげることで繰り返しの時間を削減することもできます。

AWS Batchが皆さんの代わりにリソースを管理してくれるので、ハイパーパラメータの最適化のために数十、数百の検索を並列で実行して、問題に最適なモデルパラータを探すといったワークロードの実装を簡単に行うことができます。さらに、ジョブはDockerコンテナの中で実行されるので、必要に応じて最適なツールとライブラリを選択し、Dockerイメージをビルドし、好きなイメージを使ってジョブを登録することができます。

ご自身の手でこれを試してみることをおすすめします。そして感想をぜひ教えて下さい！

原文: Deep Learning on AWS Batch (翻訳: SA岩永)