Storage Efficient and Dynamic Flexible Runtime Channel Pruning via Deep Reinforcement Learning

1 INTRODUCTION . In recent years , convolutional neural networks ( CNNs ) have been proven to be effective in a wide range of computer vision tasks , such as image classification ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2015 ; He et al. , 2016 ) , objection detection ( He et al. , 2017 ; Zhou et al. , 2019 ; Law & Deng , 2018 ) , segmentation ( He et al. , 2017 ; Zhu et al. , 2019 ) . Therefore , nowadays , many computervision-based systems , such as automatic-driving cars , security surveillance cameras , and robotics , are built on the power of CNNs . However , since most state-of-the-art CNNs require expensive computation power for inference and huge storage space to store large amount of parameters , the limitation of energy , computation and storage on mobile or edge devices has become the major bottleneck on real-world deployments of CNNs . Existing studies have been focused on speeding up the execution of CNNs for inference on edge devices by model compression using matrix decomposition ( Denil et al. , 2013 ; Masana et al. , 2017 ) , network quantization ( Courbariaux et al. , 2016 ) , network pruning ( Dong et al. , 2017 ) , etc . Among these approaches , channel pruning has shown promising performance ( He et al. , 2017 ; Luo et al. , 2017 ; Zhuang et al. , 2018 ; Peng et al. , 2019 ) . Specifically , channel pruning discards an entire input or output channel and keep the rest of the model with structures . Most channel pruning approaches can be categorized into two types : runtime approaches and static approaches . Static channel pruning approaches aim to design a measurement to evaluate the importance of each channel over the whole training dataset and remove the least important channels to minimize the loss of performance after pruning . By permanently pruning a number of channels , the computation and storage cost of CNNs can be dramatically reduced when being deployed , and the inference execution can be accelerated consequently . Runtime channel pruning approaches have been recently proposed to achieve dynamic channel pruning on each specific instance ( Gao et al. , 2019 ; Luo & Wu , 2018 ) . To be specific , the goal of runtime approaches aims to evaluate the channel importance at runtime , which is assumed to be different on different input instances . By pruning channels dynamically , different pruned structures can be considered as different routing of data stream inside CNNs . This kind of approaches is able to significantly improve the representation capability of CNNs , and thus achieve better performance in terms of prediction accuracy compared with static approaches . However , previous runtime approaches trade storage cost off dynamic flex- ibility of pruning . To achieve dynamic pruning on different specific instances , all parameters of kernels are required to be stored ( or even more parameters are introduced ) . This makes runtime approaches not applicable on resource-limited edge devices . Moreover , most of previous runtime approaches only evaluate the importance among channels in each single layer independently , without considering the difference in efficiency among layers . In this paper , to address the aforementioned issues of runtime channel pruning approaches , we propose a deep reinforcement learning ( DRL ) based pruning framework . Basically , we aim to apply DRL to prune CNNs by maximizing received rewards , which are designed to satisfy the overall budget constraints along side with network ’ s training accuracy . Note that automatic channel pruning by DRL is a difficult task because the action space is usually very huge . Specifically , the discrete action space for the DRL agent is as large as the number of channels at each layer , and the action spaces may vary among layers since there are different numbers of channels in different layers . To facilitate pruning CNNs by DRL , for each layer , we first design a novel prediction component to estimate the importance of channels , and then develop a DRL-based component to learn the sparsity ratio of the layer , i.e. , how many channels should be pruned . More specifically , different from previous runtime channel pruning approaches , which only learn runtime importance of each channel , we propose to learn both runtime importance and additionally static importance for each channel . While runtime importance maintains the saliency of specific channels for each given specific input , the static importance captures the overall saliency of the corresponding channel among the whole dataset . According to each type of the channel importance , we further design different DRL agents ( i.e. , a runtime agent and a static agent ) to learn a sparsity ratio in a layer-wise manner . The sparsity ratio learned by the runtime agent together with the estimated runtime importance of channels are used to generate runtime pruning structures , while the sparsity ratio learned by the static agent together with the estimated static importance of channels are used to generate static ( permanent ) pruning structures . By considering both the pruning structures , our framework is able to provide a trade-off between storage efficiency and dynamic flexibility for runtime channel pruning . In summary , our contributions are 2-fold . First , we propose to prune channels by taking both runtime and static information of the environment into consideration . Runtime information endows pruning with flexibility based on different input instances while static information reduces the number of parameters in deployment , leading to storage reduction , which can not be achieved by conventional runtime pruning approaches . Second , we propose to use DRL to determine sparsity ratios , which is different from the previous pruning approaches which manually set sparsity ratios . Extensive experiments demonstrate the effectiveness of our method . 2 RELATED WORK AND PRELIMINARY . 2.1 STRUCTURE PRUNING . Wen et al . ( 2016 ) pioneered structure pruning in deep neural network by imposing the L2,1 norm in training . Under the same framework , Liu et al . ( 2017 ) regarded parameters in batch normalization as channel selection signal , which is minimized to achieve pruning during training . He et al . ( 2017 ) formulated channel pruning into a two-step iterative process including LASSO regression based channel selection and least square reconstruction . Luo et al . ( 2017 ) formulated channel pruning as minimization of difference of output features , which is solved by greedy selection . Zhuang et al . ( 2018 ) further considered early prediction , reconstruction loss and final loss to select importance channels . Overall , structure pruning methods accelerate inference by producing regular and compact model . However , this brought regularness requires preserving more parameters to ensure performance . 2.2 DYNAMIC PRUNING . Dynamic pruning provides different pruning strategies according to input data . Wang et al . ( 2018 ) proposed to reduce computation by skipping layers or channels based on the analysis of input features . Luo & Wu ( 2018 ) proposed to use layer input to learn channel importance , which is then binarized for pruning . Gao et al . ( 2019 ) applied the same framework while extended features selection in both input and output features . Similarly , Liu & Deng ( 2018 ) introduced multiple branches for runtime inference according to inputs . A gating module is learnt to guide the flow of feature maps . Bolukbasi et al . ( 2017 ) learned to choose the components of a deep network to be evaluated for each input adaptively . Early exit is introduced to accelerate computation . Dynamic pruning adaptively takes different actions for different inputs , which is able to accelerate the overall inference time . However , the original high-precision model needs to be stored , together with extra parameters for making specified pruning actions . Rosenbaum et al . ( 2018 ) proposed to learn routers to route layers output to different next layers , in order to adjust a network to multi-task learning . 2.3 DEEP REINFORCEMENT LEARNING IN PRUNING . Channel selection is on trial using deep reinforcement learning . Lin et al . ( 2017 ) trained a LSTM model to remember and provide channel pruning strategy for backbone CNN model , which is conducted using reinforcement learning techniques . He et al . ( 2018 ) proposed to determine the compression ratio in each layer by training an agent regarding the pruning-retraining process as an environment . 2.4 PRELIMINARY . Reinforcement Learning We consider a standard setup of reinforcement learning : an agent sequentially takes actions over a sequence of time steps in an environment , in order to maximize the cumulative reward ( Sutton & Barto , 1998 ) . This problem can be formulated as a Markov Decision Process ( MDP ) of a tuple ( S , A , P , R , γ ) , where S is the state space , A is the action space , P : S × A × S → [ 0 , 1 ] is transition probabilities , R : S × A → R is the reward function , and γ ∈ [ 0 , 1 ) is the discount factor . The goal of reinforcement learning is to learn a policy π ( a|s ) that maximizes the objective of cumulative rewards over finite time steps , max π T∑ t=0 R ( st , at ) , where st ∈ S and at ∈ A are state and taken action at time step t respectively . 3 DRL-BASED RUNTIME PRUNING FRAMEWORK . The overview of our proposed framework is presented in Fig . 1 . To prune convolutional layer t , we learn two types of learnable channel importance : runtime channel importance ur ∈ RC×1 and static channel importance us ∈ RC×1 , where C is the number of channels in layer t. The runtime channel importance ur is generated by a subnetwork importance predictor f ( · ) , which takes the input feature map Fin as input , while the static channel importance us is randomly initialized and updated during training . Both ur and us indicate the channel importance of the full precision output feature map Fout through a convolution layer . Channels are selected to be pruned according to the values of each element in ur and us , and how many channels to be selected is decided by the sparsity ratios dr and ds , respectively . To learn the sparsity ratios dr and ds , two DRL agents , the runtime agent and the static agent , are introduced , where actions art and a s t are defined to set values of dr and ds , respectively . The detail of the two DRL agents are described in Sec . 3.3 . Consequently , a trade-off pruner g ( · ) is performed to balance the runtime and static pruning results , and output a decision mask M of binary values ( 1/0 ) to indicate which channels to be pruned ( 1 : pruned , 0 : preserved ) , as well as a unified channel importance vector u ∈ RC×1 as follows , [ M , u ] = g ( ur , us , dr , ds ) . ( 1 ) The final output after pruning is constructed by multiplying the full precision output feature map Fout , by 1−M and u as , F̂out = Fout ⊗ ( 1−M ) ⊗ u , ( 2 ) where⊗ is the broadcast element-wise multiplier , and 1 is the matrix of the same size as M with all the elements being 1 . In the following , we introduce how to learn the runtime channel importance vector ur and the static channel importance vector us in Sec . 3.1 , how to construct the trade-off pruner g ( · ) in Sec . 3.2 , and how to design the two DRL agents in Sec . 3.3 .

Summary: This paper proposes to learn static and dynamic channel pruning policies for convolutional neural networks. Static pruning depends only on the training dataset and is computed once before the model is deployed. Dynamic pruning is input-dependent. The policies are obtained with deep reinforcement learning on the training dataset using a combination of the loss function and storage/computation resource budgets as a reward signal. The key novelty in this paper is to combine static and dynamic pruning which can obtain the benefits from both worlds. Experimentally, the learned pruning policies are competitive with recent dynamic pruning approaches on CIFAR-10 and ILSVRC2012, in terms of both final test accuracy and number of parameters/inference time.