Jump to Content

Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs

Aditya Paliwal
Felix Gimeno
Vinod Gopal Nair
Yujia Li
Miles Lubin
International Conference on Learning Representations (ICLR) (2020)


We present a deep reinforcement learning approach to minimizing the execution cost of neural network computation graphs in an optimizing compiler. Unlike earlier learning-based works that require training the optimizer on the same graph to be optimized, we propose a learning approach that trains an optimizer offline and then generalizes to previously unseen graphs without further training. This allows our approach to produce high-quality execution decisions on real-world TensorFlow graphs in seconds instead of hours. We consider two optimization tasks for computation graphs: minimizing running time and peak memory usage. In comparison to an extensive set of baselines, our approach achieves significant improvements over classical and other learning-based methods on these two tasks.