Discrimination-aware Network Pruning
for Deep Model Compression

Since 2012, deep neural networks (DNNs) have achieved great success in many computer vision tasks, e.g., image classification [35, 68, 20], face recognition [62, 69, 9], object detection [59, 60, 43], image generation [12, 5, 14] and video analysis [64, 77, 84]. However, the large model size and high computational costs remain great obstacles for many applications, especially on some constrained devices with limited memory and computational resources. To address this problem, model compression is an effective approach, which aims to reduce the model redundancy without significant degeneration in performance.

Recent studies on model compression mainly contain three categories: quantization [18, 58, 86], sparse or low-rank compression [16, 18, 85], and network pruning [45, 50, 82, 80]. Network quantization seeks to represent weights and activations with low bitwidth fixed-point integers, thus convolution operations can be implemented by efficient XNOR-popcount bitwise operations for substantial speedup. However, the training can be very difficult since the non-differentiable quantizer that transforms the continuous weights/activations into the discrete ones would inevitably bring errors and hamper the model performance [58]. Sparse connections methods can obtain a high compression rate in theory, but they may generate irregular convolutional kernels that need carefully designed sparse matrix operations. Low-rank compression methods seek to approximate the original filters of a pre-trained model with low-rank filters. Nevertheless, they are often inefficient for the convolutions with small kernels sizes, e.g., $1{\times}1$ [71]. In contrast, network pruning reduces the model size and speeds up the inference by removing the redundant modules (channels [24, 91, 45] or kernels [2, 52]). In particular, channel pruning can be well supported by existing deep learning libraries with little additional effort compared with network quantization and sparse or low-rank connections. More critically, most compression methods, such as quantization, can be easily applied on top of network pruning. For example, pruning redundant channels/kernels is able to further reduce the model size and accelerate the inference speed of the quantized models by reducing the number of parameters [18].

In network pruning, how to identify the informative (or important) channels/kernels (also known as channel/kernel selection) is an important problem. Existing methods can be divided into two categories, namely, training-from-scratch methods [1, 45, 78] and reconstruction-based methods [24, 26, 38, 50]. Training-from-scratch methods directly learn the importance of channels/kernels with sparsity regularization, but it is very difficult to train very deep networks on large-scale datasets [1, 45]. The reconstruction-based methods seek to perform network pruning by minimizing the reconstruction error of feature maps between the pruned model and the pre-trained one [24, 50]. However, the performance is highly affected by the quality of the pre-trained model. If the pre-trained model is not well trained, the pruning performance can be limited. More importantly, these methods suffer from a critical limitation: the redundant channels/kernels may be mistakenly kept to minimize the reconstruction error of feature maps. Consequently, these methods may incur severe performance degradation on more compact and deeper models, such as MobileNet [25, 61] for large-scale datasets.

In this paper, we aim to overcome the drawbacks of both strategies. In contrast to existing methods [24, 26, 38, 50], we assume and highlight that an informative channel/kernel, no matter where it is, should have sufficient discriminative power; otherwise, it should be removed. Based on this intuition, we propose a discrimination-aware channel pruning (DCP) method to find the channels that actually contribute to the discriminative power of the network. In DCP, relying on a pre-trained model, we first introduce multiple additional discrimination-aware losses into the network to increase the discriminative power of the intermediate layers. Then, we perform channel selection to find the most discriminative channels for each layer by considering both the discrimination-aware loss and the reconstruction error of feature maps. In this way, we are able to make a balance between the discriminative power of the channels and feature maps reconstruction. Note that a channel (3D tensor) consists of a set of kernels (each with a 2D matrix). In practice, some kernels in the selected channels may be redundant and fail to contribute to the discriminative power of the network, which leads to kernel level redundancy. To solve this issue, we propose a discrimination-aware kernel pruning (DKP) method to find the kernels with discriminative power.

Our main contributions are summarized as follows.

• •

  We propose a discrimination-aware channel/kernel pruning (DCP/DKP) scheme to compress deep models with the introduction of additional discrimination-aware losses. The proposed methods first fine-tune the model with the additional losses and the final objective. Then, we conduct channel/kernel selection by simultaneously considering the additional loss and the reconstruction error of feature maps. In this way, the proposed method is able to find the channels/kernels that actually contribute to the discriminative power of the network.

• •

  We formulate the channel/kernel selection problem as a constrained optimization problem and propose a greedy method by solving the resultant convex optimization problem to select informative channels/kernels.

• •

  We propose a new adaptive stopping condition to prevent DCP/DKP from selecting too many channels/kernels when determining the number of selected channels/kernels. Specifically, the stopping condition incorporates an additional constraint which enforces the number of selected channels to be no more than a predefined value for each layer.

• •

  Extensive experiments demonstrate the superior performance of our methods on a variety of architectures. For example, on ILSVRC-12 [8], when pruning 30% channels from ResNet-50, DCP improves the original model by 0.36% in terms of Top-1 accuracy. To demonstrate the effectiveness of our methods, we also deploy the pruned models on a smartphone (equipped with a Qualcomm Snapdragon 845 processor) and show significant acceleration on the mobile CPU.

