SonyがImageNet上でResNet-50を224秒でトレーニング

Sonyの研究者グループは、WordNet階層に従って構成された画像データベースであるImageNet上で、ResNet 50アーキテクチャのマシンラーニングを、わずか224秒でトレーニングしたと発表した。作成されたネットワークは、ImageNetの検証セット上で75パーセントのtop-1精度を持つ。研究者らはこの記録を、NVIDIAの2.100 Tensor Core GPUを使用することで達成した。新たな時間記録と合わせて、1.088 Tesla V100 Tensor Core GPUを使用して、90パーセントのGPUスケーリング効率も達成している。

ニューラルネットワークのトレーニングを複数マシンでスケーリングするためには、おもに2つの課題がある – トレーニングを行うバッチのサイズと、勾配(gradient)の同期だ。

ニューラルネットワークのトレーニング時間を短縮するためには、ネットワークの重みを“正しい方向”に短時間でプッシュするために、小さなミニバッチ(mini-batch)サイズを使用する方法がある。ウェイトを更新するためには、エラーによってウェイトが更新される“方向”を示す勾配を与える。ミニバッチを小さくすることで、この方向を迅速かつ頻繁に決めることができるのだ。しかしながら、ミニバッチが小さいと、ニューラルネットワークはパフォーマンスの局所的な極小点(local minumum)に向かう傾向がある。Sonyはこの問題を、最近になって使用されることの多くなったテクニックである、バッチサイズコントロールを使って解決した。

バッチサイズコントロールにより、損失状況(loss landscape)が極小点を回避するように、ミニバッチサイズが徐々に増加する。より多くの画像を見ることによって、ニューラルネットワークは、少数のサンプルに基づいてエラー報告を決定するのみの場合よりも、改善に必要な平均エラー報告に関するよりよいアイデアを得ることができるようになる。Sonyは最初の5つのエポックをウォームアップエポックとして低いバッチサイズで処理し、移行のエポックにはより大きなバッチサイズを使用した。さらに、FP16でフォワード/バックワード計算が行なわれる場面では、混合精度(mixed-precision)トレーニングも採用している。

第2の問題は、マシン間の勾配の同期によって遅くなる可能性があることだ。異なるGPU間の通信には、多くの帯域幅が必要となる。Sonyの研究陣は、通信のオーバヘッドを削減する手段として、2D-Torus Allreduceアルゴリズムを採用した。この通信アルゴリズムでは、GPUは仮想グリッドに配置される。最初に勾配を水平方向に通知し、次にすべての列に対して垂直方向に通信し、最後のパスで再び水平方向に通信する。これはつまり、水平方向のGPU数がXとした場合、2(X-1)のGPU-to-GPUオペレーションが必要であるということだ。

研究者グループは、DNNトレーニングフレームワークとしてNeural Network Libraries(NLL)とそのCUDA拡張を使用し、GPU間通信にはNVIDIA Collective Communications Library (NCCL) バージョン2を使用した。

昨年、複数のグループが、ResNet 50アーキテクチャを最短時間でトレーニングする試みを行っている。2017年9月にInfoQは、IBMが同じニューラルネットワークアーキテクチャを50分でrトレーニングしたことをお伝えしている。当時IBMは、より高いスケーリング効率を達成していたが、256GPUのみだった。Q&AでHillery Hunter氏は、バッチサイズは最も難しい部分のひとつではあるが、同社のアプローチがより多くのGPUにも拡張可能であると予測していた。今回の論文では、さまざまなGPU数におけるGPUスケーリング効率が調査されており、3,264GPUを使用したトレーニングでの効率は52.74パーセントだった。

Sonyの研究陣は、研究結果を‘ImageNet/ResNet-50 Training in 224 Seconds’という論文にまとめており、arXivで見ることができる。