Dynamically pruning segformer for efficient semantic segmentation
This work addresses the problem of deploying efficient semantic segmentation models on edge devices, representing an incremental improvement through pruning and distillation.
The paper tackles the high computational cost of SegFormer for semantic segmentation on edge devices by proposing a dynamic gated linear layer that prunes uninformative neurons per input and uses two-stage knowledge distillation, achieving 36.9% mIoU with 3.3G FLOPs on ADE20K, saving over 60% computation with only a 0.5% mIoU drop.
As one of the successful Transformer-based models in computer vision tasks, SegFormer demonstrates superior performance in semantic segmentation. Nevertheless, the high computational cost greatly challenges the deployment of SegFormer on edge devices. In this paper, we seek to design a lightweight SegFormer for efficient semantic segmentation. Based on the observation that neurons in SegFormer layers exhibit large variances across different images, we propose a dynamic gated linear layer, which prunes the most uninformative set of neurons based on the input instance. To improve the dynamically pruned SegFormer, we also introduce two-stage knowledge distillation to transfer the knowledge within the original teacher to the pruned student network. Experimental results show that our method can significantly reduce the computation overhead of SegFormer without an apparent performance drop. For instance, we can achieve 36.9% mIoU with only 3.3G FLOPs on ADE20K, saving more than 60% computation with the drop of only 0.5% in mIoU