DMCP: Differentiable Markov Channel Pruning for Neural Networks

Introduction

Channel pruning has been widely used for model acceleration and compression. The core idea behind is that large CNN models are regarded as over-parameterized. By removing the large model’s unnecessary or less important weights, we can obtain a more efficient and compact model with a marginal performance drop. Conventional channel pruning methods mainly rely on the human-designed paradigm. A typical pipeline of conventional pruning method can be summarized as three stages: pre-train a large model, prune “unimportant” weights of the large model according to the pre-defined criterion, fine-tune the pruned model [6, liu2017learning, 16, 11].

Recent work showed a new perspective of channel pruning that the structure of the pruned model is the key of determining the performance of a pruned model, rather than the inherited “important” weights. Based on this observation, some works try to design a pruning process to directly search optimal sub-structure from the unpruned structure. AMC adopted reinforcement learning (RL) to train a controller to output the pruning ratio of each layer in the unpruned structure, while MetaPruning used evolution algorithm to search structures. However, the optimization of these pruning process need to train and evaluate a large number of structures sampled from the unpruned network, thus the scalability of these methods is limited. Although AMC don’t fine-tune the pruned structures and MetaPruning trained a meta-network to predict network’s weights to avoid training the searched structures, the limitation of scalability still remains.

A similar problem in neural architecture search (NAS) has been tackled by differentiable method DARTS . However, the differentiable method proposed by DARTS cannot be directly applied to channel pruning. First, the definition of search space is different. The search space of DARTS is a category of pre-defined operations (convolution, max-pooing, etc), while in the channel pruning, the search space is the number of channels in each layer. Second, the operations in DARTS are independent with each other. But in the channel pruning, if a layer has $k+1$ channels, it must have at least $k$ channels first, which has a logical implication relationship.

In this paper, we propose a novel differentiable channel pruning method named Differentiable Markov Channel Pruning (DMCP) to perform efficient optimal sub-structure searching. Our method makes the channel pruning differentiable by modeling it as a Markov process. In the Markov process for each layer, the state $S_{k}$ represents the $k^{th}$ channel is retained, and the transition from $S_{k}$ to $S_{k+1}$ represents the probability of retaining the $(k+1)^{th}$ channel given that the $k^{th}$ channel is retained. Note that the start state is always $S_{1}$ in our method. Then the marginal probability for state $S_{k}$, i.e. the probability of retaining $k^{th}$ channel, can be computed by the product of transition probabilities and can also be viewed as a scaling coefficient. Each scaling coefficient is multiplied to its corresponding channel’s feature map during the network forwarding. So the transition probabilities parameterized by learnable parameters can be optimized in an end-to-end manner by gradient descent with respect to task loss together with budget regularization (e.g. FLOPs constraint). After the optimization, the model within desired budgets can be sampled by the Markov process with learned transition probabilities and will be trained from scratch to achieve high performance. The details of our design will be presented in Section 3.

Finally, to demonstrate the effectiveness of our method, we conduct exhaustive classification experiments on ImageNet . At the same FLOPs, our method outperforms all the other pruning methods both on MobileNetV2 and ResNet, as shown in Figure 1. With our method, MobileNetV2 has 0.1% accuracy drop with 30% FLOPs reduction and the FLOPs of ResNet-50 is reduced by 44% with only 0.4% drop.

Related Work

In this section, we discuss related works from network architecture search (NAS) and channel pruning.

Neural Architecture Search. first proposed to search for neural architectures with reinforcement learning to achieve competitive accuracy with the given inference cost. But the searching cost is too expensive to be applied broadly. Recent works try to reduce the searching cost by gradient-based methods. DARTS used a set of learnable weights to parameterize the probabilities of each candidate operation, the output of a layer is the linear combination of probabilities and feature maps of corresponding operation. After training, the operation with the highest probability is chosen to be the final architecture. However, DARTS is performed on a small proxy task (e.g. CIFAR10) and transfer the searched architecture to large scale target tasks (e.g. ImageNet). ProxylessNAS avoided using proxy tasks by only sampling two paths to search for architecture on large scale target tasks. Different from searching architecture with different types of operations in the NAS methods mentioned above, our method focuses on searching structures with a different number of channels.

