论文阅读笔记1-LUCF-Net

3 minute read

Published:

门黑魆魆的洞开,我要歌唱: 这里我们曾经怎样生活?


LUCF-Net: Lightweight U-shaped Cascade Fusion Network for Medical Image Segmentation

我们今天来读一篇IEEE Fellow发的医学分割。读下来神清气爽,有一说一,这篇其实没有那么难看懂,里面很多内容都是很基础扎实的,没有很复杂的网络结构,只是使用卷积的组合操作就能很有效的提高网络精度和效果,可能这就是我想成为的目标吧。

背景

Transformer 架构在提取全局信息方面功能强大,但由于其复杂性高,其捕获局部信息的能力有限。

为了应对这一挑战,我们提出了一种用于医学图像分割的新型轻量级 U 型级联融合网络 (LUCF-Net)。它采用非对称结构设计,并结合局部和全局模块,以增强其局部和全局建模能力。此外,还设计了一个多层级联融合解码网络,以进一步增强网络的信息融合能力。

在 CT 格式的多器官数据集、MRI 格式的心脏分割数据集和图像格式的皮肤病学数据集上获得的验证结果表明,所提出的模型在处理局部全局信息方面优于其他最先进的方法,在多器官分割中实现了 Dice 系数 1.54% 的提高和 Hausdorff 距离 2.6 毫米的提高。

此外,作为结合了卷积神经网络和 Transformer 架构的网络,它仅用 693 万个参数和 6.6 GB 浮点运算就实现了具有竞争力的分割性能,而无需进行预训练。

(这个摘要看下来就感觉世界要毁灭了,这就是IEEE Fellow的含金量吗,6.93M的Params,这还玩你妈呢,直接结束了,我可以换研究方向了)

U-Net 仅由具有上采样、下采样和跳过连接的 CNN 组成,在不同领域以最小的复杂度展示了令人印象深刻的性能。它在分割多个器官和皮肤病变等任务上表现尤为出色,巩固了其在医学图像分割中的独特作用。然而,CNN 在医学图像分析中的应用仍然存在挑战。

医学图像通常包含大量大范围的上下文信息,可以捕捉图像的整体结构、形状和分布。这种全面的视图对于精确的诊断和治疗计划至关重要,因为它要考虑器官或组织内的整体布局、大小和空间关系等因素。通过利用这些远距离依赖关系,可以对医学图像进行更准确、更详细的分析。

CNN 在有效捕捉远距离相关性方面存在局限性,这可能导致忽略全局信息并影响分割精度。为了解决这个问题,利用自注意力机制的模型 Transformer 得到了广泛的认可。其出色的长距离依赖建模能力已被引入计算机视觉领域,在图像分割任务中取得了显著的成就。与 CNN 相比,Transformer 在医学图像领域的某些方面具有一定的优势。首先,Transformer 可以捕捉图像中像素之间的全局依赖关系,从而更好地理解整体结构。其次,Transformer 可以提供更高的灵活性。传统的 CNN 模型通常需要手动设计网络结构,而 Transformer 模型则可以通过简单的修改(例如增加或减少层或头)来适应不同的任务。 因此,Transformer 模型在处理各种视觉任务时更加灵活。尽管与 CNN 相比,Transformer 模型具有这些优势。

Transformer有一个致命的缺陷:基于Transformer的网络计算效率往往远低于CNN网络,导致计算成本过高。因此,如何高效利用Transformer模型成为一个关键问题。

为了进一步提升医学图像分割的效果,研究人员开始探索将CNN与Transformer结合起来的方法。通过整合各自的优势,有可能增强对医学图像中复杂属性和长距离依赖关系的处理,最终在降低模型复杂度的同时获得更准确、可靠的分割结果。尽管如此,早期将CNN与Transformer结合起来的研究只是将两者简单地合并在一起,并没有从根本上解决Transformer网络的复杂性问题。