This paper extends our preliminary version [91] from several aspects. 1) We propose two training techniques to reduce the computational overhead of DCP while still maintaining comparable or even better performance. 2) We apply the improved DCP to more compact models (i.e., MobileNetV1 and MobileNetV2) and achieve promising performance on ILSVRC-12. We further deploy the pruned models on a smartphone with a Qualcomm Snapdragon 845 processor to investigate the inference acceleration. 3) We apply the improved DCP to compress the latest face recognition models. 4) We extend DCP for kernel pruning to further compress models at kernel level. 5) We propose a new adaptive stopping condition for the optimization. 6) We provide more ablative studies to investigate the effectiveness of our methods.

Let $\{{\bf x}_{i},y_{i}\}_{i=1}^{N}$ be the training samples, where $N$ indicates the number of samples. Given an $L$-layer deep network $M$, let ${\bf W}\in{\mathbb{R}}^{n\times c\times h_{f}\times z_{f}}$ represents the model parameters w.r.t. the $l$-th convolutional layer (or block). Here, $h_{f}$ and $z_{f}$ denote the height and width of the filters, respectively; $c$ and $n$ denote the number of input channels and output filters, respectively. The parameter ${\bf W}$ contains $n\times c$ kernels in total. Each kernel, for example, the kernel ${\bf W}_{j,k}\in{\mathbb{R}}^{h_{f}\times z_{f}}$ w.r.t. the $k$-th input channel and $j$-th filter, is a matrix with the dimension of $h_{f}\times z_{f}$. Let ${\bf X}\in{\mathbb{R}}^{N\times c\times h_{in}\times z_{in}}$ and ${\bf O}\in{\mathbb{R}}^{N\times n\times h_{out}\times z_{out}}$ be the input feature maps and the involved output feature maps, respectively. Here, $h_{in}$ and $z_{in}$ denote the height and width of the input feature maps, respectively; $h_{out}$ and $z_{out}$ represent the height and width of the output feature maps, respectively. Moreover, let ${\bf X}_{i,k}\in{\mathbb{R}}^{h_{in}\times z_{in}}$ be the input feature map of the $k$-th channel for the $i$-th sample. The output feature map w.r.t. the $j$-th filter of $\bf{W}$ (i.e., ${\bf{W}}_{j}$) for the $i$-th sample, denoted by ${\bf O}_{i,j}\in{\mathbb{R}}^{h_{out}\times z_{out}}$, is obtained by performing convolution on the input feature maps ${\bf X}_{i}$ of the $i$-th sample using the kernel ${\bf{W}}_{j}$:

where $*$ denotes the convolutional operation.

Given a pre-trained model $M^{b}$, Channel Pruning aims to prune those redundant channels in ${\bf W}$ to reduce the model size and accelerate the inference speed in Eq. (1). In order to choose channels, we introduce a variant of the $\ell_{2,0}$-norm¹¹1We can also use other norms to compute the number of selected channels. $||{\bf W}||_{2,0}=\sum_{k=1}^{c}\Omega(\sum_{j=1}^{n}||{\bf W}_{j,k}||_{F})$, where $\Omega(a)=1$ if $a\neq 0$ and $\Omega(a)=0$ if $a=0$, and $||\cdot||_{F}$ represents the Frobenius norm. In order to induce sparsity, we impose an $\ell_{2,0}$-norm constraint on ${\bf W}$:

where $\kappa_{c}^{l}$ denotes the desired number of channels at layer $l$. Or equivalently, given a predefined pruning rate $\eta\in(0,1)$, it follows that $\kappa_{c}^{l}=\lceil(1-\eta)c\rceil$. When some channels of ${\bf W}$ are removed, the computation w.r.t. these channels can be effectively avoided. As a result, the pruned models would have fewer parameters and lower computational costs than the original models.

Identifying the informative channels, also known as channel selection, is an important problem in channel pruning. Most existing methods [24, 49] conduct channel pruning by minimizing the reconstruction error of feature maps between the pre-trained model and the pruned one. However, merely minimizing the reconstruction error may cause some redundant channels to be mistakenly selected, even though they are actually irrelevant to the discriminative power of the network. This issue will be even severer when the network becomes deeper.

In this paper, we highlight that an informative channel, no matter where it is, should contribute to the discriminative power of the network; otherwise, it should be removed. Based on this intuition, we propose a discrimination-aware channel pruning (DCP) scheme to find the channels that actually contribute to the discriminative power. To this end, relying on a pre-trained model, we first introduce multiple discrimination-aware losses into the network to increase the discriminative power of the intermediate layers. Then, we conduct channel selection to select the most discriminative channel by considering both the discrimination-aware loss and the reconstruction error of the feature maps. In the following, we will illustrate the details of our method.

The proposed DCP can select the channels that actually contribute to the discriminative power of the network. However, there exists a limitation for DCP. Specifically, channel pruning assumes that all kernels in a channel are equally important, which may hamper the performance of the pruned model. In fact, some kernels may not contribute to the discriminative power of the network, resulting in performance degradation of the pruned model. More critically, in DCP, once a channel is pruned (or not selected), even if there may still exist some informative kernels related to it, we are no longer able to find and keep those useful kernels related to the removed channel, which may result in suboptimal performance.

