MetaPruning: Meta Learning for Automatic Neural Network Channel Pruning
Zechun Liu, Haoyuan Mu, Xiangyu Zhang, Zichao Guo, Xin Yang, Tim Kwang-Ting Cheng, Jian Sun
Introduction
Channel pruning has been recognized as an effective neural network compression/acceleration method and is widely used in the industry. A typical pruning approach contains three stages: training a large over-parameterized network, pruning the less-important weights or channels, finetuning or re-training the pruned network. The second stage is the key. It usually performs iterative layer-wise pruning and fast finetuning or weight reconstruction to retain the accuracy .
Conventional channel pruning methods mainly rely on data-driven sparsity constraints , or human-designed policies . Recent AutoML-style works automatically prune channels in an iterative mode, based on a feedback loop or reinforcement learning . Compared with the conventional pruning methods, the AutoML methods save human efforts and can optimize the direct metrics like the hardware latency.
Apart from the idea of keeping the important weights in the pruned network, a recent study finds that the pruned network can achieve the same accuracy no matter it inherits the weights in the original network or not. This finding suggests that the essence of channel pruning is finding good pruning structure - layer-wise channel numbers.
However, exhaustively finding the optimal pruning structure is computationally prohibitive. Considering a network with 10 layers and each layer contains 32 channels. The possible combination of layer-wise channel numbers could be . Inspired by the recent Neural Architecture Search (NAS), specifically One-Shot model , as well as the weight prediction mechanism in HyperNetwork , we propose to train a PruningNet that can generate weights for all candidate pruned networks structures, such that we can search good-performing structures by just evaluating their accuracy on the validation data, which is highly efficient.
To train the PruningNet, we use a stochastic structure sampling. As shown in Figure 1, the PruningNet generates the weights for pruned networks with corresponding network encoding vectors, which is the number of channels in each layer. By stochastically feeding in different network encoding vectors, the PruningNet gradually learns to generate weights for various pruned structures. After the training, we search for good-performing Pruned Networks by an evolutionary search method which can flexibly incorporate various constraints such as computation FLOPs or hardware latency. Moreover, by directly searching the best pruned network via determining the channels for each layer or each stage, we can prune channels in the shortcut without extra effort, which is seldom addressed in previous channel pruning solutions. We name the proposed method as MetaPruning.
We apply our approach on MobileNets and ResNet . At the same FLOPs, our accuracy is - higher than MobileNet V1, - higher than MobileNet V2, and - higher than ResNet-50. At the same latency, our accuracy is - higher than MobileNet V1, and - higher than MobileNet V2. Compared with state-of-the-art channel pruning methods , our MetaPruning also produces superior results.
We proposed a meta learning approach, MetaPruning, for channel pruning. The central of this approach is learning a meta network (named PruningNet) which generates weights for various pruned structures. With a single trained PruningNet, we can search for various pruned networks under different constraints.
Compared to conventional pruning methods, MetaPruning liberates human from cumbersome hyperparameter tuning and enables the direct optimization with desired metrics.
Compared to other AutoML methods, MetaPruning can easily enforce constraints in the search of desired structures, without manually tuning the reinforcement learning hyper-parameters.
The meta learning is able to effortlessly prune the channels in the short-cuts for ResNet-like structures, which is non-trivial because the channels in the short-cut affect more than one layers.
Related Works
There are extensive studies on compressing and accelerating neural networks, such as quantization , pruning and compact network design . A comprehensive survey is provided in . Here, we summarize the approaches that are most related to our work.
Pruning Network pruning is a prevalent approach for removing redundancy in DNNs. In weight pruning, people prune individual weights to compress the model size . However, weight pruning results in unstructured sparse filters, which can hardly be accelerated by general-purpose hardware. Recent works focus on channel pruning in the CNNs, which removes entire weight filters instead of individual weights. Traditional channel pruning methods trim channels based on the importance of each channel either in an iterative mode or by adding a data-driven sparsity . In most traditional channel pruning, compression ratio for each layer need to be manually set based on human experts or heuristics, which is time consuming and prone to be trapped in sub-optimal solutions.
AutoML Recently, AutoML methods take the real-time inference latency on multiple devices into account to iteratively prune channels in different layers of a network via reinforcement learning or an automatic feedback loop . Compared with traditional channel pruning methods, AutoML methods help to alleviate the manual efforts for tuning the hyper-parameters in channel pruning. Our proposed MetaPruning also involves little human participation. Different from previous AutoML pruning methods, which is carried out in a layer-wise pruning and finetuning loop, our methods is motivated by recent findings , which suggests that instead of selecting “important” weights, the essence of channel pruning sometimes lies in identifying the best pruned network. From this prospective, we propose MetaPruning for directly finding the optimal pruned network structures. Compared to previous AutoML pruning methods , MetaPruning method enjoys higher flexibility in precisely meeting the constraints and possesses the ability of pruning the channel in the short-cut.
Meta Learning Meta-learning refers to learning from observing how different machine learning approaches perform on various learning tasks. Meta learning can be used in few/zero-shot learning and transfer learning . A comprehensive overview of meta learning is provided in . In this work we are inspired by to use meta learning for weight prediction. Weight predictions refer to weights of a neural network are predicted by another neural network rather than directly learned . Recent works also applies meta learning on various tasks and achieves state-of-the-art results in detection , super-resolution with arbitrary magnification and instance segmentation .
Neural Architecture Search Studies for neural architecture search try to find the optimal network structures and hyper-parameters with reinforcement learning , genetic algorithms or gradient based approaches . Parameter sharing and weights prediction methods are also extensively studied in neural architecture search. One-shot architecture search uses an over-parameterized network with multiple operation choices in each layer. By jointly training multiple choices with drop-path, it can search for the path with highest accuracy in the trained network, which also inspired our two step pruning pipeline. Tuning channel width are also included in some neural architecture search methods. ChamNet built an accuracy predictor atop Gaussian Process with Bayesian optimization to predict the network accuracy with various channel widths, expand ratios and numbers of blocks in each stage. Despite its high accuracy, building such an accuracy predictor requires a substantial of computational power. FBNet and ProxylessNas include blocks with several different middle channel choices in the search space. Different from neural architecture search, in channel pruning task, the channel width choices in each layer is consecutive, which makes enumerate every channel width choice as an independent operation infeasible. Proposed MetaPruning targeting at channel pruning is able to solve this consecutive channel pruning challenge by training the PruningNet with weight prediction, which will be explained in Sec.3
Methodology
In this section, we introduce our meta learning approach for automatically pruning channels in deep neural networks, that pruned network could meet various constraints easily.
We formulate the channel pruning problem as
where is the network before the pruning. We try to find out the pruned network channel width () for layer to layer that has the minimum loss after the weights are trained, with the cost meets the constraint (i.e. FLOPs or latency).
To achieve this, we propose to construct a PruningNet, a kind of meta network, where we can quickly obtain the goodness of all potential pruned network structures by evaluating on the validation data only. Then we can apply any search method, which is evolution algorithm in this paper, to search for the best pruned network.
Channel pruning is non-trivial because the layer-wise dependence in channels such that pruning one channel may significantly influence the following layers and, in return, degrade the overall accuracy. Previous methods try to decompose the channel pruning problem into the sub-problem of pruning the unimportant channels layer-by-layer or adding the sparsity regularization . AutoML methods prune channels automatically with a feedback loop or reinforcement learning . Among those methods, how to prune channels in the short-cut is seldom addressed. Most previous methods prune the middle channels in each block only, which limits the overall compression ratio.
Carrying out channel pruning task with consideration of the overall pruned network structure is beneficial for finding optimal solutions for channel pruning and can solve the shortcut pruning problem. However, obtaining the best pruned network is not straightforward, considering a small network with 10 layers and each layer containing 32 channels, the combination of possible pruned network structures is huge.
Inspired by the recent work , which suggests the weights left by pruning is not important compared to the pruned network structure, we are motivated to directly find the best pruned network structure. In this sense, we may directly predict the optimal pruned network without iteratively decide the important weight filters. To achieve this goal, we construct a meta network, PruningNet, for providing reasonable weights for various pruned network structures to rank their performance.
The PruningNet is a meta network, which takes a network encoding vector as input and outputs the weights of pruned network:
A PruningNet block consists of two fully-connected layers. In the forward pass, the PruningNet takes the network encoding vector (i.e., the number of channels in each layer) as input, and generates the weight matrix. Meanwhile, a Pruned Network is constructed with output channels width in each layer equal to the element in the network encoding vector. The generated weight matrix is cropped to match the number of input and output channel in the Pruned Network, as shown in Figure 2. Given a batch of input image, we can calculate the loss from the Pruned Network with generated weights.
In the backward pass, instead of updating the weights in the Pruned Networks, we calculate the gradients w.r.t the weights in the PruningNet. Since the reshape operation as well as the convolution operation between the output of the fully-connect layer in the PruningNet and the output of the previous convolutional layer in the Pruned Network is also differentiable, the gradient of the weights in the PruningNet can be easily calculated by the Chain Rule. The PruningNet is end-to-end trainable. The detailed structure of PruningNet connected with Pruned Network is shown in Figure 3.
To train the PruningNet, we proposed the stochastic structure sampling. In the training phase, the network encoding vector is generated by randomly choosing the number of channels in each layer at each iteration. With different network encodings, different Pruned Networks are constructed and the corresponding weights are provided with the PruningNet. By stochastically training with different encoding vectors, the PruningNet learns to predict reasonable weights for various different Pruned Networks.
2 Pruned-Network search
After the PruningNet is trained, we can obtain the accuracy of each potential pruned network by inputting the network encoding into the PruningNet, generating the corresponding weights and doing the evaluation on the validation data.
Since the number of network encoding vectors is huge, we are not able to enumerate. To find out the pruned network with high accuracy under the constraint, we use an evolutionary search, which is able to easily incorporate any soft or hard constraints.
In the evolutionary algorithm used in MetaPruning, each pruned network is encoded with a vector of channel numbers in each layer, named the genes of pruned networks. Under the hard constraint, we first randomly select a number of genes and obtain the accuracy of the corresponding pruned network by doing the evaluation. Then the top k genes with highest accuracy are selected for generating the new genes with mutation and crossover. The mutation is carried out by changing a proportion of elements in the gene randomly. The crossover means that we randomly recombine the genes in two parent genes to generate an off-spring. We can easily enforce the constraint by eliminate the unqualified genes. By further repeating the top k selection process and new genes generation process for several iterations, we can obtain the gene that meets constraints while achieving the highest accuracy. Detailed algorithm is described in Algorithm.1.
Experimental Results
In this section, we demonstrate the effectiveness of our proposed MetaPruning method. We first explain the experiment settings and introduce how to apply the MetaPruning on MobileNet V1 V2 and ResNet , which can be easily generalized to other network structures. Second, we compare our results with the uniform pruning baselines as well as state-of-the-art channel pruning methods. Third, we visualize the pruned network obtained with MetaPruning. Last, ablation studies are carried out to elaborate the effect of weight prediction in our method.
The proposed MetaPruning is very efficient. Thus it is feasible to carry out all experiments on the ImageNet 2012 classification dataset .
MetaPruning method consists of two stages. In the first stage, the PruningNet is train from scratch with stochastic structure sampling, which takes epochs as training a network normally. Further prolonging PruningNet training yields little final accuracy gain in the obtained Pruned Net. In the second stage, we use an evolutionary search algorithm to find the best pruned network. With the PruningNet predicting the weights for all the PrunedNets, no fine-tuning or retraining are needed at search time, which makes the evolution search highly efficient. Inferring a PrunedNet only takes seconds on 8 Nvidia 1080Ti GPUs. The best PrunedNet obtained from search is then trained from scratch. For the training process in both stages, we use the standard data augmentation strategies as to process the input images. We adopt the same training scheme as for experiments on MobileNets and the training scheme in for ResNet. The resolutions of the input image is set to 224 224 for all experiments.
At training time, we split the original training images into sub-validation dataset, which contains 50000 images randomly selected from the training images with 50 images in each 1000-class, and sub-training dataset with the rest of images. We train the PruningNet on the sub-training dataset and evaluating the performance of pruned network on the sub-validation dataset in the searching phase. At search time, we recalculate the running mean and running variance in the BatchNorm layer with 20000 sub-training images for correctly inferring the performance of pruned networks, which takes only a few seconds. After obtaining the best pruned network, the pruned network is trained from scratch on the original training dataset and evaluated on the test dataset.
2 MetaPruning on MobileNets and ResNet
To prove the effectiveness of our MetaPruning method, we apply it on MobileNets and ResNet .
MobileNet V1 is a network without shortcut. To construct the corresponding PruningNet, we have the PruningNet blocks equal to the number of convolution layers in the MobileNet v1, and each PruningNet block is composed of two concatenated fully-connected(FC) layers.
The input vector to the PruningNet is the number of channels in each layer. Then this vector is decoded into the input and output channel compression ratio of each layer, i.e., . Here, denotes the number of channels, is layer index of current layer and denotes the previous layer, means output of the original network and is the pruned output. This two dimensional vector is then inputted into each PruningNet block associated with each layer. The first FC layer in the PruningNet block output a vector with 64 entries and the second FC layer use this 64-entry encoding to output a vector with a length of . Then we reshape it to () as the weight matrix in the convolution layer, as shown in Figure.3.
In stochastic structure sampling, an encoding vector of output channel numbers is generated with its each entry independently and randomly selected from [, ], with the step being . More refined or coarse step can be chosen according to the fineness of pruning. After decoding and the weight generation process in the PruningNet, the top left part of generated weight matrix is cropped to () and is used in training, and the rest of the weights can be regards as being ‘untouched’ in this iteration, as shown in Figure.4 (a). In different iterations, different channel width encoding vectors are generated.
2.2 MobileNet V2
In MobileNet V2, each stage starts with a bottleneck block matching the dimension between two stages. If a stage consists of more than one block, the following blocks in this stage will contain a shortcut adding the input feature maps with the output feature maps, thus input and output channels in a stage should be identical, as shown in Figure 4 (b). To prune the structure containing shortcut, we generate two network encoding vectors, one encodes the overall stage output channels for matching the channels in the shortcut and another encodes the middle channels of each blocks. In PruningNet, we first decode this network encoding vector to the input, output and middle channel compression ratio of each block. Then we generate the corresponding weight matrices in that block, with a vector inputting to the corresponding PruningNet blocks, where denotes the block index. The PruningNet block design is the same as that in MobileNetV1, and the number of PruningNet block equals to the number of convolution layers in the MobileNet v2.
2.3 ResNet
As a network with shortcut, ResNet has similar network structure with MobileNet v2 and only differs at the type of convolution in the middle layer, the downsampling block and number of blocks in each stage. Thus, we adopt similar PruningNet design for ResNet as MobileNet V2.
3 Comparisons with state-of-the-arts
We compare our method with the uniform pruning baselines, traditional pruning methods as well as state-of-the-art channel pruning methods.
Table 1 compares our accuracy with the uniform pruning baselines reported in . With the pruning scheme learned by MetaPruning, we obtain 6.6% higher accuracy than the baseline 0.25 MobileNet V1. Further more, as our method can be generalized to prune the shortcuts in a network, we also achieves decent improvement on MobileNet V2, shown in Table.2 Previous pruning methods only prunes the middle channels of the bottleneck structure , which limits their maximum compress ratio at given input resolution. With MetaPruning, we can obtain 3.7% accuracy boost when the model size is as small as 43M FLOPs. For heavy models as ResNet, MetaPruning also outperforms the uniform baselines and other traditional pruning methods by a large margin, as is shown in Table.3.
In Table 4, we compare MetaPruning with the state-of-the-art AutoML pruning methods. MetaPruning achieves superior results than AMC and NetAdapt . Moreover, MetaPruning gets rid of manually tuning the reinforcement learning hyper-parameters and can obtain the pruned network precisely meeting the FLOPs constraints. With the PruningNet trained once using one-fourth epochs as normally training the target network, we can obtain multiple pruned network structures to strike different accuracy-speed trade-off, which is more efficient than the state-of-the-art AutoML pruning methods . The time cost is reported in Sec.4.1.
3.2 Pruning under latency constraint
There is an increasing attention in directly optimizing the latency on the target devices. Without knowing the implementation details inside the device, MetaPruning learns to prune channels according to the latency estimated from the device.
As the number of potential Pruned Network is numerous, measuring the latency for each network is too time-consuming. With a reasonable assumption that the execution time of each layer is independent, we can obtain the network latency by summing up the run-time of all layers in the network. Following the practice in , we first construct a look-up table, by estimating the latency of executing different convolution layers with different input and output channel width on the target device, which is Titan Xp GPU in our experiments. Then we can calculate the latency of the constructed network from the look-up table.
We carried out experiments on MobileNet V1 and V2. Table 5 and Table 6 show that the prune networks discovered by MetaPruning achieve significantly higher accuracy than the uniform baselines with the same latency.
4 Pruned result visualization
In channel pruning, people are curious about what is the best pruning heuristic and lots of human experts are working on manually designing the pruning policies. With the same curiosity, we wonder if any reasonable pruning schemes are learned by our MetaPruning method that contributes to its high accuracy. In visualizing the pruned network structures, we find that the MetaPruning did learn something interesting.
Figure 5 shows the pruned network structure of MobileNet V1. We observe significant peeks in the pruned network every time when there is a down sampling operation. When the down-sampling occurs with a stride 2 depth-wise convolution, the resolution degradation in the feature map size need to be compensated by using more channels to carry the same amount of information. Thus, MetaPruning automatically learns to keep more channels at the downsampling layers. The same phenomenon is also observed in MobileNet V2, shown in Figure 6. The middle channels will be pruned less when the corresponding block is in responsible for shrinking the feature map size.
Moreover, when we automatically prune the shortcut channels in MobileNet V2 with MetaPruning, we find that, despite the 145M pruned network contains only half of the FLOPs in the 300M pruned network, 145M network keeps similar number of channels in the last stages as the 300M network, and prunes more channels in the early stages. We suspect it is because the number of classifiers for the ImageNet dataset contains 1000 output nodes and thus more channels are needed at later stages to extract sufficient features. When the FLOPs being restrict to 45M, the network almost reaches the maximum pruning ratio and it has no choice but to prune the channels in the later stage, and the accuracy degradation from 145M network to 45M networks is much severer than that from 300M to 145M.
5 Ablation study
In this section, we discuss about the effect of weight prediction in the MetaPruning method.
We wondered about the consequence if we do not use the two fully-connected layers in the PruningNet for weight prediction but directly apply the proposed stochastic training and crop the same weight matrix for matching the input and output channels in the Pruned Network. We compare the performance between the PruningNet with and without weight prediction. We select the channel number with uniformly pruning each layer at a ratio ranging from [0.25, 1], and evaluate the accuracy with the weights generated by these two PruningNets. Figure 8 shows PruningNet without weight prediction achieves 10% lower accuracy. We further use the PruningNet without weight prediction to search for the Pruned MobileNet V1 with less than 45M FLOPs. The obtained network achieves only 55.3% top1 accuracy, 1.9% lower than the pruned network obtained with weight prediction. It is intuitive. For example, the weight matrix for a input channel width of 64 may not be optimal when the total input channels are increased to 128 with 64 more channels added behind. In that case, the weight prediction mechanism in meta learning is effective in de-correlating the weights for different pruned structures and thus achieves much higher accuracy for the PruningNet.
Conclusion
In this work, we have presented MetaPruning for channel pruning with following advantages: 1) it achieves much higher accuracy than the uniform pruning baselines as well as other state-of-the-art channel pruning methods, both traditional and AutoML-based; 2) it can flexibly optimize with respect to different constraints without introducing extra hyperparameters; 3) ResNet-like architecture can be effectively handled; 4) the whole pipeline is highly efficient.
Acknowledgement
The authors would like to acknowledge HKSAR RGC’s funding support under grant GRF-16203918, National Key R&D Program of China (No. 2017YFA0700800) and Beijing Academy of Artificial Intelligence (BAAI).