本研究受到EdgeViTs的启发,提出了一种基于局部-全局特征级联的非对称CNN-Transformer网络,在下采样后加入patch-wise的自注意力机制,完成局部和全局特征提取,同时显著降低了网络复杂度。通过在U型网络编码器中构建高效的局部-全局特征提取模块,将CNN提取的局部特征与Transformer提取的全局特征进行有效融合。

本研究主要贡献如下:

  1. 通过在U型网络编码器中引入高效的局部-全局特征提取模块,使CNN得到的局部特征与Transformer提取的全局特征无缝结合。
  2. 设计非对称U型网络架构,降低模型复杂度,在解码器中进行多层特征融合,在训练过程中逐层计算损失,加快网络收敛速度,增强网络融合局部和全局信息的能力。
  3. 采用新的多种损失函数组合解决数据集样本不平衡问题,并通过在线硬样本学习策略进一步提高分割精度。

相关工作

基于CNN的网络

早期的医学图像分割方法大多采用纯CNN结构,U-Net无疑是该领域的开创性工作,它结合了解码器、编码器和跳跃连接,为U型网络架构奠定了基础。

U-Net提出后,各种基于U-Net的方法相继问世。

在 UNet++ 、IR-UNet++ 和 UNet3+等改进版本中,利用跳过连接、多级特征融合和上采样结构,模型的信息传播和特征提取能力得到进一步增强。

在三维医学图像分割中,引入了基于三维卷积的 3D-UNet和 VNet,使医学图像分割网络适用于体数据。

上述基于 CNN 的方法主要采用了多层特征融合、注意力机制等技术来弥补 CNN 网络固有的全局建模能力的局限性,在一定程度上有助于性能的提升。

基于Transformer的网络

Transformer最早是在自然语言处理(NLP)领域引入的,因其能够捕捉广泛相互依赖关系的能力而闻名。

Dsosovitskiy等人将Transformer的运用扩展到计算机视觉领域,将图像分割成token用于Transformer网络,这一突破大大提升了网络提取全局特征的能力。

作为一项开创性的努力,TransUNet将Transformer融入到U型网络架构中,不仅通过将图像特征编码为序列来编码强全局上下文,还通过U-Net混合网络设计充分利用了低级CNN特征。

曹等进一步结合Swin Transformer,用Transformer网络代替解码器和编码器,生成纯Transformer UNet,以弥补CNN网络在全局特征方面的不足。

类似地,DS TransUNet在 TransUNet 的基础上,利用密集网络构建了一个密集连接的纯 Transformer U 形网络。

面对Transformer本身计算能力的限制,越来越多的研究者开始研究基于Transformer的更高效的U-Net架构。

Huang等提出MISSFormer,重新设计了encoder结构中的前馈网络,从而能够更高效地提取局部和全局上下文特征。

Reza等提出DAEFormer,重新定义了自注意力机制,跳过了连接路径,保证了整个特征维度上空间和通道连接的包含,保持了特征的可回收性,从而降低了自注意力机制的计算负担。

大多数基于Transformer的U-Net架构要么是CNN与Transformer的结合,要么是单纯的Transformer结构,这些方法要么没有考虑CNN在局部特征提取中的作用,要么将CNN用于局部特征提取、将Transformer用于全局特征提取,对CNN网络的特征进行自注意力操作,而没有解决Transformer自注意力机制的计算成本问题。

虽然它们能够保持一定的局部和全局特征建模水平,但计算成本和模型参数往往很高。平衡网络性能和规模是一项具有挑战性的任务。鉴于这些考虑,我们致力于构建一个高效的 CNN-Transformer U 形网络。

方法

图 1 展示了 LUCF-Net 的完整结构,它采用了非对称 CNN-Transformer U 形框架。核心组件是局部-全局特征提取模块 (LG Block),无缝集成在编码器的下采样结构中。每个组件的详细信息将在后续章节中描述。

图 1. 所提出的 LUCF-Net 的总体布局,其中涉及采用多输出级联聚合并在训练期间加入多层损失。

LUCF-Net 由多层 CNN-Transformer 编码器和 CNN 解码器组成;

LG 块不会改变特征图大小;

