PSE-Net: Channel Pruning for Convolutional Neural Networks with Parallel-subnets Estimator

Shiguang Wang,Tao Xie,Haijun Liu,Xingcheng Zhang,Jian Cheng
2024-08-29
Abstract:Channel Pruning is one of the most widespread techniques used to compress deep neural networks while maintaining their performances. Currently, a typical pruning algorithm leverages neural architecture search to directly find networks with a configurable width, the key step of which is to identify representative subnet for various pruning ratios by training a supernet. However, current methods mainly follow a serial training strategy to optimize supernet, which is very time-consuming. In this work, we introduce PSE-Net, a novel parallel-subnets estimator for efficient channel pruning. Specifically, we propose a parallel-subnets training algorithm that simulate the forward-backward pass of multiple subnets by droping extraneous features on batch dimension, thus various subnets could be trained in one round. Our proposed algorithm facilitates the efficiency of supernet training and equips the network with the ability to interpolate the accuracy of unsampled subnets, enabling PSE-Net to effectively evaluate and rank the subnets. Over the trained supernet, we develop a prior-distributed-based sampling algorithm to boost the performance of classical evolutionary search. Such algorithm utilizes the prior information of supernet training phase to assist in the search of optimal subnets while tackling the challenge of discovering samples that satisfy resource constraints due to the long-tail distribution of network configuration. Extensive experiments demonstrate PSE-Net outperforms previous state-of-the-art channel pruning methods on the ImageNet dataset while retaining superior supernet training efficiency. For example, under 300M FLOPs constraint, our pruned MobileNetV2 achieves 75.2% Top-1 accuracy on ImageNet dataset, exceeding the original MobileNetV2 by 2.6 units while only cost 30%/16% times than BCNet/AutoAlim.
Computer Vision and Pattern Recognition
What problem does this paper attempt to address?
The main problem that this paper attempts to solve is to improve the efficiency and accuracy of channel pruning in convolutional neural networks (CNNs). Specifically, existing channel pruning methods usually adopt a serial training strategy when training a supernet, which results in a very time - consuming and inefficient training process. In addition, due to the long - tailed distribution characteristics of network configurations, finding samples that meet resource - constraint conditions is also a challenge. To address these issues, the paper proposes PSE - Net, a new parallel sub - net estimator, aiming to solve the above problems in the following ways: 1. **Propose PSE - Net**: PSE - Net is an efficient channel pruning method. Through the parallel - subnets training algorithm, it simulates the forward - backward propagation of multiple sub - nets in a single - round training, thereby significantly improving the training efficiency of the supernet. This method not only speeds up the training but also enables the network to interpolate the accuracy of unsampled sub - nets, so that it can more reliably evaluate the performance of sampled sub - nets. 2. **Introduce a prior - distributed - based sampling algorithm**: In order to effectively search for the optimal sub - net and solve the problem of finding samples that meet resource - constraint conditions, the paper proposes a prior - distributed - based sampling algorithm. This algorithm uses the loss function and computational cost (such as FLOPs) in the supernet training stage to efficiently estimate the prior distribution. By directly performing rejection sampling from the prior distribution, it significantly improves the efficiency and quality of sampling. 3. **Experimental verification**: The paper verifies the effectiveness of PSE - Net through extensive experiments. For example, on the ImageNet dataset, when the FLOPs are limited to 300M, the MobileNetV2 pruned by PSE - Net achieves a Top - 1 accuracy of 75.2%, which is 2.6 percentage points higher than the original MobileNetV2, and the time required to construct the supernet estimator is only 30%/16% of that of BCNet/AutoSlim. In conclusion, through the innovative parallel - subnets training algorithm and the prior - distributed - based sampling algorithm, PSE - Net effectively solves the problems of low training efficiency and difficulty in finding samples that meet resource - constraint conditions in existing channel pruning methods, providing a new solution for efficient and accurate channel pruning.