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 work introduces a Reinforcement Learning based framework that simultaneously learns both a static and dynamic pruning strategy. The combination allows the static pruner to decrease the required storage while the dynamic pruning can optimize the required compute using input-dependent pruned weights. The RL agent can dynamically learn the optimal sparsity distribution among the different layer of the networks while staying under a resource constraint as opposed to other methods which often enforce a layer level sparsity ratio. It demonstrates the efficacy of the algorithm on CIFAR10 and ILSVRC2012 and showed the effect of the tradeoff between static and dynamic pruning.