Pruning Convolutional Neural Networks: Decide the Plan, Prune, Then Finetune

Model pruning is a common technique for accelerating and compressing neural networks. It dates back at least to LeCun’s Optimal Brain Damage in 1989. The idea is to remove the “unimportant” parts of a network, and the whole process is usually split into three steps:

1. First, based on the overall objective, determine the model’s pruning plan—that is, the pruning rate for each layer (sometimes also called the compression rate or sparsity rate).
2. Then prune the corresponding layers at the given pruning rates.
3. Finally, finetune the network to recover the accuracy lost during pruning.

Steps 2 and 3 are often iterated many times, pruning only a portion each time and progressively approaching the overall objective. The above is the pruning workflow with a pretrained model; when no pretrained model is available, you can also prune while training. Let’s go through these in turn.

Determining the Pruning Plan

Before pruning there is usually an overall objective (parameter count, peak memory, FLOPs, latency, etc.). The purpose of determining the plan is to allocate each layer’s pruning rate reasonably while meeting the objective, so as to minimize accuracy loss. There are two kinds of approaches: manual design and automatic design.

Manual Design

The most naive approach is to rely on intuition. For example, in a CNN the earlier convolutional layers are responsible for extracting the most primitive image features, and the deeper layers all depend on the outputs of the shallow ones, so the shallow layers should not be pruned too aggressively. Likewise, convolutional layers that perform downsampling (e.g. $stride=2$) produce smaller output feature maps and need to retain more filters to maintain enough channels, lest too much feature information be lost.

A more principled approach is pruning sensitivity analysis. For each layer individually, prune different proportions of its parameters and measure how much accuracy drops: if a layer loses a lot of accuracy after pruning only a few parameters, it is very “sensitive,” and you should keep more parameters in that layer when pruning the whole network. The figure below shows the sensitivity analysis for the layers of AlexNet, and the results match intuition: the early convolutional layer with only 3 input channels is the most sensitive to pruning; convolutional layers have fewer parameters than fully connected layers and bear the main burden of feature extraction, so they are also more sensitive.

Automatic Design

One automatic approach is to set a global importance threshold and prune all weights (unstructured) or filters (structured) below it—the more low-importance ones a layer has, the more aggressively it is naturally pruned, with no need to manually specify each layer’s ratio.

Another is search. It closely resembles neural architecture search (NAS): both aim to find the most accurate structure under resource constraints, except that automatic pruning usually focuses on only one dimension, the number of filters. Its search space is smaller and it can start directly from a pretrained model, so it is considerably simpler than NAS.

AMC uses DDPG from reinforcement learning to search the compression rate layer by layer. From a reinforcement learning perspective: the environment is model pruning, the state is the current pruning progress, the action is the pruning rate for this layer, and the reward is the optimization objective.

What is worth noting here is the state and the reward. The state includes the index of the layer being pruned, the convolution weight dimensions, the layer’s FLOPs, the FLOPs already pruned, the FLOPs still remaining to be pruned, and the previous layer’s pruning rate—because pruning proceeds layer by layer, the “remaining FLOPs to prune” must be added to the state to constrain the lower bound of compression for the current layer, so that the overall objective is still met after pruning. The reward includes not only the error rate but also a FLOPs penalty term. As in NAS, the most time-consuming part of automatic pruning is model validation (the segment from action to reward); AMC uses a small amount of calibration data and least squares to recalibrate the weights of the pruned layer, avoiding the time-consuming finetune at the cost of limited accuracy gain.

NetAdapt, on the other hand, is iterative: in each round it prunes different layers to generate several candidate proposals, briefly finetunes them, and selects the most accurate one to advance to the next round, until the latency meets the requirements of a specific piece of hardware.

Both MetaPruning and EagleEye devised ways to “speed up candidate validation”: MetaPruning uses a PruningNet to directly generate the pruned weights, sparing a great deal of finetuning; EagleEye, like SPOS in NAS, performs just one forward pass on the calibration data after pruning to recalibrate BN’s mean and variance, likewise bypassing the time cost of finetuning.

Pruning the Model

Once the pruning rate for each layer is fixed, it is time to actually prune. There are two routes: unstructured pruning, which sparsifies the weights, and structured pruning, which removes whole filters/modules.

Unstructured Pruning

Unstructured pruning sparsifies the weights: it zeroes out the unimportant ones (equivalent to removing connections between neurons). This saves space when storing the network, but to actually accelerate computation it relies on special software implementations (such as high-performance sparse matrix multiplication) or hardware support (such as the NVIDIA A100 Tensor Core GPU).

A classic method is Learning both Weights and Connections: during training it trains not only the weights but also “trains” the connections, pruning through learning, in three steps—

