Novel channel attention-based filter pruning methods for low-complexity semantic segmentation models - Md. Bipul Hossain; Na Gong; Mohamed Shaban (Machine Learning with Applications)
Semantic segmentation is the area of classifying each pixel in an image using a deep learning model. Examples of widely used semantic segmentation models are the U-Net and DeeplabV3+ models. While the aforementioned models have been deemed very successful in segmenting medical targets including organs and diseases in high resolution images, the computational complexity represents a burden for the real-time application of the algorithms or the deployment of the models on resource-constrained platforms. Until recently, few methods have been introduced for optimizing or pruning of the parameters of the semantic segmentation models. In this paper, we propose two novel channel attention-based filter pruning techniques (i.e., Sub-Sampling Channel Attention (SACA) and Self-Attention Based Attention (SBCA)) in order to reduce the complexity of the semantic segmentation models while maintaining high performance with respect to the benchmark models. This is realized by recognizing the contextual importance of the feature maps in each layer of the models and the significance of each filter to the final model performance. The proposed optimization methods have been validated on the U-Net and DeeplabV3+ models using both lung and skin lesion datasets. The proposed approaches achieved a pruned model performance (i.e., dice coefficient) of up to 96%, as well as an extensively reduced complexity (i.e., percentage of remaining parameters down to 1.1%, model size down to 1.22 MB and number of GFLOPS down to 1.06), outperforming the benchmark magnitude based (i.e., l1-norm, and l2-norm) and the attention-based (i.e., SE, ECA, and CBAM CA) filter pruning methods.