Scheduling Large-scale Distributed Training via Reinforcement Learning

Abstract

© 2018 IEEE. Scheduling the training procedure of deep neural networks (DNNs) such as tuning the learning rates is crucial to the successes of deep learning. Previous strategies such as piecewise and exponential learning rate schedulers have different arguments (hyper-parameters) that need to be tuned manually. With the expanding of data scale and model computation, searching for these arguments spends lots of empirical efforts. To address this issue, this work proposes policy schedular that determines the arguments of learning rate (lr) by reinforcement learning, significantly reducing costs to tune them. The policy schedular possesses several appealing benefits. First, instead of manually defining the values of initial lr and ultimate lr, it autonomously determines these values in training. Second, rather than using predefined functions to update lr, it adaptively oscillates lr by monitoring learning curves without human intervention. Third, it is able to select lr for each block or layer of a DNN. Experiments show that the DNNs trained with policy schedular achieve superior performances, outperforming previous work on various tasks and benchmarks such as ImageNet, COCO, and learning-to-learn.