CIE 头将多层输出统一为通用格式;(这里和FCT最后的deepwise卷积很像)

训练损失计算为与位于解码部分的四层特征相关的损失的累积和;

网络的最终输出来自解码部分所有四层特征的聚合。

局部-全局特征提取

图 2. 提出的 LG Block 结构。

LocalAggregation 表示局部特征聚合算子,局部特征聚合由多层 CNN 网络实现;

Conv MLP(CMLP)表示由两个卷积层组成的双层感知器;

Global Sparse Attention 是采样操作后的自注意力;

Transposed Convolution 恢复特征图分辨率;

MLP 是由两个全连接层组成的感知器。

为了实现这一点,我们引入了一个局部-全局特征提取模块,称为 LG Block,如图 2 所示。它获取传入的特征信息并启动局部特征聚合操作以将信息收敛到局部窗口中。随后,对通过均匀窗口采样获得的标记进行注意。最后,通过转置卷积通过邻域扩散传播从注意操作中获得的全局上下文信息。该模块的公式可以描述如下:

编码器和解码器

在编码器部分,初始输入图像经过两层卷积层,随后进行降采样。经过此过程后,原始输入的分辨率降低一半,而通道数相应增加。

随后,将降采样后的图像输入到LG Block中,并进行自注意力操作。此序列重复进行四层卷积降采样和四个LG Block。

解码层仅使用卷积和上采样操作。同样,每经过一层卷积和上采样,图像分辨率都会翻倍。此过程重复进行四层上采样,最终得到与原始输入匹配的图像大小。

需要强调的是,解码器部分不使用LG Block,可以使用跳跃链接和多层级联模块融合来自编码器的局部和全局层面的信息,从而避免在解码器中使用Transformer并减小模型大小。

特征融合

传统的 U-Net 网络通常利用解码端的最后一层作为综合网络输出,并在训练期间计算损失。为了在多尺度图像输出的情况下改善分割,通过跳过连接将来自不同编码器层的多尺度信息集成到解码器中,增强了该架构。来自解码器的每一层的上采样输出都被馈送到单独的解耦头中。

这个解耦头对应于图 1 中的 CIE 头,它将不同尺度的图像协调为一致的输出大小。CIE 头使用双线性插值运算,这与解码器上的上采样操作相同。在训练期间,我们通过比较每一层的输出及其相应的标签来计算损失。最终输出是四个不同阶段的输出的总和。这种结构通过多级级联加强了像素之间的空间关系,从而加快了训练期间的模型收敛。

损失函数

在医学图像分割中,交叉熵损失和 Dice 损失是最常用的损失函数。Dice 系数是视觉计算领域中常用的度量标准,用于衡量两幅图像之间的相似性。然而,Dice 损失的训练表现出明显的波动,因此,它经常与交叉熵损失函数结合在一起。

在这里,我们引入 Lovász Softmax 损失来替代 Dice 损失。我们之所以做出这样的选择,是因为这种损失也直接优化了基于区域的指标。它是一个凸函数,确保它在训练期间不会陷入局部最小值。此外,Lovász Softmax 损失在处理对象边界像素方面表现良好,避免了模糊边缘的产生。Lovász Softmax 损失源自 Jaccard 指数损失的变体,属于 c 类,可以使用以下公式表示:

其中,y_i^c 和 p_i^c 分别表示第i个像素的标签和网络的预测。c是C的子类,表示类别总数,M是一批像素的数量。公式(5)是一个离散函数,不适合直接优化损失。利用Lovász扩展赋予Jaccard指数损失可微性,从而将离散的输入值转化为连续的值。采用该扩展后,Lovász Softmax损失函数可以通过以下公式计算:

其中$f_i(p_i^c)$ 为网络输出在类别c上的概率分布,由Softmax函数得到。 $e_i(c) $为类别c的像素误差,向量$e(c)$为类别c的Jaccard指数的代数。$\bar{\Delta J_c}$为Jaccard指数的Lovász扩展。 为了缓解数据集内样本不平衡的问题,引入了在线难例挖掘(OHEM)损失函数。