Channel Pruning. Previous works on channel pruning can be roughly classified into two categories, i.e. hard pruning and soft pruning. Hard pruning removes channels during iterative pruning and fine-tuning process, while soft pruning only makes the pruned channels to be or approach to zero. Hard pruning methods mainly depend on different pruning criteria, for example, weight norm , the average percentage of zeros in the output or the influence of each channel to the final loss . Soft pruning methods mainly make the pruned channels to be or approach to zero so that those channels’ influence is decreased. first zero some filters by intra-layer criterion and a calculated layer-wise ratio. Then it increases the ratio of pruned filters gradually until reaching the given computation budget. [liu2017learning] add L1 regularization on Batch Normalization’s coefficients when training, and after training the channels with small coefficients will be pruned. search for the least important filters in a binary search manner. use generative adversarial learning to learn a sparse soft mask to scaled the output of pruned filters toward zero.

Our method can be seen as soft pruning. The major difference among DMCP and the above methods is the elimination of duplicated solutions by our Markov modeling. For example, given a layer with $C$ channels, the solution space of our method is $O(C)$, but for methods mentioned above, the solution space is $O(2^{C})$ for different combinations even with the same number of channels.

Based on recent work , some work designed a search process to directly search the optimal sub-structures from the unpruned net. AMC used reinforcement learning to determinate the ratio of channels each layer should retain. MetaPruning used an evolution algorithm to search network structures and a meta network is trained to predict weights for network structures during searching. These methods need to train or evaluate a large number structures, which makes them inefficient, while our method can be optimized by gradient descent and avoid the problem.

Method

In this section, we will give a detailed explanation of the proposed Differentiable Markov Channel Pruning (DMCP) method. As illustrated in Section 3.1, the channel pruning is first formulated as a Markov process parameterized by architecture parameters and can be optimized in an end-to-end manner. Then in Section 3.2, the training procedure of DMCP can be divided into two stages: in stage 1, the unpruned network is updated by our proposed variant sandwich rule, while in stage 2, the architecture parameters are wrapped into the unpruned network and get updated, as shown in Figure 2 (a). After the optimization, we propose two ways to sample the pruned network in Section 3.3.

Let $M(L^{(1)},L^{(2)},...,L^{(N)})$ denote the $N$-layer unpruned network, where $L^{(i)}$ is the $i^{th}$ layer. In layer $L^{(i)}$ with $C_{out}^{(i)}$ convolutional filters (i.e. channels), given input $x$, the output $O^{(i)}$ can by computed by:

As illustrated in Section 1, we perform the channel pruning in a reversed way, which can be represented by a directed ascyclic graph, as shown in Figure 3, where the state $S_{k}(1\leq k\leq C_{out})$ represent $k^{th}$ channel is retained during the pruning process, and the transition $p_{k}$ from $S_{k}$ to $S_{k+1}$ means $(k+1)^{th}$ channel is retained if $k^{th}$ channel is retained. The pruning process can be ended by transferring to the terminal state $T$ from any other state. This process has the property that if $k$ out of $C_{out}$ channels are retained in layer $L$, they must be first $k$ channels. In other words, given $k^{th}$ channel is retained, then first $(k-1)$ channels must be retained, and we can further conclude that retaining $(k+1)^{th}$ channel is conditional independent of first $(k-1)$ channels give $k^{th}$ channel is retained, which follows the Markov property.

We model transition in aforementioned ascyclic graph as a stochastic process and parameterized the transition probabilities by a set of learnable parameters. We name the learnable parameters as architecture parameters for distinguishing them from network weights. Let $p(w_{1},w_{2},...,w_{k-1})$ be the probability that first $k-1$ channels are retained. The probability of retaining first $k$ channels can be represented as:

where $p(w_{k}|w_{1},...,w_{k-1})$ is the probability of retaining $k^{th}$ channel given first $(k-1)$ channels are retained. Since retaining $w_{k}$ is conditionally independent of $\{w_{1},w_{2},...w_{k-2}\}$ given $w_{k-1}$ is retained, hence we can rewrite Equation 2 as:

in which $\neg w_{k-1}$ means $(k-1)^{th}$ channel is discarded.

Therefore, in Figure 3, the transitions can be represented by a set of transition probabilities $P=\{p_{1},p_{2},..p_{C_{out}}\}$ that defined by Equation 3.