1. Train the network in the normal way (usually adding L1 or L2 penalty to the weights during the process);
2. Remove the weights below a certain threshold (zero them out, yielding sparse weights);
3. Retrain to adjust the remaining weights (i.e. finetune).

These three steps are usually iterated many times, pruning only a portion of the connections each time, progressively compressing the network to a very sparse state. Experiments show that retraining (finetuning) after pruning, and using iterative rather than one-shot pruning, are both crucial for reducing accuracy loss.

Another very interesting line of work is the Lottery Ticket Hypothesis: within a randomly initialized, unpruned dense network, there always exists a subnetwork that—keeping the same initialization parameters—can reach the same accuracy as the original network after the same or fewer training steps. This subnetwork is called the winning ticket, and it is found as follows:

1. Randomly initialize the network $f(x;\theta_0)$;
2. Train for $j$ iterations to obtain parameters $\theta_j$;
3. Prune the smallest $p\%$ of parameters in $\theta_j$, obtaining a mask $m$;
4. Reset the remaining parameters back to $\theta_0$, obtaining the winning ticket $f(x;m\odot\theta_0)$.

As above, the iterative approach is again better: to achieve a $p\%$ pruning rate, you can split it into $n$ rounds, pruning $p^{1/n}\%$ each time. Experiments show that the iterative approach outperforms one-shot, and keeping the initial parameters ($m\odot\theta_0$) outperforms re-randomizing the initialization.

Structured Pruning

Structured pruning achieves real compression and acceleration without special software or hardware, so it is more commonly used. It usually removes entire filters from convolutional layers, and the problem splits into two halves: how many filters should be pruned in a given layer? And by what criterion should the filters to prune be chosen? The former has already been covered above (sensitivity analysis, global importance ranking, automated methods), so here we focus on the latter—the importance criteria for filters. Depending on whether they rely on data, they fall into two categories.

Data-free methods require only the pretrained weights and rank filters based on the convolution weights themselves. L1 pruning is the simplest: it directly uses the L1 norm of the filter weights as importance. With the $i$-th filter weight being $\mathcal{F}_i$, its importance is $\lVert\mathcal{F}_i\rVert_1$, and filters with a small norm are deemed unimportant; L2 pruning is analogous, using $\lVert\mathcal{F}_i\rVert_2$ instead. The two actually differ little.

FPGM instead uses the geometric median of the filter weights: if a layer has $N$ filters, the importance of the $i$-th is $\sum_{j\in[1,N]}\lVert\mathcal{F}_i-\mathcal{F}_j\rVert_2$. Unlike L1/L2, FPGM keeps both large-norm and small-norm filters—what it prunes are those that are “too similar to others, i.e. redundant.”

Network Slimming judges via the BN layer. A convolution is generally followed by BN for normalization and an affine transformation; if a filter’s output is ultimately multiplied by a $\gamma$ (channel scaling factor) close to 0, then its contribution to the subsequent layers is also close to 0, and it can be considered unimportant. To make $\gamma$ sufficiently sparse, an L1 regularization is usually added to $\gamma$ during training.

The L1/L2, FPGM, and Network Slimming methods above all look only at the pretrained weights. Data-dependent methods, on the other hand, require calibration data and rank filters based on those data’s outputs in the convolutional layers; because the importance criterion is highly coupled with the target data, the accuracy loss is often smaller at the same compression ratio. Network Trimming proposes using APoZ (Average Percentage of Zeros)—after the convolution output passes through BN and ReLU, if a filter’s activations contain many zeros, it contributes little and can be safely pruned; Molchanov et al. also propose using the mean or standard deviation of the activations as the criterion. HRank notices that each filter corresponds to one output feature map (a matrix) and uses that matrix’s rank as importance: the higher the rank, the more information it carries, and the more it should be kept.

A caveat: the calibration data for data-dependent methods must be randomly sampled from the training set to keep the same distribution; otherwise the accuracy loss may be even larger than with data-free methods.

One more point: structured pruning changes the network’s original structure, so you must ensure the input and output dimensions of each layer still line up after pruning, especially for structures with residual connections. Take a ResNet residual block as an example: if the main path and the shortcut have mismatched dimensions, they cannot be added. A common practice is to make the number of filters in the last convolutional layer of a block equal to the number of channels of the input feature map, thereby ensuring the channel counts match when adding.

Pruning During Training

Everything above has been about pruning a trained model. If there is no pretrained model, then besides “train first, then prune” and “directly use NAS to search a more compact structure,” you can also prune while training.

Slimmable Networks uses a special training method that lets a single network run at different widths. During training, each batch is forwarded separately through subnetworks of multiple widths, computing the loss and backpropagating gradients, and finally these gradients are accumulated before updating the parameters. Two key points: one is using switchable BN—each width uses its own independent set of BN parameters ($\mu,\sigma,\gamma,\beta$); the other is aggregating the gradients of different widths before updating.