在深度学习模型的训练过程中,利用这种损失函数策略来解决类别分布不平衡引起的问题。OHEM损失的目的是将注意力集中在难以分类的样本上,鼓励模型更好地学习困难的案例,从而提高整体性能。在整个训练阶段,OHEM损失的核心思想是从批次中挑选出难以分类的样本进行反向传播。这有效地将模型的注意力引向具有挑战性的实例,帮助模型更好地区分不同的类别。自然地,我们将 OHEM 损失函数中硬样本的定义扩展到像素级。对于每组训练批次,初始损失函数计算当前批次中训练的所有像素的平均交叉熵损失。交叉熵损失表示为:

基于交叉熵损失,OHEM 损失公式可以表述为:

这里,$l_org$ 和 $l_re$ 分别表示所有像素的损失和与所选硬像素相关的损失。变量 K 表示硬像素的数量,通过过滤掉置信度较低的预测像素来确定。 OHEM 损失选择这些置信度较低的像素并计算平均交叉熵损失。随后,将硬像素的平均损失与所有像素的平均交叉熵损失相加。 总之,我们的混合损失可以确定如下:

实验

实现细节

实验使用PyTorch 2.0.0框架进行,在拥有24 GB内存的Nvidia RTX 3090 GPU上进行训练。(太良心了哥们,这里的显卡配置是可复现的)

对于Synapse数据集和ACDC数据集,指定的输入图像尺寸配置为224×224。只有一个通道,训练时使用的batch size为16。

对于两个ISIC数据集,输入图像尺寸配置为512×512,使用3个通道,训练时使用的batch size为4。

我们采用动态学习率进行训练。具体而言,初始学习率为0.05。在训练过程中,随着训练次数的增加,学习率逐渐降低。

LUCF-Net使用SGD优化器进行微调,动量为0.9,权重衰减为0.0001。

在实验中,使用翻转、旋转等数据增强方法来增加数据的多样性。此外,我们在绘制结果图像时使用了 matplotlib 函数。

数据集和评估指标

为了评估网络的泛化性能,使用三个不同的数据集类别进行了实验测试。 Synapse 多器官腹部数据集由 CT 格式数据组成,自动心脏诊断挑战 (ACDC) 数据集由 MRI 格式数据组成,ISIC2016 和 ISIC2018 数据集由图像格式数据组成。除了采用五倍交叉验证方法的 ISIC2018 外,其余实验结果均来自五次实验的平均值和标准差。

  1. Synapse 数据集:Synapse 腹部多器官数据集 包括 30 次腹部 CT 扫描和 3779 张轴向平面捕获的腹部临床 CT 图像。数据集已分为 18 次训练扫描和 12 次随机测试扫描。我们使用了与 相同的处理方法。我们将最后一轮训练结果作为测试权重。 评估指标是 8 个腹部器官的平均 DSC 和平均 HD。
  2. ACDC数据集:自动心脏诊断挑战(ACDC)数据集:该数据集包括从不同患者获得的100张MRI扫描,每张扫描都标注了三个器官:左心室(LV)、右心室(RV)和心肌(MYO)70个案例被分配用于训练,10个用于验证,20个用于测试。平均DSC作为评估的性能指标。
  3. ISIC数据集:对于ISIC-2016 数据集,共有900个训练样本和379个验证样本。ISIC2018 数据集包含2594幅图像及其相应的标签,图像分辨率范围从720×540到6708×4439。如参考文献所述,进行了五倍交叉验证以确保公平评估。评估基于平均 DSC 和平均并集交点 (IoU) 分数。

(如果我选的话可能也就是这三个数据集吧应该,这三个数据集比较基础也满受人偏爱的)

实验结果

Synapse数据集上的结果

