Image Classification at Supercomputer Scale

TL;DR

This paper presents system-level optimizations enabling the training of deep learning models at petaflop scale, achieving rapid training times and high throughput on large-scale TPU clusters.

Contribution

The paper introduces three novel systems optimizations—distributed batch normalization, input pipeline enhancements, and 2-D torus all-reduce—for large-scale deep learning training.

Findings

Trained ResNet-50 on ImageNet to 76.3% accuracy in 2.2 minutes.

Achieved over 1.05 million images/second throughput.

No accuracy drop despite large-scale distributed training.

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.