EMCAD- Efficient Multi-scale Convolutional Attention Decoding for Medical Image Segmentation
论文地址:arxiv
摘要
背景
一个高效且有效的解码机制在医学图像分割中至关重要,而当前的解码机通常会有着高计算成本。
创新点
提出了 EMCAD,是一种新的高效多尺度卷积注意力解码器,旨在优化性能与计算效率。
EMCAD 利用了独特的多尺度深度卷积块,通过多尺度卷积显著增强特征图。EMCAD 还用了通道,空间与分组门控注意力机制,这些机制在捕捉复杂的空间关系并聚集于显著区域方面非常有效。
通过使用分组与深度卷积,EMCAD 非常高效且具有良好的扩展性。
正文
贡献
- 新型高效多尺度卷积解码器:引入了高效的多尺度级联全卷积注意力解码器(EMCAD)来实现 2 D 医学图像分割。它可以接收视觉编码器的多阶段特征,并逐步增强多尺度和多分辨率的表示。
- 高效多尺度卷积注意力模块:引入了 MSCAM,这是一种新的高效多尺度卷积注意力模块。可以在多个尺度上执行深度卷积,优化了视觉编码器生成的特征图,并通过抑制不相关区域来捕捉多尺度显著特征。
- 大核分组注意力门:引入了一种新的分组注意力门,以融合优化后的特征与来自跑跃连接(skip connections)的特征。
- 性能的提升。
当前技术的不足
CNN:由于其固有的局部感受野,通常缺乏捕捉图像中长距离依赖关系的能力。
ViTs:解决了 CNNs 在捕捉像素长距离依赖关系这方面的局限性,但是它们在捕捉像素之间的局部关系这方面仍有困难。
高效多尺度卷积注意力解码(EMCAD)
EMCAD 的架构如 (b)
所示:
EMCAD 由以下模块组成:
- 高效多尺度卷积注意力模块(MSCAMs)组成,以稳健地增强特征图;
- 大核分组注意力门(LGAGs)通过门控注意力机制与跳跃连接结合来细化特征图;
- 高效上卷积块(EUCBs)用于上采样并随后增强特征图;
- 以及分割头(SHs)来生成分割输出。
使用四个 MSCAMs 来细化从编码器的四个阶段提取的金字塔特征(即图 2 中的 X 1、X 2、X 3、X 4)。在每个 MSCAM 之后,使用一个 SH 来生成该阶段的分割图。随后,使用 EUCBs 对细化后的特征图进行上采样,并将它们添加到相应 LGAGs 的输出中。最后,将四个不同的分割图相加以生成最终的分割输出。
大核分组注意力门 (LGAG)
大核分组注意力门,将特征图与注意力系数逐步结合起来。通过网络学习,使相关特征得到更高的激活,而无关特征则被抑制。这一过程利用从高层特征中提取的门控信号来控制网络不同阶段的信息流,从而提高了医学图像分割的精确度。
在本次论文中,$q_{att}(.)$ 函数通过分别应用 3*3
组卷积 $GC_{g}$ 与 $GC_{x}$ 来处理 g 和 x。然后使用批量归一化 ($BN(.)$)来对这些卷积特征进行归一化处理,并通过逐元素加法合并。合并后的特征图通过 ReLU
层($R(.)$)激活。之后应用一个 1*1
卷积($C(.)$)后接 $BN(.)$ 层,来得到单通道特征图。之后将得到的单通道特征图通过 sigmoid
激活函数($\sigma(.)$)生成注意力系数。将转移后的输出用于通过逐元素乘法绽放输入特征 x,生成注意力门控特征 $LGAG(g,x)$,以上内容可通过以下公式表示:
$$
q_{att}(g,x)=R(BN(GC_{g}(g)+BN(GC_{x}(x)))))
$$
$$
LGAG(g,x)=x \odot \sigma(BN(C(q_{att}(g,x))))
$$
这种方法的优势在于,通过使用 3 × 3 核分组卷积,LGAG 能够在保持较低计算成本的同时捕获更大的空间上下文信息,从而提高医学图像分割的精度。
多尺度卷积注意力模块 (MSCAM)
MSCAM 由一个通道注意力块 (CAB (·)) 来强调相关通道,一个空间注意力块 (SAB (·)) 来捕捉局部上下文信息,以及一个高效的多尺度卷积块 (MSCB (.)) 来增强特征图并保留上下文关系。
该模块结构为上图的 (d)
多尺度卷积块(Multi-scale Convolution Block, MSCB)
该模块用于增强级联扩展路径生成的特征。该模块遵循 MobileNetV2
的倒置残差块(IRB)的设计,不过与之不同的是:
- 该模块在多个尺度上执行深度卷积(DWConv)
- 使用通道混洗(channel shuffle)在组间混洗通道
首先使用 1*1
卷积层($PWC(.)$)扩展通道数(扩展因子为 2),然后使用批量化层($BN(.)$)与 ReLU 6 激活层($R6(.)$),之后使用多尺度深度卷积 ($MSDC(.)$)来捕捉多尺度和多分辨率上下文。由于深度卷积层忽略了通道之间的关系,所以要使用通道混洗操作来结合通道之间的关系。之后再使用另一个逐点卷积($PWC2(.)$)以及一个 BNn 层来将通道数变回原来的数目,这同时也可以编码了通道之间的依赖关系。
以上过程可以用以下公式表示:
$$
MSCB(x) = BN(PWC_2(CS(MSDC(R6(BN(PWC_1(x)))))))
$$
其中,不同核大小的并行 MSDC,可以用以下公式表示
$$
MSDC(x) = \sum {ks \in KS} DWCB{ks}(x)
$$
$$
DWCBks(x) = R6(BN(DW Cks(x))).
$$
此外,$MSDC(.)$ 使用递归更新的输入 x,输入的 x 与前一个 $DWCBks(.)$ 残差连接可以实现更好的正则化,公式如下:
$$
x=x+DWCB_{ks}(x)
$$
通道注意力块(Channel Attention Block, CAB)
我们使用通道注意力块来为每个通道分配不同的重要性,从而强调更相关的特征,同时抑制不太有用的特征。
在 CAB 中,首先对空间维度(即高度与宽度)应用自适应最大池化($P_{m}(.)$)与自适应平均池化($P_{a}(.)$),来提取每个通道整个特征图中最显著的特征。之后,对于每个池化后的特征图,分别使用点卷积($C_{1}(.)$)将通道数减少到原来的 $\frac{1}{16}$,然后使用 ReLU 激活函数($R$)。之后使用另一个点卷积($C_{2}(.)$),恢复原始的通道数。将两个恢复后的特征图相加,并应用 sigmoid
($\sigma$)激活函数来估计注意力权重。最后使用 Hadamard
积 ($\circledast$),将这些权重与输入 x 结合。
以上公式可以如下表示:
$$
SAB(x) = \sigma (LKC([Ch_{max}(x), Ch_{avg}(x)])) \circledast x
$$
高效上卷积模块(Efficient up-convolution block,EUCB)
该模块的作用是对当前阶段的特征图进行上采样,来匹配来自下一个跑跃连接的特征图的维度与分辨率。
EUCB 先使用一个上采样操作($Up(.)$)将特征图按给因子 2 进行上采样。然后,通过应用一个 3*3
深度卷积($DWC(.)$),接着是批量归一化($BN(.)$)和 ReLU 激活函数($ReLU(.)$),来增强上采样后的特征图。最后,使用一个 1*1
卷积($C_{1*1}(.)$)来减少通道数,来匹配下一阶段的特征图。
以上过程可以用以下公式表示:
$$
EUCB(x) = C_{1\times 1}(ReLU(BN(DWC(Up(x)))))
$$
通过使用深度卷积,可以使得运行速度很快。
分割头 (Segmentation head, SH)
使用分割头从解码器的四个阶段的精细特征图中生成分割输出。
SH 层对具有 $x_{i}$ 通道数( $x_{i}$ 是第 $i$ 阶段特征图的通道数)的精细特征图应用一个 1×1 卷积操作 Conv 1×1 (·),并为多类别分割生成通道数等于目标数据集中类别数的输出,而对于二分类分割则生成 1 个通道的输出。
以上过程可由以下公式描述:
$$
SH(x)=Conv_{1*1}(x)
$$
整体架构
该解码器是可以自适应的,并且可以无缝兼容其他分层骨干网络。论文中将 EMCAD 解码器与 PVTv 2 的微型与标准网络集成在一起。
通过使用 PVTv2-B0
和 PVTv2-B2
作为编码器,开发了 PVT-EMCAD-B0
和 PVT-EMCAD-B2
两种架构。在这些架构中,首先从 PVTv2
网络的四层中提取特征(X 1, X 2, X 3, X 4),然后将这些特征输入到 EMCAD 解码器中。具体来说,X 4 在向上采样路径中使用,而 X 3, X 2, X 1 则在跳跃连接中使用。
EMCAD 解码器处理这些特征后,会生成四个分割图,每个分割图对应编码器网络的一个阶段。
多阶段损失与输出聚合
模型使用了四个分割头(segmentation heads),每个头在不同的阶段产生一个预测图(prediction map)。为了提高分割的准确性,采用了多阶段的损失聚合策略,即 MUTATION 方法,该方法计算所有可能的预测组合的损失,并将这些损失相加,以最小化总的组合损失。对于二值分割任务,还引入了一个额外的损失项,以进一步优化模型。最后,模型的输出是通过对最后一个阶段的预测图应用 Sigmoid 或 Softmax 函数来得到的,这取决于任务是二值分割还是多类分割。
模型评估
二值医学图像分割结果
可以看到,达到了最高的平均 DICE 分数。
息肉分割
从表 1 中可以看到:该模型( PVT-EMCAD-B2
)在五个息肉分割数据集上超越了所有 SOTA 方法。
皮肤病变分割
可以从表 1 中看到 PVT-EMCAD-B2
在 ISIC 17
和 ISIC 18
皮肤病变分割数据集上效果显著。
细胞分割
使用 DSB 18进行细胞核分割,使用 EM 进行细胞结构分割,从表 1 中显示,效果依然不错。
乳腺癌分割
在超声图像的 BUSI 数据集上进行了乳腺癌分割实验。效果不错。
腹部器官分割结果
PVT-EMCAD-B2
在 Synapse 多器官数据集上的腹部器官分割表现出色,同时,EMCAD 解码器提升了单个器官分割,显著超越了 SOTA 方法在八个器官中的六个。
心脏器官分割结果
表明,计算量更少,DICE 分数更好。
消融实验
从仅使用编码器开始,然后添加不同的模块,如级联结构、LGAG 和 MSCAM,得出以下表:
可以看到:解码器的级联结构有助于提高性能,加入 LGAG 和 MSCAM 可以提高性能,而 MSCAM 更为有效。当同时使用 LGAG 和 MSCAM 模块时,效果最好。
为了了解 MSDC 中用于深度卷积的不同多尺度核的效果,做了以下的实验:
因此,在所有实验中选择 [1, 3, 5]
核。
下表是 EMCAD 解码器和基准解码器(即 CASCADE)的计算复杂度的实验结果
可以看到,用了更少的计算量达到了更好的效果。