如表一所示,在腹部多器官数据集上,我们提出的LUCF-Net与SOTA方法的比较结果如下。最后两列代表8个器官的平均Dice相似系数(DSC)和平均Hausdorff距离(HD)。下面不同器官的值代表平均DSC。 与其他基于CNN或Transformer的模型相比,LUCF-Net在DSC方面领先TransCASCADE 1.54%,在HD方面领先2.60毫米(mm)。图3中的数字描绘了各种方法在多器官CT数据集上实现的分割结果。这些图像表明,LUCF-Net在大多数器官分割任务中准确地描绘出复杂结构,产生更精确的分割结果,即使在具有挑战性的背景下也表现出有竞争力的性能。

图 3. 通过视觉分析比较 Synapse 腹部多器官数据集上的各种技术。每行代表不同的受试者,第一列是每个受试者的真实标签,不同的颜色代表不同的器官。

ACDC 数据集上的结果

表二展示了 LUCF-Net 和 SOTA 方法在 ACDC 数据集上的性能比较。最后一列表示心脏三个节段的平均 DSC,而前三列表示不同节段的平均 DSC。值得注意的是,LUCF-Net 获得了最高的平均 DSC,为 92.19%,如图 4 所示。

(ACDC数据集也太高了,妈的,这还怎么玩)

皮肤病变分割结果

表三显示了 LUCF-Net 和其他网络在 ISIC 2016 和 ISIC 2018 上的平均 DSC 和平均 IoU 性能的比较结果。值得注意的是,LUCF-Net 在处理图像数据集方面也表现出了竞争力。图 5 还展示了结果的视觉比较。结果表明,LUCF-Net 能够捕捉复杂的细节并产生更精确的轮廓。与纯 Transformer 网络相比,我们的方法可以捕捉到更精细的局部细节,表明它在提取局部和全局特征方面的有效性。

注意层分析

为了进一步验证网络中局部和全局特征建模的有效性,我们从注意层的角度进行了研究。我们选择了最先进的U型网络架构,其下采样编码设计与Swin Unet、MISSFormer和LUCF-Net有相似之处。我们选择了经过自注意层处理的四层特征图进行比较。

对于DAEFormer,其在编码阶段只有三层下采样,我们选择了这三层的特征图经过子注意层处理进行比较。

与其他网络相比,LUCF-Net在浅层编码阶段捕获了更详细的图像信息,从而能够捕捉到更精细的局部特征细节。

在深度编码阶段,LUCF Net对目标信息实现了更好的全局建模,从而对特征信息中所蕴含的对象有了更完整的表示。

模型参数对比

在表Ⅳ中,我们将本文提出的方法与医学图像分割模型的参数数量进行了比较。网络的输入形状标准化为1×1×224×224,并使用Params和GFLOPs量化模型的计算强度。Params表示神经网络的参数个数,GFLOPs表示模型在推理或训练时每秒执行的浮点运算次数的数量级。与基于SOTA Transformer的网络不同,我们的模型在复杂度方面表现出一定的优势,具体来说,我们在利用CNN架构的简单性的同时,实现了超越现有高复杂度Transformer网络的性能。

(别人都是20M+,30M+的Params,这篇不到7M直接结束游戏了)

消融研究

首先从网络架构和损失函数两个方面分析了所提框架的有效性。研究了超损失参数对模型性能的影响。 所有实验均在 Synapse 数据集上进行,使用 DSC 和 HD 作为评估指标。 消融实验的细节如下。

1)网络架构消融实验

为了评估网络架构的有效性,在网络结构消融研究中采用了由交叉熵和 Dice 损失组成的双重损失函数。 超参数配置遵循中概述的指导方针,如表 V 中所述。 结合局部和全局特征提取模块(单独或组合进行特征融合)可显着提高网络性能。 同时使用这两个模块可使 DSC 增加 1.24%,HD 增加 6.54 毫米,超过了使用单个模块所获得的结果。

2)损失函数消融实验

为评估所采用损失组合的有效性,在一致的实验条件下,对LUCF-Net进行训练,并使用不同的损失函数进行实验。考虑到损失计算涉及多级特征融合,将特征融合融入到损失函数变量的消融研究中。 OHEM损失源自交叉熵损失,Lovász Softmax损失和Dice损失量化集合相似性。我们分别用OHEM损失和Lovász Softmax损失替换交叉熵损失和Dice损失。如表Ⅵ所示,Lovász Softmax损失和OHEM损失的组合被证明更适合特征融合网络。与交叉熵损失和Dice损失的组合相比,DSC增加了1.16%,HD减少了1.69 mm。