We use a set of architecture parameters $A=\{\alpha_{1},\alpha_{2},...,\alpha_{c_{out}}\}$ to parameterize $P$, therefore $p_{k}$ can be computed as follows:

Note that we leave at least one channel for each layer, so $p_{1}=p(w_{1})=1$.

The marginal probability of sampling channel $w_{k}$ denoted by $p(w_{k})$ can be computed as:

Then the architecture parameters are wrapped into the unpruned network by following equation:

where $\hat{O_{k}}$ is the actual output of $k^{th}$ channel. Therefore, pruning $w_{k}$ can be represented by setting $p(w_{k})$ to zero.

However, we cannot directly implement Equation 7 right after the convolutional layer, because the batch normalization layer can scale up the value of $i^{th}$ channel such that the latter layer will not be affected. So the pruning process should be put after the batch normalization layer. An example of how to combine architecture parameters with an unpruned network is given in Figure 2 (b).

By the above definition, the pruned model can be sampled by the Markov process, while the transitions can be optimized by gradient descent, which will be illustrated in Section 3.2.

1.2 Solution to Shortcut Issue

Note that both MobilenetV2 and ResNet have residual blocks with shortcut connections. For the residual blocks with identity shortcuts, the number of channels in the last convolutional layer must be the same as the one in previous blocks due to the element-wise summation. Many previous works don’t prune the last convolutional layer of the residual block. In our method, we adopt a weight sharing strategy to solve this issue such that the layers, whose output channels must be equal, will share the same set of architecture parameters.

1.3 Budget Regularization

FLOPs and latency are commonly used in evaluating pruning methods. To perform an easy-to-implement and fair comparison, we use accuracy at certain FLOPs as budget regularization. However, budget regularization like FLOPs cannot be naturally optimized by gradient descent. In this section, we introduce our solution to handle the non-differentiable budget regularization problem.

In layer $L$, the expected channel $E(channel)$ can be computed as:

where $C_{out}$ is the number of output channels in $L$ and $p(w_{i})$ is the marginal probability defined in Equation 6.

In layer $L$, given expected input channels $E(in)$ and output channels $E(out)$ computed as Equation 8, the expected FLOPs $E(L_{FLOPs})$ can be computed by:

where $groups=1$ for normal convolution and $groups=E(in)$ for depth-wise convolution. $S_{I}$ and $S_{K}$ indicate input width/height and kernel width/height respectively, while $S_{P}$ is padding size and $stride$ is convolution stride.

Then the expected flops of the model $E(N_{FLOPs})$ is:

in which N is the number of convolutional layers. With Equation 12, we can optimize FLOPs by gradient descent.

1.4 Loss Function

Given target FLOPs $FLOPs_{target}$, we formulated the differentiable budget regularization loss $loss_{reg}$ as follows:

To make $E(N_{FLOPs})$ strictly lower than $FLOPs_{target}$ but not too sensitive around the target, we add single side margin to the loss function, i.e. when $\gamma\times FLOPs_{target}\leq E(N_{FLOPs})\leq FLOPs_{target}$ is satisfied, the loss will be zero. $\gamma<1$ is the tolerance ratio that can be adjusted by users.

When updating weights, the FLOPs loss has no effect on weights, so the loss function is:

where $loss_{cls}$ is cross entropy loss for classification. When updating the architecture parameters, the loss function is formulated as below:

where $\lambda_{reg}$ is hyper-parameters to balance two loss terms.

Note that we don’t add weight decay to architecture parameters. Because when the probability of keeping some channels approaching to zero or one, the norm of learnable parameters $\alpha$ will become very large, which will make them move forward to zero and hinge the optimization.

2 Training Pipeline

As illustrated in Figure 2 (a), the training procedure of DMCP can be divided into two stages, i.e. weight updating of the unpruned network and architecture parameters updating. The stage 1 and stage 2 are called iteratively during the training.

