EGE-UNet: an Efficient Group Enhanced UNet for skin lesion segmentation
This work addresses the need for efficient models for skin lesion segmentation in mobile health applications, offering a significant reduction in resource usage while maintaining performance.
The paper tackled the problem of high parameter and computational costs in transformer-based models for medical image segmentation, proposing EGE-UNet, which achieved superior segmentation performance on ISIC2017 and ISIC2018 datasets while reducing parameters by 494x and computation by 160x compared to TransFuse, with a parameter count limited to 50KB.
Transformer and its variants have been widely used for medical image segmentation. However, the large number of parameter and computational load of these models make them unsuitable for mobile health applications. To address this issue, we propose a more efficient approach, the Efficient Group Enhanced UNet (EGE-UNet). We incorporate a Group multi-axis Hadamard Product Attention module (GHPA) and a Group Aggregation Bridge module (GAB) in a lightweight manner. The GHPA groups input features and performs Hadamard Product Attention mechanism (HPA) on different axes to extract pathological information from diverse perspectives. The GAB effectively fuses multi-scale information by grouping low-level features, high-level features, and a mask generated by the decoder at each stage. Comprehensive experiments on the ISIC2017 and ISIC2018 datasets demonstrate that EGE-UNet outperforms existing state-of-the-art methods. In short, compared to the TransFuse, our model achieves superior segmentation performance while reducing parameter and computation costs by 494x and 160x, respectively. Moreover, to our best knowledge, this is the first model with a parameter count limited to just 50KB. Our code is available at https://github.com/JCruan519/EGE-UNet.