此外,图 7 还描绘了所提出的最终网络架构上各种损失函数组合的训练损失曲线。使用 OHEM 损失和 Lovász Softmax 损失训练的网络与使用其他损失函数训练的网络相比表现出增强的稳定性。

表 Ⅶ 扩展了 OHEM 损失和 Lovász Softmax 损失的使用,概述了融合不同层对模型性能的影响,而图 8 则说明了不同层的训练损失曲线。输入到 CIE Head 的四个分层特征来自解码器输出的四个特征图,所有这些特征图都经过 LeakyReLU 操作。

CIE Head 值的确定仅依赖于卷积和双线性插值技术,确保输出形状的一致性。我们分别取最后一层、最后两层、最后三层和最后四层的输出之和作为融合不同层深度特征的结果。结果表明,随着融合层的加深,模型性能和网络收敛速度都有明显改善。

3)损失超参数消融实验

我们采用网格搜索的方式探究不同的损失函数超参数组合对模型性能的影响,如图9所示。单独使用Lovász Softmax loss或OHEM loss训练的模型性能并不理想,尤其是当仅使用OHEM loss作为训练损失函数时。OHEM loss更注重像素级的分类准确率,而Lovász Softmax loss则注重衡量两组之间的相似性。通过优化组合,可以取得更好的效果。

讨论

目前融合 CNN 和 Transformer 的方法通常需要大量的计算和存储资源。 为了解决这些问题,我们加入了一个稀疏自注意力 Transformer 模块,将稀疏自注意力集成到每个编码器级别,并将 CNN 层与 Transformer 融合,以有效捕获局部和全局特征。当比较 LUCF-Net 在 CT、MRI 和图片格式的四个数据集上的性能时,它始终优于当前的 SOTA 方法。可视化结果表明,LUCF-Net 有效地捕获了局部细节和总体全局数据,展示了其强大的性能。此外,由于 LUCF-Net 的设计复杂,其复杂性优于其他网络,其特点是简单的 Transformer 模块和非对称网络架构。

由于其低复杂度和出色的分割性能,LUCF-Net 有可能成为可靠的骨干网络。现有的方法虽然如HiFormer和PVT-CASCADE将CNN与Transformer进行了整合,但并未充分解决自注意力机制带来的计算负担这一根本问题

高效的纯Transformer U型网络(如MISSFormer)可能会忽视CNN网络在区域特征提取方面的优势,而对称的Transformer模块设计也增加了网络的复杂度。本文提出的LUCF-Net通过利用高效、稀疏的自注意力机制,有效降低了Transformer模块的计算要求,同时结合非对称设计和多层特征级联融合机制,在简化网络的同时提升了CNN提取局部特征的能力,因此LUCF-Net在医学图像分割任务中表现出色。

结论

本研究通过结合 CNN 和 Transformer 提出了一种新的医学图像分割方法,称为 LUCF-Net。与其他基于 CNN 和 Transformer 的 SOTA 模型相比,LUCF-Net 不仅可以捕获更详细的图像信息,还可以对目标信息进行更好的全局建模。对这些全局特征进行建模有助于网络更好地理解整个图像上下文,从而提高分割性能。LUCF-Net 还在降低模型复杂度的情况下表现出更好的分割性能,展现出医学图像分割应用的潜力。然而在医学图像处理中,由于可用样本数量相对有限,经常存在数据不足的挑战,难以支持完全监督的训练。为了应对这一挑战,Wang 等人 采用一种创新方法,将视觉 Transformer 与一致性正则化框架相结合,与基于带有有限注释数据的 CNN 网络的半监督分割框架相比取得了卓越的性能。在未来的工作中,我们的方法可以与半监督医学图像分割相结合,以进一步提高分割性能。