Algorithm  Training slimmable neural network M.

Require: Define switchable width list for slimmable network M,
         for example, [0.25, 0.5, 0.75, 1.0]x.
 1: Initialize shared convolutions and fully-connected layers for slimmable network M.
 2: Initialize independent batch normalization parameters for each width in switchable width list.
 3: for i = 1, ..., n_iters do
 4:   Get next mini-batch of data x and label y.
 5:   Clear gradients of weights, optimizer.zero_grad().
 6:   for width in switchable width list do
 7:     Switch the batch normalization parameters of current width on network M.
 8:     Execute sub-network at current width, y_hat = M'(x).
 9:     Compute loss, loss = criterion(y_hat, y).
10:     Compute gradients, loss.backward().
11:   end for
12:   Update weights, optimizer.step().
13: end for

Each width is forwarded/backpropagated separately, and the gradients are accumulated before a unified update.

Universally Slimmable Networks builds on this by proposing the sandwich rule and inplace distillation, further improving performance.

Algorithm 1  Training universally slimmable network M.

Require: Define width range, for example, [0.25, 1.0]x.
Require: Define n as number of sampled widths per training iteration, for example, n = 4.
 1: Initialize training settings of shared network M.
 2: for t = 1, ..., T_iters do
 3:   Get next mini-batch of data x and label y.
 4:   Clear gradients, optimizer.zero_grad().
 5:   Execute full-network, y' = M(x).
 6:   Compute loss, loss = criterion(y', y).
 7:   Accumulate gradients, loss.backward().
 8:   Stop gradients of y' as label, y' = y'.detach().
 9:   Randomly sample (n-2) widths, as width samples.
10:   Add smallest width to width samples.
11:   for width in width samples do
12:     Execute sub-network at width, y_hat = M'(x).
13:     Compute loss, loss = criterion(y_hat, y').
14:     Accumulate gradients, loss.backward().
15:   end for
16:   Update weights, optimizer.step().
17: end for

Universally Slimmable: sandwich rule + inplace distillation.

AutoSlim’s automatic pruning idea is almost identical to NetAdapt’s—both prune greedily and iteratively until FLOPs/latency meet the target; the biggest difference is that AutoSlim starts from a trained slimmable network, whereas NetAdapt starts from an ordinary pretrained model.

References

• LeCun, Yann, et al. Optimal Brain Damage. NeurIPS, 1989.
• Han, Song, et al. Learning both Weights and Connections for Efficient Neural Networks. NeurIPS, 2015.
• Frankle, Jonathan, Carbin, Michael. The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks. arXiv:1803.03635, 2018.
• Li, Hao, et al. Pruning Filters for Efficient ConvNets. arXiv:1608.08710, 2016.
• Liu, Zhuang, et al. Learning Efficient Convolutional Networks through Network Slimming. ICCV, 2017.
• He, Yang, et al. Filter Pruning via Geometric Median for Deep Convolutional Neural Networks Acceleration (FPGM). CVPR, 2019.
• Hu, Hengyuan, et al. Network Trimming: A Data-Driven Neuron Pruning Approach. arXiv:1607.03250, 2016.
• Molchanov, Pavlo, et al. Pruning Convolutional Neural Networks for Resource Efficient Inference. arXiv:1611.06440, 2016.
• Lin, Mingbao, et al. HRank: Filter Pruning using High-Rank Feature Map. CVPR, 2020.
• He, Yihui, et al. AMC: AutoML for Model Compression and Acceleration on Mobile Devices. ECCV, 2018.
• Yang, Tien-Ju, et al. NetAdapt: Platform-Aware Neural Network Adaptation for Mobile Applications. ECCV, 2018.
• Liu, Zechun, et al. MetaPruning: Meta Learning for Automatic Neural Network Channel Pruning. ICCV, 2019.
• Li, Bailin, et al. EagleEye: Fast Sub-net Evaluation for Efficient Neural Network Pruning. arXiv:2007.02491, 2020.
• Guo, Zichao, et al. Single Path One-Shot Neural Architecture Search with Uniform Sampling (SPOS). ECCV, 2020.
• Yu, Jiahui, et al. Slimmable Neural Networks. arXiv:1812.08928, 2018.
• Yu, Jiahui, Huang, Thomas S. Universally Slimmable Networks and Improved Training Techniques. ICCV, 2019.
• Yu, Jiahui, Huang, Thomas. AutoSlim: Towards One-Shot Architecture Search for Channel Numbers. arXiv:1903.11728, 2019.
• Ioffe, Sergey, Szegedy, Christian. Batch Normalization. arXiv:1502.03167, 2015.
• He, Kaiming, et al. Deep Residual Learning for Image Recognition. CVPR, 2016.