A Generalist Reinforcement Learning Agent for Compressing Convolutional Neural Networks

Abstract

Over the years, researchers have proposed multiple approaches to reduce the number of parameters Deep Learning models have. Due to the complexity of compressing models, some authors have opted to train Reinforcement Learning agents that learn how to compress a particular model without losing considerable accuracy. Nonetheless, training an agent for each model can be time-consuming. We propose a methodology for training a generalist agent capable of compressing other convolutional neural networks that it was not trained to compress. Our generalist agent uses feature maps to select which compression technique to apply to convolutional and dense layers. Since the shape of the feature maps is reduced as it goes deeper into the network, we implemented a Dueling Deep Q-Network with a Region of Interest layer, allowing it to generate features of a fixed size for feature maps of various heights and widths. Our generalist agent trained to compress two LeNet models, one trained with fashion MNIST and the other with Kuzushiji-MNIST, compressed the same architecture trained on MNIST to less than 15% of its original size with an accuracy loss of less than 2.5%.

Keywords

• Computer vision
• deep learning
• model compression
• model optimization
• reinforcement learning