Stage 1: Weight updating of the unpruned network. In the first stage, we only update weights in the unpruned network. As defined in Equation 6, the probability of retaining the $k^{th}$ channel can also be regarded as the probability of retaining the first $k$ channels. Then our method can be seen as soft sampling all sub-structures in a single forwarding when updating architecture parameters. In general, all channels in a layer are equal and it is not intuitive to modeling the channel selection as Markov process. Inspired by previous work , which proposed a “sandwich rule” training method that the $0.75\times$ parts of the trained MobileNetV2 $1.0\times$ can get similar performance to it trained from scratch. , we introduce a variant sandwich rule, into the training scheme to make the channel groups in the unpruned model more “important” than the channel groups right after it. So that channels in a layer are not equal. The best choice of a layer with $k$ channels will be the first $k$ channels instead of other possible combinations. Based on this channel importance ranking property in the unpruned model, when sampling a sub-network with $k(k<C)$ channels, selecting the first $k$ channels can better indicate the true performance of the sub-network (trained from scratch individually). Therefore, it is reasonable to introduce Markov modeling.

There are two differences between our variation and the original “sandwich rule”. First, the randomly sampled switch (the ratio of retained channels) in each layer is not the same. Because the pruned network may have different switches in different layers. Second, the random sampling of switches obeys distribution from architecture parameters with the Markov process, instead of uniform distribution. Because the possible number of architecture in our method is much more than . And to make all architectures reflect their true performance will need too much costs. Thus we only focus on the frequently sampled architectures.

Stage 2: Architecture parameter updating. In the second stage, we only update architecture parameters. For each convolutional layer in the unpruned net, an architecture parameter is incorporated with its original output tensors by Equation 7. So that gradients could be backpropagated to the architecture parameters. And the gradients will be backpropagated to $\alpha$ by following formulas:

Such that all components of our method can be trained in an end-to-end manner. To further reduce the search space, we divide the channels into groups ($\geq 10$ groups) uniformly and each architecture parameter $\alpha$ is responsible for one group instead of only one channel. Each layer has the same number of groups.

Warmup process. Before iteratively called stage 1 and stage 2, DMCP first runs stage 1 for several epochs to warm up, in which the sub-networks are sampled by Markov process with randomly initialized architecture parameters. This process aims to avoid the network dropping into bad local minima when updating architecture parameters caused by weights’ insufficient training. We also conduct ablation study in section 4.2 to show the effectiveness of using warmup.

3 Pruned Model Sampling

After DMCP training done, we then produce models that satisfy the given cost constrain from it. In this section, we will introduce two producing methods. The first method, named Direct Sampling (DS), is to sample in each layer independently by the Markov process with optimized transition probabilities. We sample several structures and only keep the structures that lie in the target FLOPs budget.

The second method, named Expected Sampling (ES), is to set the number of channels in each layer to be the expected channels computed by Equation 8. In our experiment, $loss_{reg}$ is always optimized to zero, so the FLOPs of the expected network is equal or less than the given FLOPs constraint. Thus the expected network also satisfies the requirements.

In Section 4, we perform plenty of experiments to compare these two methods. The best performance of the pruned model sampled from Direct Sampling is a little bit higher than the one produced by Expected Sampling method, but it takes a much longer time to find such a model. So in our experiments, we use the Expected Sampling method to produce the final pruned model.

Experiments

In this section, we perform a large number of experiments to validate and analyze our method. We first describe the implementation details of DMCP in Section 4.1. To study the effectiveness of each component in our method, we conduct ablation experiments in Section 4.2. Finally in Section 4.3, we compare our results with state-of-the-art channel pruning methods. More visualizations and experiments will be shown on Supplemental Materials.

We demonstrate the effectiveness of our proposed differentiable pruning method on ImageNet classification , which contains 1000 classes. We perform experiments on both light (MobileNetV2 ) and heavy (ResNet ) models. For MobilenetV2, we use MobilenetV2 1.5x as the unpruned net, and the channels in each layer are divided into 15 groups (0.1x for each group). While for ResNet, we use standard ResNet50 (1.0x) and ResNet18 (1.0x) as unpruned structures, the channels in each layer are divided into 10 groups (0.1x for each group).

DMCP training. As described in Section 3.2, the training pipeline of DMCP contains two phases: warmup and iterative training. The training is conducted on 16 Nvidia GTX 1080TI GPUs with a batch size of 1024. Both MobileNetV2 and ResNet are trained for 40 epochs in total, the initial learning rate for both unpruned net and architecture parameters updating is 0.2 and reduced to 0.02 by cosine scheduler finally.

