Training EfficientNets at Supercomputer Scale: 83% ImageNet Top-1 Accuracy in One Hour

Arissa Wongpanich,Hieu Pham,James Demmel,Mingxing Tan,Quoc Le,Yang You,Sameer Kumar

DOI: https://doi.org/10.48550/arXiv.2011.00071

2020-11-05

Abstract:EfficientNets are a family of state-of-the-art image classification models based on efficiently scaled convolutional neural networks. Currently, EfficientNets can take on the order of days to train; for example, training an EfficientNet-B0 model takes 23 hours on a Cloud TPU v2-8 node. In this paper, we explore techniques to scale up the training of EfficientNets on TPU-v3 Pods with 2048 cores, motivated by speedups that can be achieved when training at such scales. We discuss optimizations required to scale training to a batch size of 65536 on 1024 TPU-v3 cores, such as selecting large batch optimizers and learning rate schedules as well as utilizing distributed evaluation and batch normalization techniques. Additionally, we present timing and performance benchmarks for EfficientNet models trained on the ImageNet dataset in order to analyze the behavior of EfficientNets at scale. With our optimizations, we are able to train EfficientNet on ImageNet to an accuracy of 83% in 1 hour and 4 minutes.

Machine Learning,Computer Vision and Pattern Recognition,Distributed, Parallel, and Cluster Computing

What problem does this paper attempt to address?

The main problem that this paper attempts to solve is to accelerate the training time of the EfficientNets model on large - scale supercomputers while maintaining the accuracy of the model. Specifically, the authors explored techniques for scaling EfficientNets training on TPU - v3 Pods, with the goal of significantly reducing the training time without affecting the model performance. For example, through the application of optimization techniques, they were able to achieve 83% ImageNet Top - 1 accuracy within 1 hour, while typically a model of this scale requires several days to train.

To achieve this goal, the paper explored a variety of optimization techniques, including but not limited to:
- **Large - scale batch optimizer**: Use an optimizer suitable for large - scale training, such as LARS (Layer - wise Adaptive Rate Scaling), to overcome the problem of model quality degradation as the batch size increases.
- **Learning rate scheduling**: Adopt learning rate warm - up and linear scaling techniques to meet the training requirements under large - scale batch processing.
- **Distributed evaluation**: Solve the computational bottleneck in the standard cloud TPU implementation by distributing the training and evaluation steps across all TPUs.
- **Distributed batch normalization**: Improve the final model accuracy by distributing batch normalization across replicas, although this will increase the communication cost between TPUs.
- **Precision**: Implement mixed - precision training, utilize the bfloat16 format to improve hardware efficiency while maintaining model quality.

These techniques work together to enable EfficientNets to be trained efficiently on a supercomputer scale, significantly increasing the training speed while also maintaining high accuracy of the model.