Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour

By Priya Goyal, Piotr Doll'ar, Ross B. Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, Kaiming He. 2017

4,148 citations recorded by the source corpus.

Deep learning thrives with large neural networks and large datasets. However, larger networks and larger datasets result in longer training times that impede research and development progress. Distributed synchronous SGD offers a potential solution to this problem by dividing SGD minibatches over a pool of parallel workers. Yet to make this scheme efficient, the per-worker workload must be large, which implies nontrivial growth in the SGD minibatch size. In this paper, we empirically show that on the ImageNet dataset large minibatches cause optimization difficulties, but when these are addressed the trained networks exhibit good generalization. Specifically, we show no loss of accuracy when training with large minibatch sizes up to 8192 images. To achieve this result, we adopt a hyper-parameter-free linear scaling rule for adjusting learning rates as a function of minibatch size and develop a new warmup scheme that overcomes optimization challenges early in training. With these simple techniques, our Caffe2-based system trains ResNet-50 with a minibatch size of 8192 on 256 GPUs in one hour, while matching small minibatch accuracy. Using commodity hardware, our implementation achieves ~90% scaling efficiency when moving from 8 to 256 GPUs. Our findings enable training visual recognition models on internet-scale data with high efficiency.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Optimization & Training

Citations in the corpus

Masked Autoencoders Are Scalable Vision Learners (2021, methodology): “We use the linear lr scaling rule [20]: lr = base lr × batchsize / 256.”
A Simple Framework for Contrastive Learning of Visual Representations (2020, methodology): “In contrast to supervised learning (Goyal et al., 2017), in contrastive learning, larger batch sizes provide more negative examples, facilitating convergence (i.e. taking fewer epochs and steps for a given accuracy).”
Momentum Contrast for Unsupervised Visual Representation Learning (2019, background): “But optimizing with a larger mini-batch is harder [25], and it is questionable whether the trend can be extrapolated into a larger K even if memory is sufficient.”
Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers (2020, methodology): “In particular, what if we use gradient accumulation to increase the batch size rather than the model size? We vary the batch size for the 12 layer, 768H model and increase the learning rate as is common practice (Goyal et al., 2017; Liu et al., 2019b).”
An Empirical Model of Large-Batch Training (2018, methodology): “For example, batch sizes of 8 thousand [GDG + 17], 16 thousand [SKYL17], 32 thousand [YGG17, YZH + 17, ASF17], and even 64 thousand [JSH + 18] examples have been effectively employed to train ImageNet, and batch sizes of thousands have been effective for language models and generative models…”
Rethinking ImageNet Pre-Training (2018, background): “Specially, the initial learning rate is 0.02 (with a linear warm-up [12]).”
Measuring the Effects of Data Parallelism on Neural Network Training (2018, methodology): “For a fair comparison between batch sizes, batch normalization statistics should be computed over the same number of examples or else the training objective diﬀers between batch sizes (Goyal et al., 2017).”
Do Better ImageNet Models Transfer Better? (2018, methodology): “Second, we retrained one network (Inception v3) with less regularization in a setting similar to ResNets (see Appendix A.4).”
Exploring the Limits of Weakly Supervised Pretraining (2018, methodology): “All other initialization and training details match [19] and are summarized in the supplemental material.”
Adafactor: Adaptive Learning Rates with Sublinear Memory Cost (2018, background): “The step size α t is often held constant over the course of training, but recent work in large-scale optimization suggests that performance can be improved on some problems through a linear ramp-up followed by some form of decay (Goyal et al., 2017; Vaswani et al., 2017).”
Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model (2019, methodology): “…the batch size initially achieves perfect scaling (i.e. doubling the batch size halves the number of steps needed) up to a problem-dependent critical batch size, beyond which it yields diminishing returns [Balles et al., 2017, Goyal et al., 2017, Jastrzębski et al., 2018, McCandlish et al., 2018].”
Beyond human-level accuracy: computational challenges in deep learning (2019, background): “Researchers have shown how to train image and language applications quickly using data-parallel scaling to reduce per-accelerator activation memory requirements [14, 26, 40].”
Don't Decay the Learning Rate, Increase the Batch Size (2017, methodology): “Goyal et al. (2017) had already observed such a scaling rule, and exploited it to train ResNet-50 on ImageNet in one hour.”
Sustainable AI: Environmental Implications, Challenges and Opportunities (2021, background): “In addition, significant research has gone into algorithmic approaches to efficiently scale training (Goyal et al., 2017; Ott et al., 2018) by reducing communication cost via compression (Alistarh et al.”
Llemma: An Open Language Model For Mathematics (2023, methodology): “We use Tensor Parallelism (Shoeybi et al., 2019) with a world size of 2 for L LEMMA -7B , and a world size of 8 for L LEMMA -34B, alongside ZeRO Stage 1 sharded optimizer states (Rajbhandari et al., 2020) across Data Parallel (Goyal et al., 2017) replicas.”
Object-Centric Learning with Slot Attention (2020, methodology): “[77] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He.”
Exploring Self-Attention for Image Recognition (2020, background): “Momentum and weight decay are set to 0.9 and 1e-4, respectively [12, 38, 11].”