BWCP: Probabilistic Learning-to-Prune Channels for ConvNets via Batch Whitening

This work presents a probabilistic channel pruning method to accelerate Convolutional Neural Networks ( CNNs ) . Previous pruning methods often zero out unimportant channels in training in a deterministic manner , which reduces CNN ’ s learning capacity and results in suboptimal performance . To address this problem , we develop a probability-based pruning algorithm , called batch whitening channel pruning ( BWCP ) , which can stochastically discard unimportant channels by modeling the probability of a channel being activated . BWCP has several merits . ( 1 ) It simultaneously trains and prunes CNNs from scratch in a probabilistic way , exploring larger network space than deterministic methods . ( 2 ) BWCP is empowered by the proposed batch whitening tool , which is able to empirically and theoretically increase the activation probability of useful channels while keeping unimportant channels unchanged without adding any extra parameters and computational cost in inference . ( 3 ) Extensive experiments on CIFAR-10 , CIFAR-100 , and ImageNet with various network architectures show that BWCP outperforms its counterparts by achieving better accuracy given limited computational budgets . For example , ResNet50 pruned by BWCP has only 0.58 % Top-1 accuracy drop on ImageNet , while reducing 42.9 % FLOPs of the plain ResNet50 . 1 INTRODUCTION . Deep convolutional neural networks ( CNNs ) have achieved superior performance in a variety of computer vision tasks such as image recognition ( He et al. , 2016 ) , object detection ( Ren et al. , 2017 ) , and semantic segmentation ( Chen et al. , 2018 ) . However , despite their great success , deep CNN models often have massive demand on storage , memory bandwidth , and computational power ( Han & Dally , 2018 ) , making them difficult to be plugged onto resource-limited platforms , such as portable and mobile devices ( Deng et al. , 2020 ) . Therefore , proposing efficient and effective model compression methods has become a hot research topic in the deep learning community . Model pruning , as one of the vital model compression techniques , has been extensively investigated . It reduces model size and computational cost by removing unnecessary or unimportant weights or channels in a CNN ( Han et al. , 2016 ) . For example , many recent works ( Wen et al. , 2016 ; Guo et al. , 2016 ) prune fine-grained weights of filters . Han et al . ( 2015 ) proposes to discard the weights that have magnitude less than a predefined threshold . Guo et al . ( 2016 ) further utilizes a sparse mask on a weight basis to achieve pruning . Although these unstructured pruning methods achieve optimal pruning schedule , they do not take the structure of CNNs into account , preventing them from being accelerated on hardware such as GPU for parallel computations ( Liu et al. , 2018 ) . To achieve efficient model storage and computations , we focus on structured channel pruning ( Wen et al. , 2016 ; Yang et al. , 2019a ; Liu et al. , 2017 ) , which removes entire structures in a CNN such as filter or channel . A typical structured channel pruning approach commonly contains three stages , including pre-training a full model , pruning unimportant channels by the predefined criteria such as ` p norm , and fine-tuning the pruned model ( Liu et al. , 2017 ; Luo et al. , 2017 ) , as shown in Fig.1 ( a ) . However , it is usually hard to find a global pruning threshold to select unimportant channels , because the norm deviation between channels is often too small ( He et al. , 2019 ) . More importantly , as some channels are permanently zeroed out in the pruning stage , such a multi-stage procedure usually not only relies on hand-crafted heuristics but also limits the learning capacity ( He et al. , 2018a ; 2019 ) . To tackle the above issues , we propose a simple but effective probability-based channel pruning framework , named batch-whitening channel pruning ( BWCP ) , where unimportant channels are pruned in a stochastic manner , thus preserving the channel space of CNNs in training ( i.e . the diversity of CNN architectures is preserved ) . To be specific , as shown in Fig.1 ( b ) , we assign each channel with an activation probability ( i.e . the probability of a channel being activated ) , by exploring the properties of the batch normalization layer ( Ioffe & Szegedy , 2015 ; Arpit et al. , 2016 ) . A larger activation probability indicates that the corresponding channel is more likely to be preserved . We also introduce a capable tool , termed batch whitening ( BW ) , which can increase the activation probability of useful channels , while keeping the unnecessary channels unchanged . By doing so , the deviation of the activation probability between channels is explicitly enlarged , enabling BWCP to identify unimportant channels during training easily . Such an appealing property is justified by theoretical analysis and experiments . Furthermore , we exploit activation probability adjusted by BW to generate a set of differentiable masks by a soft sampling procedure with Gumbel-Softmax technique , allowing us to train BWCP in an online “ pruning-from-scratch ” fashion stably . After training , we obtain the final compact model by directly discarding the channels with zero masks . The main contributions of this work are three-fold . ( 1 ) We propose a probability-based channel pruning framework BWCP , which explores a larger network space than deterministic methods . ( 2 ) BWCP can easily identify unimportant channels by adjusting their activation probabilities without adding any extra model parameters and computational cost in inference . ( 3 ) Extensive experiments on CIFAR-10 , CIFAR-100 and ImageNet datasets with various network architectures show that BWCP can achieve better recognition performance given the comparable amount of resources compared to existing approaches . For example , BWCP can reduce 68.08 % Flops by compressing 93.12 % parameters of VGG-16 with merely accuracy drop and ResNet-50 pruned by BWCP has only 0.58 % top-1 accuracy drop on ImageNet while reducing 42.9 % FLOPs . 2 RELATED WORK . Weight Pruning . Early network pruning methods mainly remove the unimportant weights in the network . For instance , Optimal Brain Damage ( LeCun et al. , 1990 ) measures the importance of weights by evaluating the impact of weight on the loss function and prunes less important ones . However , it is not applicable in modern network structure due to the heavy computation of the Hessian matrix . Recent work assesses the importance of the weights through the magnitude of the weights itself . Specifically , ( Guo et al. , 2016 ) prune the network by encouraging weights to become exactly zero . The computation involves weights with zero can be discarded . However , a major drawback of weight pruning techniques is that they do not take the structure of CNNs into account , thus failing to help scale pruned models on commodity hardware such as GPUs ( Liu et al. , 2018 ; Wen et al. , 2016 ) . Channel Pruning . Channel pruning approaches directly prune feature maps or filters of CNNs , making it easy for hardware-friendly implementation . For instance , relaxed ` 0 regularization ( Louizos et al. , 2017 ) and group regularizer ( Yang et al. , 2019a ) impose channel-level sparsity , and filters with small value are selected to be pruned . Some recent work also propose to rank the importance of filters by different criteria including ` 1 norm ( Liu et al. , 2017 ; Li et al. , 2017 ) , ` 2 norm ( Frankle & Carbin , 2018 ) and High Rank channels ( Lin et al. , 2020 ) . For example , ( Liu et al. , 2017 ) explores the importance of filters through scale parameter γ in batch normalization . Although these approaches introduce minimum overhead to the training process , they are not trained in an end-to-end manner and usually either apply on a pre-trained model or require an extra fine-tuning procedure . Recent works tackle this issue by pruning CNNs from scratch . For example , FPGM ( He et al. , 2019 ) zeros in unimportant channels and continues training them after each training epoch . Furthermore , both SSS and DSA learn a differentiable binary mask that is generated by channel importance and does not require any additional fine-tuning . Our proposed BWCP is most related to variational pruning ( Zhao et al. , 2019 ) and SCP ( Kang & Han , 2020 ) as they also employ the property of normalization layer and associate the importance of channel with probability . The main difference is that our method adopts the idea of whitening to perform channel pruning . We will show that the proposed batch whitening ( BW ) technique can adjusts the activation probability of different channels according to their importance , making it easy to identify unimportant channels . Although previous work SPP ( Wang et al. , 2017 ) and DynamicCP ( Gao et al. , 2018 ) also attempt to boost salient channels and skip unimportant ones , they fail to consider the natural property inside normalization layer and deign the activation probability empirically . 3 PRELIMINARY . Notation . We use regular letters , bold letters , and capital letters to denote scalars such as ‘ x ’ , and vectors ( e.g.vector , matrix , and tensor ) such as ‘ x ’ and random variables such as ‘ X ’ , respectively . We begin with introducing a building layer in recent deep neural nets which typically consists of a convolution layer , a batch normalization ( BN ) layer , and a rectified linear unit ( ReLU ) ( Ioffe & Szegedy , 2015 ; He et al. , 2016 ) . Formally , it can be written by xc = wc ∗ z , x̃c = γcx̄c + βc , yc = max { 0 , x̃c } ( 1 ) where c ∈ [ C ] denotes channel index and C is channel size . In Eqn . ( 1 ) , ‘ ∗ ’ indicates convolution operation and wc is filter weight corresponding to the c-th output channel , i.e . xc ∈ RN×H×W . To perform normalization , xc is firstly standardized to x̄c through x̄c = ( xc − E [ xc ] ) / √ D [ xc ] where E [ · ] and D [ · ] indicate calculating mean and variance over a batch of samples , and then is re-scaled to x̃c by scale parameter γc and bias βc . Moreover , the output feature yc is obtained by ReLU activation that discards the negative part of x̃c . Criterion-based channel pruning . For channel pruning , previous methods usually employ a ‘ smallnorm-less-important ’ criterion to measure the importance of channels . For example , BN layer can be applied in channel pruning ( Liu et al. , 2017 ) , where a channel with a small value of γc would be removed . The reason is that the c-th output channel x̃c contributes little to the learned feature representation when γc is small . Hence , the convolution in Eqn . ( 1 ) can be discarded safely , and filter wc can thus be pruned . Unlike these criterion-based methods that deterministically prune unimportant filters and rely on a heuristic pruning procedure as shown in Fig.1 ( a ) , we explore a probability-based channel pruning framework where less important channels are pruned in a stochastic manner . Activation probability . To this end , we define an activation probability of a channel by exploring the property of the BN layer . Those channels with a larger activation probability could be preserved with a higher probability . To be specific , since x̄c is acquired by subtracting the sample mean and being divided by the sample variance , we can treat each channel feature as a random variable following standard Normal distribution ( Arpit et al. , 2016 ) , denoted as X̄c . Note that only positive parts can be activated by ReLU function . Proposition 1 gives the activation probability of the c-th channel , i.e.P ( X̃c ) > 0 . Proposition 1 Let a random variable X̄c ∼ N ( 0 , 1 ) and Yc = max { 0 , γcX̄c + βc } . Then we have ( 1 ) P ( Yc > 0 ) = P ( X̃c > 0 ) = ( 1 + Erf ( βc/ ( √ 2|γc| ) ) /2 where Erf ( x ) = ∫ x 0 2/ √ π · exp−t2dt , and ( 2 ) P ( X̃c > 0 ) = 0⇔ βc ≤ 0 and γc → 0 . Note that a pruned channel can be modelled by P ( X̃c > 0 ) = 0 . With Proposition 1 ( see proof in Appendix A.2 ) , we know that the unnecessary channels satisfy that γc approaches 0 and βc is negative . To achieve channel pruning , previous compression techniques ( Li et al. , 2017 ; Zhao et al. , 2019 ) merely impose a regularization on γc , which would deteriorate the representation power of unpruned channels ( Perez et al. , 2018 ; Wang et al. , 2020 ) . Instead , we adopt the idea of whitening to build a probabilistic channel pruning framework where unnecessary channels are stochastically disgarded with a small activation probability while important channels are preserved with a large activation probability .

Summary: This paper presents a probabilistic channel pruning method (BWCP) for accelerating CNNs. The key newly proposed technique is “batch whitening”. They evaluate their method on CIFAR-10/100 and ImageNet compared to other recent filter pruning methods.