In the warmup phase, only the network weights are trained for 20 epochs using a variant of “sandwich rule”. In the iterative training phase, architecture parameters and unpruned net are both trained in a total of 20 epochs. The $\lambda_{reg}$ of budget regularization is set to 0.1 in all experiments. The tolerance ratio $\gamma$ is set to be 0.95 in all the experiments. To make the explanation brief in the following sections, we use the shortened form of experiment settings. For example, MBV2 1.0x-59M means the unpruned net is MobileNetV2 1.0x with target FLOPs equals to 59M.

Pruned network training. The pruned networks are produced by Direct Sampling or Expected Sampling. The details of the pruned model producing methods are illustrated in Section 3.3. Note that all pruned models are trained from scratch. The training of pruned models is performed on 32 Nvidia GTX 1080TI GPUs with a batch size of 2048. The pruned MobileNetV2 is trained for 250 epochs and pruned ResNet is trained for 100 epochs. The initial learning rate for training all pruned models is first warming up from 0.2 to 0.8 within one epoch, then is reduced to 0 by cosine scheduler.

2 Ablation Study

Recoverability verification. One property of our method should have is that it should retain nearly all channels when searching on a pre-trained model without FLOPs constraint. We use pre-trained MobileNetV2 1.0x and randomly initialize the architecture parameters. Note that only the iterative training phase is performed. We freeze the weight of MobileNetV2 1.0x and trained architecture parameters with only task loss. The result in Figure 4 shows that the FLOPs and top-1 training accuracy of our method can recover to those of the pre-trained model within 500 iterations.

Expected sampling and Direct Sampling. As described in Section 3.3, we can sample pruned models by Direct Sampling (DS) and Expected Sampling (ES). We verify the effectiveness of two model producing methods on MobileNetV2-210M and ResNet50-1.1G. We also train MobilenetV2 0.75x and ResNet50 0.5x, whose FLOPS is 210M and 1.1G respectively, as baselines for comparison. The performance of these two baselines are 70.4% and 71.9% separately. For DS, we sample five models and the results are reported in Table 1. The table shows that the performance of all models produced by DS is better than baseline models, which means the architecture parameters converge to a high-performance sub-space. And the performance of model produced by ES is very close to the best model produced by DS, which shows the effectiveness of ES. Besides, results from Table 2 and Table 3 also show the robustness of the ES. For saving the cost of fine-tuning, we use the ES to produce a model if not indicated.

The scale of the unpruned network. In this section, we evaluated the influence of scaling the unpruned network. We use two scales of MobileNetV2, i.e. MobileNetV2 1.0x and MobileNetV2 1.5x, as unpruned network, and prune them into 59M and 210M FLOPs. Note that in the experiments, the channels in each layer are divided into 10 groups to maintain the same group size. The results showed in Table 2 indicate that our method is not sensitive to the unpruned network scale. Using the larger unpruned network can lead to a little bit better performance. So we use MobileNetV2 1.5x and ResNet50 1.0x as our default unpruned network in the remaining paper.

We also visualize the difference computed by subtracting the number of channels each layer in MBV2 1.0x-210M from that in MBV2 1.5x-210M in Figure 5. From the figure, we can observe that MBV2 1.0x-210M tends to retain more channels in shallow layers while MBV2 1.5x-210M retains more channels in deep layers, even they only have a tiny difference in accuracy. This indicates that there exist multiple local minima in the search space of channel pruning.

Influence of warmup phase. We train MobileNetV2-210M with and without the warmup phase and evaluate their performance of the corresponding pruned models. To keep other settings the same, we double the epochs of the iterative training phase for the experiment without warming up. In the setting without warming up, the models are trained for 100 epochs, the initial learning rate and the scheduler in the first 50 epochs are the same as the warmup phase. The results in Table 3 shows that using warmup leads to better performance. One possible reason is that using warmup makes the weights trained more sufficiently before updating architecture parameters, which makes weights more discriminable and prevents architecture parameters from trapping into bad local minima.

Impact of the variant sandwich rule. We ablate the impact of the sandwich rule in the MobileNetV2-210M setting. The original sandwich rule and our variant sandwich rule are adopted solely in DMCP for comparison. The results are tabulated in Table 4. We can see that using the variant sandwich rule leads to better performance. The possible reason is that the weights corresponding to higher probability will be optimized better by the variant sandwich rule. And in these weights, each of them will be optimized better to represent their true importance with less influence of other weights. Thus, when updating architecture parameters, the “competition” is mainly centered on them, which makes updating more accurate.

