富士通が画像認識ニューラルネットワークの学習を74.7秒で完了：産総研のスーパーコンピュータABCIで達成

課題

画像認識をはじめとしたニューラルネットワークは、ネットワークの構造やデータセットの数など規模が大きくなるにつれて性能が向上することが明らかになっていますが、学習に時間がかかるなど応用上の問題も残っています。

ResNet-50と呼ばれる代表的な画像認識モデルに関してはこの2年間だけでも30倍以上の高速化が図られており大規模ニューラルネットワークの学習高速化のベンチマークとして使用されています。この学習速度向上の競争は業界全体の大規模クラスタのハードウェア・ソフトウェア両面の技術向上に寄与しています。

これまで、認識性能をおおよそ同等とした上で、Sonyが224秒、Googleが132秒での学習完了を達成し、さらなる向上が期待されていました。

解決方法

2018年4月1日、富士通研究所はResNet-50の学習において新記録を達成したと発表しました。

このプロジェクトでは、Sonyが記録を達成したのと同じ産業技術総合研究所のABCI（AI橋渡しクラウド）の2,048個のNVIDIA Tesla V100 GPU（マザーボードへはSXM2で接続）を並列に用いて計算を行いました。高速化のために、確率的勾配降下法（Stochastic Gradient Decent）の学習率を学習が不安定にならないように適応的に増加させる技術やネットワークの層ごとに異なる学習率となるようにする手法（Layer-wise Adaptive Rate Scaling: LARS）を取り入れるなど学習手法への工夫を行いました。

また、学習データのまとまりであるバッチサイズを81,920と増加させても学習結果の劣化を抑えることに成功しました。一般にバッチサイズは大きすぎると性能の低下を起こしますが、並列実行する大規模ネットワークの場合にはバッチサイズが大きい方が計算が効率的となります。この大きさは、同じシステムを用いたSonyの実験55,296に比較して約1.5倍となっています。

ニューラルネットワークを構築するフレームワークとしてMXNetを使用しました。このフレームワークはC++とCUDA Cにより構築されており柔軟性とスケーラビリティに優れているとされています。研究グループはMXNetのボトルネックもプロファイラにより見つけ出して修正を施したとしています。

学習時の勾配の通信においても最適化を行い、計算だけでなく通信での時間の削減も行なっています。

どうなったか

このプログラムによりResNet-50の学習が74.7秒という短時間で終了することに成功しました。これはSonyがABCIにおいて1088ノードに4台ずつ計4352個のNVIDIA Tesla V100 GPUを用いて達成した224秒や、その後Googleが1024個のTPU v3（Googleが開発したニューラルネットワーク専用アクセラレータ）を用いて達成した132秒（少し識別精度の劣る設定で108秒）を凌駕する成果になっています。

また識別精度に関しても75.08%となり、Sonyの75.29%やGoogleの76.3%（132秒）、75.2%（108秒）と見劣りしない成果となっています。

Sonyと同じABCIという同じスーパーコンピュータで達成されたのは、独自の工夫が功を奏したといえます。

まとめ

ResNet-50の速度向上は実用上区別のつかないところまで到達したのではないでしょうか？

今後は、さらに大規模なネットワークや時間のかかる対象、例えばBERTなどの大規模な自然言語処理などに移行するのではないでしょうか？オリジナル論文によると大きなタイプのBERTでは合計64個のTPUを用いて4日間の学習が必要でした。BERTでは、今の所ネットワークの規模が大きくなればなるほど性能が向上するとされるため、さらなる大規模化と学習の効率化が求められるでしょう。