To solve this issue, we propose a kernel pruning method called discrimination-aware kernel pruning (DKP) to further compress deep networks by removing redundant kernels. Similar to DCP, we introduce a variant of the $\ell_{2,0}$-norm constraint on ${\bf W}$ to conduct kernel pruning:

where $\kappa_{ker}^{l}$ is the desired number of kernels at layer $l$. When some kernels are removed, the corresponding computational cost w.r.t. the kernels can be effectively reduced.

Starting from a pre-trained model, DKP introduces $P$ additional losses $\{{\mathcal{L}}^{p}_{S}\}_{p=1}^{P}$ evenly to the intermediate layers. Then, DKP fine-tunes the model using the addition losses $\{{\mathcal{L}}^{p}_{S}\}_{p=1}^{P}$ and the final loss ${\mathcal{L}}_{f}$ to improve the discriminative power of the intermediate layers. After that, DKP conducts kernel selection for each layer in a layer/stage-wise manner.

Similar to DCP, we introduce a variant of the $\ell_{2,0}$-norm constraint into the loss function in Eq. (6). Thus, the optimization problem for DKP can be formulated as:

To solve Problem (13), we propose a greedy algorithm for kernel selection similar to Algorithm 2. At the beginning of the kernel selection, DKP removes all the kernels by setting ${\bf W}^{0}={\bf 0}$. Instead of choosing the channels, we choose $B$ kernels with the largest gradients $||{\bf G}_{j,k}^{t-1}||_{F}={\partial{\mathcal{L}}}/{\partial{\bf W}_{j,k}^{t-1}}$ and put their indices into ${\mathcal{J}}_{t}$. Then, we update ${\mathcal{A}}_{t}$ by ${\mathcal{A}}_{t}={\mathcal{A}}_{t-1}\cup{\mathcal{J}}_{t}$. Once ${\mathcal{A}}_{t}$ is determined, we optimize ${\bf W}^{t-1}$ w.r.t. the selected kernels by minimizing Subproblem (8), which is similar to DCP. In practice, we may directly apply DKP to compress models (denoted by DKP Only) or sequentially perform kernel selection after performing channel pruning (denoted by DCP+DKP). We investigate the effectiveness of DKP in Section 6.3.

To demonstrate the effectiveness of the proposed method, we apply DCP to various architectures, such as ResNet [20], MobileNetV1 [25] and MobileNetV2 [61], etc. We conduct experiments on both image classification and face recognition. In order to verify the effectiveness of DKP, we apply DKP to ResNet [20] and VGGNet [65]. All implementations are based on PyTorch [56].

We organize the experiments as follows. First, we evaluate the proposed DCP on image classification in Section 6.1. Second, we apply DCP to face recognition in Section 6.2. Last, we evaluate the proposed DKP in Section 6.3.

In this section, we conduct ablative studies for the proposed DCP. 1) We investigate the effect of using different pruning rates in Section 7.1. 2) We explore the effect of using different $\lambda$ in Section 7.2. 3) We explore the effect of efficient training strategies in Section 7.3. 4) We explore the effect of using different $B$ in Section 7.4. 5) We discuss the effect of the tolerance $\epsilon$ in the proposed stopping conditions in Section 7.5. 6) We compare different stopping conditions in Section 7.6. 7) We visualize the feature maps w.r.t. the pruned/selected channels of ResNet-18 in Section 7.7. 8) We explore the influence of the number of additional losses in Section 7.8.

In this paper, we have proposed a discrimination-aware channel pruning (DCP) method for the compression of deep neural networks. Specifically, we first introduce additional losses to improve the discriminative power of the network and then perform channel selection in a layer-wise manner. In order to select informative channels, we have formulated the channel pruning as a sparsity-induced optimization problem and proposed a greedy algorithm to solve it. Based on DCP, we have further proposed several techniques to accelerate the channel pruning process. Moreover, we have also proposed a discrimination-aware kernel pruning (DKP) method to perform model compression at the kernel level. Extensive results on both image classification and face recognition tasks have demonstrated the effectiveness of the proposed methods.

In the future, we may extend our method in two aspects. First, in the proposed DCP, we perform channel/kernel selection in a layer-wise manner. This strategy, however, may result in suboptimal performance as only a single layer is considered each time. To address this, we may consider multiple layers/blocks each time to improve the pruning performance. Second, we will extend the proposed DCP scheme for model quantization. Specifically, based on DCP, we can conduct the quantization for each selected channel/kernel, and use the reconstruction residual to choose the next channel/kernel. In this way, we can perform channel/kernel pruning and quantization simultaneously.

This work was partially supported by Key-Area Research and Development Program of Guangdong Province 2018B010107001, National Natural Science Foundation of China (NSFC) 61836003 (key project), National Natural Science Foundation of China (NSFC) 62072190, and Program for Guangdong Introducing Innovative and Enterpreneurial Teams 2017ZT07X183.