Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs

Abstract

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.

arXiv commentAccepted to ICLR 2020 https://openreview.net/forum?id=rkxDoJBYPB

References in corpus (9)

• Neural Message Passing for Quantum Chemistry
• On the difficulty of training Recurrent Neural Networks
• MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
• Learning Combinatorial Optimization Algorithms over Graphs
• Combinatorial Optimization with Graph Convolutional Networks and Guided Tree Search
• Device Placement Optimization with Reinforcement Learning
• Learning to Optimize Tensor Programs
• Beyond Data and Model Parallelism for Deep Neural Networks
• Learning to superoptimize programs

Cited by in corpus (7)

• A Learned Performance Model for Tensor Processing Units
• A View on Deep Reinforcement Learning in System Optimization
• Efficient Algorithms for Device Placement of DNN Graph Operators
• ConfuciuX: Autonomous Hardware Resource Assignment for DNN Accelerators using Reinforcement Learning
• DLFusion: An Auto-Tuning Compiler for Layer Fusion on Deep Neural Network Accelerator
• Resource Allocation in Disaggregated Data Centre Systems with Reinforcement Learning
• Scaling Up Graph Neural Networks Via Graph Coarsening