Training scheme. We verify the effectiveness of the updating scheme. We conduct three experiments on MobileNetV2-59M. All experiments use the same setting in the warmup phase, while the settings in the iterative training phase are as follows: In the first experiment, we only update architecture parameters with respect to budget regularization (FLOPs loss); in the second experiment, we only update architecture parameters with respect to both budget regularization and task loss; and in the last experiment, we update both unpruned net and architecture parameters with respect to the full loss function. The results are shown in Table 5. The first experiment is a naive baseline of FLOPs guided pruning. The layers with the same FLOPs may be pruned to the same extent. The result is far worse than the other experiments. Comparing with the first experiment and the second experiment, we know that the task loss can help to discriminate the importance of different layer even they have same FLOPs. Finally, by comparing with the last two experiments, we can conclude that when architecture parameters changed, the weights should also be adapted.

3 Comparison with state-of-the-art

In this section, we compare our method with various pruning methods, including reinforcement learning method AMC , evolution method MetaPruning , one-shot method AutoSlim , and traditional channel pruning methods SFP and FPGM . The training settings of our method in all FLOPs settings are illustrated in Section 4.1, and our pruned models are sampled by Expected Sampling.

All methods are evaluated on MobileNetV2, ResNet18, and ResNet50, in each type of model, we trained a set of baseline model with setting 4.1 for comparison.

From the Table 6, we can see that our method outperforms all other methods under the same settings, which show the superiority of our method. Note that AMC, MetaPruning and our method train the pruned model from scratch by standard hard label loss. While AutoSlim adopts a in-place distillation method in which the pruned network share weights with unpruned net and mimic the output of the unpruned net. To fairly compare with AutoSlim, we also train our pruned model with the slimmable training method. Results show that this training method can further boost the performance, and our method surpasses AutoSlim in different FLOPs models.

Conclusion

In this paper, we propose a novel differentiable method for channel pruning, named Differentiable Markov Channel Pruning (DMCP), to solve the defect of existing methods that they need to train and evaluate a large number of sub-structures. The proposed method is differentiable by modeling the channel pruning as the Markov process, thus can be optimized with respect to task loss by gradient descent. After optimization, the required model can be sampled from the optimized transitions by a simple Expected Sampling and trained from scratch. Our method achieves state-of-the-art performance with ResNet and MobileNet V2 on ImageNet in various FLOPs settings.

References

Appendix A Visualization

We sample 3000 structures from the trained MobileNetV2-210M via the Markov process, and the distribution of their FLOPs is showed in Figure 6. From the figure, we can find that the mean of FLOPs lies around 210M, which means that the expected FLOPs converged to the desired budget 210M.

A.2 The Channel Distribution of Pruned Layers.

In this section, we examine the channel distribution in each layer of the pruned model. We sample 3000 models from MobileNetV2-210M whose FLOPs are within the desired budget (210M) by Markov process. Figure 7 shows the channel distribution of 12 layers sampled from different blocks. From the figure, we can observe that the number of channel in most layers follows an uni-modal distribution, and some layers choose to retain all the channels (e.g. LinearBottleneck6 and 7).

Appendix B Comparison between using warm-up and using pre-trained model

As described in Section 3.2, the warm-up is performed by only running stage 1 that updating the weights of the unpruned network by our proposed variant sandwich rule, which makes the channel group more “important” than the channel group right after it, providing a good initialization for iterative training. However, using a pre-trained model cannot provide initialization with the property. Our experiment also shows the superiority of using warm-up, on DMCP-MBV2 with 210M FLOPs by replacing warm-up with using pre-trained models, a 0.6% accuracy drop was observed.

Appendix C Modeling architecture parameters as independent Bernoulli variables

Given a layer with $C$ channels, the solution space of our method is $O(C)$, by modeling architecture parameters as Bernoulli variables, the solution space becomes $O(2^{C})$, as there are $2^{C}$ possible channel combinations, which makes it much harder to optimize. To demonstrate our analysis, we experiment on MobileNet-v2 with 210M FLOPs, by replacing Markov modeling with Bernoulli Modelling of architecture parameters, the performance of the pruned model is 70.1%, which is 2.3% lower than DMCP.