[論文レビュー] Image Classification at Supercomputer Scale

Chris Ying, Sameer Kumar+3 more|arXiv (Cornell University)|Nov 16, 2018

Advanced Neural Network Applications参考文献 1被引用数 95

ひとこと要約

tldr: 本論文は、画像ネットの大規模で ResNet-50 を訓練するためのシステム最適化を提示し、1024チップ TPU v3 Pod 上で 2.2 分で 76.3% の精度を達成し、スループットは 1.05 百万画像/秒を超える。

ABSTRACT

Deep learning is extremely computationally intensive, and hardware vendors have responded by building faster accelerators in large clusters. Training deep learning models at petaFLOPS scale requires overcoming both algorithmic and systems software challenges. In this paper, we discuss three systems-related optimizations: (1) distributed batch normalization to control per-replica batch sizes, (2) input pipeline optimizations to sustain model throughput, and (3) 2-D torus all-reduce to speed up gradient summation. We combine these optimizations to train ResNet-50 on ImageNet to 76.3% accuracy in 2.2 minutes on a 1024-chip TPU v3 Pod with a training throughput of over 1.05 million images/second and no accuracy drop.

研究の動機と目的

• Petascale での深層ネットワーク訓練の動機付けと、精度を保ちながら wall-clock 時間を短縮すること。
• Accelerator 上での大規模バッチ同期 SGD を妨げるシステムのボトルネックを特定すること。
• グローバルバッチサイズとリプリカごとのバッチサイズのバランスを、モデル品質を損なわずに設計・検証すること。
• 最先端のスループットを達成しつつ精度を維持する組み合わせ最適化を提示すること。

提案手法

• 畳み込みには bf16、非畳み込み演算には 32-bit の混合精度訓練を使用する。
• ウォームアップとデカイを伴う線形学習率スケーリングを適用し、LARS を用いてバッチサイズを最大で 32768 へスケールする。
• グローバルなバッチサイズとは独立して BN 統計を制御する分散バッチ正規化を導入する。
• データセットのシャーディング、キャッシュ、プリフェッチ、融合 JPEG デコードとクロップ、並列解析を組み合わせた入力データパイプラインを最適化する。
• 勾配和の 2-D トーラス全減法を採用して通信遅延を低減する。
• 1024-chip TPU v3 Pod での性能を実証し、従来の結果と比較する。

Figure 1: Cloud TPU v2 device with four chips, 180 teraFLOPS of peak floating point throughput and 64 GB of High Bandwidth Memory (HBM).

実験結果

リサーチクエスチョン

• RQ1 分散 BN、入力パイプラインの最適化、および 2-D 全減法を組み合わせて、巨大スケールでの同期 SGD を可能にできるか。
• RQ2 大規模バッチ訓練におけるリプリカごとのバッチサイズ、グローバルバッチサイズ、モデル精度のトレードオフはどうなるか。
• RQ3 混合精度とスケーリング戦略は、非常に大きなバッチで ImageNet の ResNet-50 の精度を維持できるか。
• RQ4 大規模な TPU ポッド上で、精度を損なうことなく実行時間とスループットのエンドツーエンドの成果はどれくらいか。

主な発見

Hardware	Chips	Batch	Optimizer	BN	Accuracy	Time
Goyal et al. [6]	P100	256	8192	Momentum	Local	76.3%	1 hour
Smith et al. [16]	TPU v2	128	8192 → 16384	Momentum	Local	76.1%	30 mins.
Akiba et al. [2]	P100	1024	32768	RMS + Mom.	Local	74.9%	15 mins.
Jia et al. [10]	P40	1024	65536	LARS	Local	76.2%	8.7 mins.
Baseline TPU v2	4	1024	Momentum	Local	76.3%	7.2 hours
Ours TPU v2	256	16384	Momentum	Local	75.3%	9.7 mins.
Ours TPU v2	256	32768	LARS	Local	76.3%	8.0 mins.
Ours TPU v3	512	32768	LARS	Local	76.4%	3.3 mins.
Ours TPU v3	1024	65536	LARS	Local	75.2%	1.8 mins.
Ours TPU v3	1024	32768	LARS	Distributed	76.3%	2.2 mins.

• 大規模なスケールでも、ImageNet の ResNet-50 を 76.3% の精度で訓練し、精度の低下なし。
• 1024-chip TPU v3 Pod で 2.2 分の訓練時間と 1.05 百万画像/秒超のスループットを達成。
• 2-D 勾配和はトーラスリンクを用いることで、レイテンシとスループットの両方で 1-D リング全減法を上回り、スケーラブルな同期を実現。
• 分散 BN はグローバルバッチサイズに依存せず BN の実効バッチサイズを制御でき、巨大規模での精度を支援。
• 入力パイプラインの最適化（キャッシュ、プリフェッチ、融合 JPEG デコード/クロッピング、並列解析）は、特にデータがワーカー間でシャーディングされる場合にスループットを大幅に向上。

Figure 2: Liquid-cooled Cloud TPU v3 device with four chips, 420 teraFLOPS of peak floating point throughput and 128 GB of HBM.