【论文阅读】EGE-UNet: anEfficient Group Enhanced UNet for skin lesion segmentation

摘要

1.解决的问题:解决Transformer和它的变体的大量参数和计算内存使他们不适合再移动端应用的问题。
2. 提出的解决方法:EGE-UNet。以轻量级的方式整合了GHPA和GAB模块。GHPA模块对输入的特征进行分组,在不同的轴线方向上执行HPA(一种注意力模块),从不同角度提取病理信息。GAB通过把低级特征、高级特征和解码器在不同的阶段生成的标签进行分组来有效地融合不同尺度的信息。
3. 实验结果:在ISIC2017和ISIC2018数据集上的实验证明EGE-UNet超过现有的最先进的模型。据我们所知这是第一个模型参数限制在50KB的模型。
4. code :https://github.com/JCruan519/EGE-UNet

方法

EGE-UNet的结构如下图所示,是在包含对称的编码器解码器部分的U形架构中建立。编码器由6个阶段,通道数为{8,16,24,32,48,64}。前三个阶段使用卷积核大小为3的普通卷积,后三个使用提出的GHPA提取不同视角的信息。与U-Net简单的跳跃连接不同的是,EGE-UNet在每个阶段的编码器和解码器之间使用GAB。进一步说,我们的模型利用深监督去生成不同尺度的预测掩码。这些预测掩码在损失函数中使用,也作为GAB的一个输入。通过融合这些提出的模块,EGE-UNet显著的降低了参数量和计算负荷,同时相比于先前的方法分割性能提升。
在这里插入图片描述

GHPA module

为了克服MHSA的二次计算复杂度的问题,我们提出了线性复杂度的HPA。给定一个输入x和一个随机初始化的可学习的张量p,先使用双线性插值去改变p的形状去匹配x的形状。然后我们使用对p深度可分离卷积DW,随后对x和p使用逐元素乘积操作得到输出。然而仅仅利用简单的HPA对于从多种视角提取特征是不足的,造成不令人满意的结果。受到MHSA中的多头模块的启发,我们基于HPA提出了GHPA,如下面Algorithm 1所示。我们把输入在通道维度上平均分成四组,在高-宽,通道-高,通道-宽的轴线上对前3组数据使用HPA。对于最后一组我们仅仅对特征图进行DW操作。最后,我们在通道维度结合了这四组,然后采用其他的DW操作去整合不同视角的信息。注意所有的DW中的卷积大小都是3。
在这里插入图片描述
看上面的伪代码,其中x1-x4是输入X在通道上平均分成的四组。然后,x1是xy平面上与Pxy进行点乘,xy是高-宽在最后两个维度上,不需要变换顺序相乘,x4只是进行简单的DW变换。其中x2,x3是在通道-高和通道-宽维度上与Pzx和Pzy进行点乘,需要先把通道-高和通道-宽分别放到最后两个维度上,再进行点乘。输出是把x1-x4在通道维度进行结合,然后进行了层归一化和深度可分离卷积操作得到输出。

GAB

在这里插入图片描述
对于密集预测任务来说关键是获得不同尺度的信息,比如:医学图像分割任务。因此,如上图所示,我们提出了GAB,它有三个输入:低级特征、高级特征和标签。首先使用DW和双线性插值变换去调整高级特征图的大小来匹配低级特征图的大小。其次,我们把两种特征图都在通道维度上分割成4组,把低级特征上的每一组和高级特征上的每一组结合起来获得4组融合特征。对于每一组融合的特征,标签是串联到上面。然后,卷积核大小为3有不同膨胀率{1,2,5,7}的膨胀卷积在不同的组使用,为了在不同的尺度上提取特征。最后,这四组在通道维度结合,之后使用卷积核大小为1的卷积保证不同尺度的特征相互联系。

Loss function

在这个研究中,因为不同的GAB需要不同的尺度的标签信息,使用深度监督去计算不同阶段的损失函数,以生成更加准确的标签信息。我们的损失函数可以以下面两个式子来表示。

在这里插入图片描述
Bce和Dice表示BCE Loss和Dice Loss。λi是不同阶段的权重。这个文章i=0-5λ设置的是1,0.5,0.4,0.3,0.2,0.1。

实验

数据集和实验

为了评估我们模型的有效性,我们选择了ISIC2017和ISIC2018,包含2150和2694张皮肤镜学图片,设置的训练集和测试集比例是7:3。
图片都归一化改变大小到256×256。我们设置了各种各样的数据增强方法,包括水平翻转、垂直翻转和随机旋转。使用了AdamW作为优化器。初始化学习率为0.001。学习率更新策略是余弦退火策略,最大迭代策略是50,最小的学习率是1e-5。训练300轮,批次大小为8。
为了评估我们的方法,我们使用mIoU和DSC的平均插值作为评价标准,我们进行了5次并且公布了每一个数据集平均值和标准差。

对比实验结果

在这里插入图片描述
在这里插入图片描述

消融实验结果

在这里插入图片描述

结论及未来工作

文章提出了两个先进的模块。我们GHPA使用一个新颖的HPA机制简化自注意力中的二次复杂度到线性复杂度。它也利用不同的分组去充分利用不同角度的信息。GAB融合了低级特征和高级特征,引入了标签去整合多尺度信息。基于这些模块,我们提出了EGE-UNet来解决皮肤损害分割模型。实验结果证明了我们的方法在显著降低资源需求的情况下实现最先进性能的有效性。我们希望我们的工作能够启发医学图像社区对轻量级模型的进一步研究。
考虑到局限性和未来工作,一方面我们主要聚焦于怎样显著减少参数和计算复杂度,同时提高分割性能。因此我们准备在真实世界环境中配置EGE-UNet。另一方面,我们扩展到其他分割任务中。

  • 1
    点赞
  • 12
    收藏
    觉得还不错? 一键收藏
  • 3
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值