Neural Network Pruning Through Constrained Reinforcement Learning

Network pruning reduces the size of neural networks by removing neurons such that the performance drop is minimal. Traditional pruning approaches focus on designing metrics to quantify the usefulness of a neuron which is often tedious and sub-optimal. More recent methods have instead focused on training auxiliary networks to automatically learn how useful each neuron is however, they often do not take computational limitations into account. We propose a general methodology for pruning neural networks. Our approach can prune neural networks to respect pre-defined computational budgets on arbitrary, possibly non-differentiable, functions. Furthermore, we only assume the ability to be able to evaluate these functions for different inputs, and hence they do not need to be fully specified beforehand. We achieve this by proposing a novel pruning strategy via constrained reinforcement learning algorithms. We prove the effectiveness of our approach via comparison with state-of-the-art methods on standard image classification datasets. Specifically, we reduce 83−92.90% of total parameters on various variants of VGG while achieving comparable or better performance. We also achieved 75.09% reduction in parameters on ResNet18 without incurring any loss in accuracy. The code is also available at our github repository1.