Device Placement Optimization for Deep Neural Networks via One-shot Model and Reinforcement Learning

With the development of deep learning that employs deep neural networks (DNN) as powerful tool, its computational requirement grows rapid together with the increasing size (e.g. depth and parameter) of DNN. Currently, model and data parallelism are employed for accelerating the training and inference process of DNN. However, the above techniques make placement decision on devices for DNN based on heuristics and intuitions by machine learning experts. In this paper, we propose an novel approach for designing device placement of DNN in an automatic way. For a DNN, we employ a sequence-to-sequence model as controller to sample device placement from a one-shot model, which contains all possible device placements with respect to a specific hardware environment (e.g. CPU and GPU). Then, reinforcement learning treats the execution time of sampled device placement as reward to guide the sequence-to-sequence model for finding better one. The proposed approach is employed to optimize the device placement for both model and data parallelism of Inception-V3 on ImageNet. Experimental results show that the optimal placements discovered by our method can outperform hand-crafted one.