论文笔记EGE-UNet: an Efficient Group Enhanced UNet for skin lesion segmentation

EGE-UNet：一种高效的群算法用于皮肤病变分割的增强型UNet算法

        Transformer及其变体已广泛用于医学图像分割。然而，这些模型的大量参数和计算负载使它们不适合移动的应用。为了解决这个问题，提出了一个更有效的方法EGE-UNet。我们以一个轻量级的方式将一组多轴Hadamard产品注意力模块（GHPA）和一组聚合桥模块（GAB）进行组合搭建。GHPA对输入特征进行分组，并在不同的轴上执行Hadamard乘积注意机制（HPA），以从不同的角度提取病理信息。GAB有效地融合多尺度信息，通过分组低级别的功能，高级别的功能和由解码器在每个阶段产生的掩模。在ISIC 2017和ISIC 2018数据集上的综合实验表明，EGEUNet优于现有的最先进的方法。（与TransFuse相比，模型实现了更好的分割性能，同时将参数和计算成本分别降低了494倍和160倍）此外，这是第一个参数计数仅限于50 KB的模型。代码可在https://github.com/JCruan519/EGE-UNet上获得。

(a)和(b)分别显示了ISIC2017和ISIC2018数据集上的比较实验结果的可视化

        具体来说，EGE-UNet利用两个关键模块：组多轴Hadamard产品注意模块（GHPA）和组聚合桥模块（GAB）:

        1.一方面，最近基于ViT的模型已经显示出希望(由于多头自我注意机制MHSA）MHSA将输入分为多个头部，并计算每个头部的自我注意力，这使得模型能够从不同的角度获取信息，整合不同的知识，提高性能。尽管如此，MHSA的二次复杂度极大地增加了模型的大小。因此，我们提出了Hadamard产品注意力机制（HPA）的线性复杂度。HPA采用一个可学习的权重，并与输入进行哈达玛乘积运算以获得输出。随后，受MHSA中的多头模式的启发，我们提出了GHPA，它将输入分成不同的组，并在每个组中执行HPA。但是，值得注意的是，我们在不同的群体中在不同的轴上执行HPA，这有助于进一步从不同的角度获取信息。

        2.另一方面，对于GAB，由于医学图像中分割目标的大小和形状不一致，因此获得多尺度信息是必不可少的。因此，GAB集成了高层次和低层次的功能与不同大小的组聚合的基础上，并额外引入掩码信息，以协助特征融合。通过将上述两个模块与UNet相结合，我们提出了EGE-UNet，它以极低的参数和计算量实现了优异的分割性能。与以前的方法，只专注于提高性能，我们的模型还优先考虑在现实世界的环境中的可用性。

        总之，我们的贡献有三个方面：

        (1)提出了GHPA和GAB算法，前者有效地获取和融合多视角信息，后者接受不同尺度的特征，并结合辅助模板进行有效的多尺度特征融合。

        (2)我们提出了EGEUNet，一个非常轻量级的模型，专为皮肤病变分割。

        (3)我们进行了广泛的实验，这表明我们的方法在实现国家的最先进的性能，显着降低资源需求的有效性。

        EGE-UNet如图所示，其构建在由对称编码器-解码器部分组成的U形架构上。我们以编码器部分为例。编码器由六个级组成，每个级具有{8，16，24，32，48，64}的信道号。虽然前三个阶段采用核大小为3的普通卷积，最后三个阶段利用所提出的GHPA从不同的角度提取表示信息。与UNet中的简单跳过连接相反，EGE-UNet针对编码器和解码器之间的每个阶段并入GAB。此外，我们的模型利用深度监督来生成不同尺度的掩码预测，这些掩码预测用于损失函数并作为GAB的输入之一。通过这些先进模块的集成，EGE-UNet显着降低了参数和计算负荷，同时提高了分割性能相比，以前的方法。

        组多轴Hadamard产品注意模块。为了克服MHSA带来的二次复杂度问题，我们提出了线性复杂度的HPA。给定输入x和随机初始化的可学习张量p，首先利用双线性插值来调整p的大小以匹配x的大小。然后，我们在p上使用深度可分离卷积（DW），然后在x和p之间进行hadamard乘积运算以获得输出。然而，单独利用简单的HPA不足以从多个视角提取信息，导致不令人满意的结果。受MHSA中的多头模式的启发，我们引入了基于HPA的GHPA，算法如下所示。我们将输入沿着通道维度平均分为四组，并分别对前三组的高度-宽度、通道-高度和通道-宽度轴执行HPA。对于最后一组，我们只在特征图上使用DW。最后，我们沿着渠道维度串联的四个组，并应用另一个数据仓库从不同的角度整合信息。注意，DW中使用的所有内核大小都是3。

         组聚合桥模块。多尺度信息的获取被认为是密集预测任务的关键，例如医学图像分割。因此，如下图所示，我们引入GAB，它需要三个输入：低级特征、高级特征和掩模。首先，采用深度可分离卷积（DW）和双线性插值来调整高层特征的大小，以匹配底层特征的大小。其次，我们沿着通道维度将两个特征映射分成四组，并将一组从低层次的特征与一个从高层次的特征连接，以获得四组融合的特征。对于每组融合特征，掩码被级联。接下来，将核大小为3的扩张卷积和不同的扩张率{1，2，5，7}应用于不同的组，以便提取不同尺度的信息。最后，四个组沿着通道维度级联，随后应用具有1的核大小的普通卷积，以实现不同尺度的特征之间的交互。

 ISIC2017和ISIC2018数据集的比较实验结果

 ISIC2017数据集上的消融研究。(a)两个模块上的宏观消融。(b)GHPA上的微消融(c)GAB的微消融

 

 

        在本文中，我们提出了两个先进的模块。我们的GHPA使用了一种新的HPA机制，以简化的二次复杂性的自我关注的线性复杂性。它还利用分组从不同的角度充分捕获信息。我们的GAB融合了低级别和高级别的功能，并引入了一个掩模来集成多尺度信息。基于这些模块，我们提出了EGE-UNet皮肤病变分割任务。实验结果表明，我们的方法在实现国家的最先进的性能显着降低资源需求的有效性。我们希望我们的工作可以启发医学影像界对轻量级模型的进一步研究。