Tailored Pruning via Rollback Learning

Abstract

Structured pruning of networks has received increased attention recently for compressing convolutional neural networks (CNNs). Existing methods are not yet optimal, however, because there is still an obvious performance loss after pruning. This conflicts with the learning theory’s parsimony principle that prefers the simpler model, when they explain the same facts, and thus maintain a consistent level of performance. In this paper, we propose rollback learning for pruning (RLP)-an algorithm to efficiently prune networks while maintaining performance. RLP provides a generic and tailored pruning framework that is easily applied to existing CNNs. Unlike existing work, we conditionally prune the kernels, and only those reaching the performance limit (high filter sparsity) are ultimately removed. This guarantees an optimized pruning by backtracking and updating the dense filters to their full potential. Our implementation adds rollback pruning into generative adversarial learning to compute a soft mask to scale the output of the network structures by defining a new objective function. The soft mask is backtracked and duplicated before and after the batch normalization layer, leading to an optimized fit to the structure pruning. Extensive experiments demonstrate that the proposed method significantly outperforms state-of-the-art methods resulting in a $\mathrm{{0.25}} \times$ FLOPs reduction and $0.04\%$ accuracy increase over the full-precision ResNet-18 model on CIFAR10, and an $\mathrm{{0.18}} \times$ FLOPs reduction and $0.54\%$ accuracy increase over the full-precision ResNet-50 model on ImageNet.