迁移学习通过使用在更大数据集上预训练的参数来初始化深度学习模型,从而提高它们的性能。最近看到了一篇效果拔群的文章,作者通过把CNN与迁移学习结合,在医学图像数据中,准确率达到了98%!像Transformer这样的注意力网络越来越多地用于图像分析,以捕捉像素之间的远距离关系。然而,CNN模型只能捕捉图像中的局部空间关系。这篇论文将Transformer和CNN相互补充,本文是对该论文的阅读记录,感兴趣的话可以参考一下,如果想要进一步了解研究工作详情内容的话,建议移步阅读原英文论文,地址在这里,如下所示:
摘要
迁移学习通过使用在大规模数据集上预训练的参数初始化深度学习模型,从而提高其性能。直观上,当预训练数据集与目标任务领域相同时,迁移学习更为有效。NASA最近的一项研究表明,使用在显微镜图像上预训练的CNN编码器的编码器-解码器算法在微观结构分割任务中比在自然图像上预训练的编码器表现更好。然而,CNN模型仅捕捉图像中的局部空间关系。近年来,注意力网络如Transformer在图像分析中被广泛用于捕捉像素之间的长程关系。在本研究中,我们比较了在显微镜图像和自然图像上预训练的Transformer和CNN模型的分割性能。我们的结果部分证实了NASA的研究,即在不同成像和样本条件下获取的分布外图像的分割性能在使用显微镜图像进行预训练时显著提高。然而,对于一次性学习和少样本学习,Transformer的性能提升较为有限。我们还发现,对于图像分割任务,预训练的Transformer和CNN编码器的结合始终优于单独使用预训练的CNN编码器。我们的数据集(约50,000张图像)结合了NASA数据集的公开部分和我们收集的额外图像。即使训练数据量较少,我们的预训练模型在图像分割任务中仍表现出显著更好的性能。这一结果表明,Transformer和CNN相互补充,当在显微镜图像上预训练时,它们对下游任务更为有益。
关键词:显微镜、Swin Transformer、CNN和CS-UNet分割。
1 引言
显微镜成像提供了关于物质的实际信息,但要获取形态、尺寸和分布的定量信息,需要在显微照片上进行手动测量,这不仅耗时且劳动密集,还容易产生偏差。材料结构和现象的长度和时间尺度在各组分之间差异显著,增加了复杂性。因此,建立过程、结构和性能之间的联系是一个具有挑战性的问题[1], [2]。
深度学习(DL)因其能够自动提取重要信息而被广泛应用于复杂系统。研究人员已将DL算法应用于图像分析,以识别结构并确定微观结构与性能之间的关系。DL已被证明可以补充基于物理的方法用于材料设计。然而,DL需要大量的训练数据,而有限的显微镜图像数量往往会降低其有效性。为此,研究人员开发了诸如迁移学习、多保真度建模和主动学习等技术,以使DL适用于较小的数据集[2], [3]。迁移学习使用在一个较大数据集上预训练的模型的参数来初始化一个在较小数据集上训练的模型,以进行下游任务。例如,一个在自然图像上预训练的卷积神经网络(CNN)可以用于初始化一个用于图像分割的神经网络,以提高其精度和减少训练时间。
图1. 用于微观结构分割的编码器-解码器架构,采用迁移学习方法,其中CNN和Swin-T模型分别在ImageNet和显微镜图像上预训练。预训练的CNN和Swin-T模型的权重用于初始化编码器,而Swin-T模型的权重用于初始化解码器。
使用自然图像(如ImageNet)进行预训练并不理想,因为这些模型识别的高级特征在显微镜图像中并不存在。Stuckner等人[1]的最新工作证明了使用名为MicroNet的显微镜数据集(包含超过110,000张图像)预训练CNN的优势。他们评估了预训练的CNN编码器用于分割镍基高温合金(Super)和环境屏障涂层(EBC)的显微镜图像。使用MicroNet进行预训练在一次性学习和少样本学习以及具有不同成分、蚀刻和成像条件的分布外图像中,交并比(IoU)测量结果显著提高。
近年来,基于注意力的神经网络(称为Transformer)在计算机视觉中被广泛采用。CNN通过卷积滤波器从图像的局部区域提取特征以捕捉像素之间的空间关系,而Transformer将图像分割成块并将其输入基于Transformer的编码器以捕捉图像中像素之间的长程关系[4], [5]。因此,结合CNN和Transformer可能比单独使用任一模型在迁移学习中更为有效。
在本文中,我们评估了使用CNN和Transformer编码器组合进行显微镜图像分割的迁移学习。迁移学习方法如图1所示,该方法包含一个用于图像分割的编码器-解码器架构。每个编码器将输入图像转换为潜在表示向量以提取语义信息。解码器将提取的信息映射回输入图像中的每个像素,以生成图像的像素级分类[1], [2], [6]。
我们使用了一个流行的Transformer版本——Swin-Transformer,特别是其_tiny_版本——Swin-T,以提高效率。我们的预训练数据集包含约50,000张显微镜图像,分为74个类别,我们将其称为_MicroLite_数据集。Swin-T模型可以使用在ImageNet上预训练的模型权重进行初始化,然后在MicroLite上进行微调。我们使用在MicroNet上预训练的CNN模型来初始化图1中蓝色编码器的权重,而Swin-T模型用于初始化橙色编码器和解码器的权重。CNN和Swin-T编码器的输出在连接到解码器之前进行融合。为了评估迁移学习的分割性能,我们比较了在ImageNet上单独预训练的模型和在显微镜图像上预训练的模型在7个数据集(Super和EBC的子集)上的图像分割IoU分数。我们的结果显示,尽管一次性学习和少样本学习的分割准确性有所提高,但提升幅度不如NASA论文[1]所示。对于分布外图像,使用显微镜图像进行预训练的准确性仍然显著更好。我们还比较了使用CNN、Swin-T及其组合进行分割的性能。我们的结果表明,在大多数情况下,组合优于单独使用CNN,并且在某些情况下优于单独使用Swin-T。
2 相关工作
在本节中,我们回顾了与图像分析中的Transformer以及我们部署和评估的基于Transformer的分割算法相关的先前研究。
Transformer
CNN基于卷积运算符提供平移等变性。然而,CNN的局部感受野限制了其捕捉像素之间长程关系的能力[5]。近年来,Transformer被用作CNN的替代方案用于计算机视觉任务。Transformer是由Vaswani等人[8]引入的一种深度神经网络,由于其能够捕捉文本等序列数据中的长程依赖关系,在自然语言处理(NLP)中取得了成功应用。这种方法在语言翻译、文本分类和文本生成等NLP应用中显著提高了性能。与NLP任务相比,在计算机视觉任务中使用基于Transformer的模型更具挑战性,因为图像具有尺寸变化、噪声和冗余模态。自注意力过程是Transformer的基本构建块,旨在学习自对齐,从而捕捉图像块之间的长程关系[5]。这导致了在计算机视觉领域中对基于Transformer方法的广泛兴趣[4],如图像识别、图像分割[9]、目标检测[10, 11]、图像超分辨率和图像生成[12]。
Dosovitskiy等人[4]提出了一种基于原始Transformer网络[8]的视觉Transformer(ViT),尽可能少地进行修改以捕捉输入图像的全局上下文。ViT将每个图像分割成块,并将其线性嵌入提供给Transformer。图像块的处理方式与NLP应用中的标记相同。使用监督学习对模型进行训练以进行图像分类。在预训练于包含3亿张图像的JFT数据集后,模型在下游识别基准(如ImageNet分类)上进行微调。ViT在ImageNet分类任务中表现优于传统CNN,达到了88.5%的准确率。然而,ViT需要更多的计算资源进行训练。此外,计算每个自注意力块的SoftMax的复杂度与输入序列长度成二次关系,限制了其应用于高分辨率图像[4, 5]。
为了改进Transformer模型以捕捉局部信息,Liu等人[7]提出了一种新的视觉Transformer,称为移位窗口Transformer(Swin Transformer)。该方法提出了一种新的通用主干网络用于图像分类和识别任务,并取得了最先进的性能。该模型使用移位窗口方案捕捉视觉实体的尺度变化和高分辨率像素,具有与输入图像尺寸成线性计算复杂度。相比之下,ViT[4]模型生成单一低分辨率的特征图,并且由于对所有块进行全局自注意力,计算复杂度与输入图像尺寸成二次关系。Swin Transformer在ImageNet分类任务中达到了87.3%的准确率,在COCO检测任务中达到了58.7%的框平均精度分数,在ADE20K数据集的分割任务中达到了53.5%的mIoU。作者[7]还提出了一种新的Swin Transformer V2版本,将其扩展到30亿参数,并允许其训练高达1536×1536像素的图像。Swin Transformer V2[14]修改了Swin注意力模块[7]以提高窗口分辨率和模型容量。通过将预归一化替换为后归一化配置,使用缩放余弦注意力代替点积注意力,并将先前的参数化方法替换为对数间隔连续相对位置偏置方法[14]。
分割
Ronneberger等人[15, 16]提出了一种称为U-Net的全卷积网络(FCN)。U-Net是一种对称的U形编码器-解码器神经网络,用于语义图像分割。U-Net由典型的下采样编码器和上采样解码器结构以及它们之间的“跳跃连接”组成。这些连接从编码器复制特征图并将其与解码器中的特征图连接。
Xie等人[17]提出了一种称为SegFormer的语义分割框架,结合了Transformer和轻量级多层感知器(MLP)解码器。SegFormer基于编码器-解码器架构,其中编码器是一个分层结构的Transformer,输出多尺度特征而不需要位置编码,轻量级全MLP解码器聚合不同层的信息并结合局部和全局注意力以生成最终的语义分割掩码。SegFormer使用4×4像素的块大小输出分割图。这种方法有助于提高密集预测任务的性能,并在ADE20K数据集上达到了50.3%的mIoU分数,在Cityscapes数据集上达到了84%的mIoU分数。
Chen等人[18]提出了TransUNet,一种采用混合CNN-Transformer编码器和多个上采样层的CNN解码器的U形架构。该方法结合了CNN特征的详细高分辨率空间信息和Transformer编码的全局上下文。TransUNet架构在编码器中包含12层视觉Transformer(ViT[4]),这些层编码从CNN层获得的图像块。这些编码特征然后通过解码器中的上采样层进行上采样以生成最终的分割图,并结合跳跃连接。TransUNet在性能上优于基于CNN的模型。
Cao等人[19]提出了一种名为Swin-Unet的纯Transformer模型,用于医学图像分割。Swin-Unet具有U形编码器-解码器架构,通过将图像块输入Swin-Unet模型进行局部-全局语义特征学习。编码器和解码器结构均受到分层Swin-Transformer[7]的启发,使用移位窗口。
Hatamizadeh等人[20]提出了UNETR模型,一种使用ViT[4]作为编码器以捕捉全局多尺度信息的U形编码器-解码器架构。Transformer编码器通过跳跃连接与CNN解码器连接以计算最终的语义分割输出。所提出的方法在各种医学数据集的图像分割任务中具有出色的准确性和效率。Hatamizadeh等人[21]还提出了一种用于多模态MRI图像脑肿瘤分割的新模型,称为Swin UNETR。它使用Swin Transformer[7]作为编码器,并通过不同分辨率的跳跃连接连接到基于CNN的解码器。在BraTS 2021分割挑战中,该团队展示了所提出的多模态3D脑肿瘤分割模型在验证阶段排名前列。
Heidari等人[22]提出了HiFormer,一种用于医学图像分割的CNN和Transformer桥接模型。HiFormer有效地使用两种多尺度特征表示和一个双层融合(DLF)模块来融合全局和局部特征。广泛的实验表明,HiFormer在计算复杂性和结果质量方面优于其他基于CNN、Transformer和混合方法。HiFormer提供了一种端到端训练策略,将来自Swin Transformer的全局上下文表示和来自CNN模块的局部代表性特征集成在编码器中,随后由解码器输出最终的分割图。
Azad等人[23]提出了一种新颖的方法,称为TransDeepLab,这是一种基于纯Transformer的DeepLabV3+架构,用于医学图像分割。所提出的模型使用具有移位窗口的分层Swin-Transformer来建模空洞空间金字塔池化(ASPP)模块。编码器模块将输入图像分割成块并应用Swin-Transformer块以编码局部语义和长程上下文表示。设计了一个具有不同窗口大小的Swin-Transformer块的金字塔用于ASPP以捕捉多尺度信息。提取的多尺度特征然后通过交叉上下文注意力机制融合到解码器模块中。最后,在解码路径中,提取的多尺度特征被上采样并与编码器中的低级特征连接以细化特征表示。
本节提到的混合架构主要集中在用Transformer替换编码器中的CNN,如UNETR和Swin UNETR,或将CNN与Transformer堆叠以形成新的编码器,如Transnet[24, 25]。
用Transformer替换编码器中的CNN赋予了模型捕捉网络中长距离依赖关系的能力。然而,由于编码器中CNN的移除,导致缺乏详细的纹理特征提取。将CNN与Transformer堆叠以形成新的编码器未能考虑到自注意力全局建模能力和卷积局部建模能力之间的互补关系。相反,它们将卷积操作和自注意力视为两个独立且不相关的操作[24, 26]。
为了克服这些缺点,CS-Unet中的编码器并行使用CNN和Transformer从显微镜图像中获取丰富的特征信息。CNN用于提取低级特征,Swin-T用于提取全局上下文特征,这些特征通过跳跃连接在不同阶段/层融合到解码器中。此外,为了减少传输过程中的特征损失并增加Swin-T编码器提取的上下文信息,两个连续的Swin-T块中的多层感知器(MLP)被替换为残差多层感知器(ResMLP)。
3 方法论
我们的目标是证明基于Transformer的显微镜图像预训练模型对下游任务(如图像分割)是有益的,并且比基于CNN的预训练模型更为鲁棒。为此,我们完成了以下任务。
-
我们收集了一个经过预处理的显微镜数据集,约50,000张图像(MicroLite),
-
我们在MicroLite上预训练Transformer编码器,并使用它们初始化几个基于Transformer的分割算法(Swin-Unet、TransDeeplabv3+和HiFormer)和基于CNN和Transformer编码器的混合分割神经网络(CS-UNet)的参数,
-
为了展示CS-UNet的优势,我们比较了基于CNN的分割算法[1]与基于Transformer的分割算法和CS-UNet的顶级性能。这些算法使用NASA团队[1]的7个测试集进行比较,其中CNN编码器在MicroNet[1]上预训练,Transformer编码器在MicroLite上预训练,
-
为了评估在领域内数据上预训练的优势,我们比较了CS-UNet在ImageNet和MicroLite上预训练时的顶级性能,
-
为了详细考察预训练对领域内数据的影响,我们比较了不同预训练设置下基于CNN的分割算法的平均性能。同样,我们比较了不同预训练设置和Transformer架构下基于Transformer和混合分割算法的平均性能,
-
最后,为了说明我们混合策略的鲁棒性,我们比较了三种类型分割算法在所有配置下的平均性能。
数据集预处理
我们的MicroLite数据集中的图像来自多个来源,包括使用不同测量技术(如光学显微镜、SEM、TEM和X射线)的不同材料和化合物的图像。MicroLite聚合了Aversa数据集[27]、超高碳钢显微照片[28]、材料数据存储库中的SEM图像以及一些最近出版物的作者提供的图像[29, 30, 31, 32, 33, 34, 35, 36, 37]。
Aversa数据集包括超过25,000张10个类别的SEM显微镜图像,每个类别包含不同尺度(包括1、2、10、20微米和100、200纳米)和对比度的图像。为了正确分类这些图像,我们使用预训练的VGG-16模型从这些图像中提取特征图,并使用K-means算法对特征图进行聚类,以便将具有相似特征图的图像分组在同一类别中。经过预处理步骤后,我们获得了53个类别。Aversa数据集的作者手动将一小部分图像(1038张)分类为层次数据集,其中10个类别进一步分为27个子类别[27]。我们对这1038张图像的分类与手动分配的子类别基本一致。请注意,我们有更多的类别,因为我们处理了整个Aversa数据集。
总的来说,MicroLite包括50,000张显微镜图像,分为74个类别,这些图像是通过以下预处理步骤获得的。
-
去除图像中的任何伪影,如比例尺,
-
根据原始图像的大小,将图像分割成512×512像素的块,有或没有重叠,
-
应用数据增强以增加数据集的大小,
-
将原始图像、图像块和增强图像聚合以形成最终数据集。
预训练
我们在显微镜图像上训练Swin Transformer以学习特征表示,以便将其迁移到分割等任务中。我们评估了两种类型的训练。
图2. 原始的Swin-T架构由四个阶段组成。每个阶段包含两个Swin Transformer块,除了第三阶段包含六个Swin Transformer块[7]。
-
在ImageNet上微调预训练模型与MicroLite(表示为ImageNet → MicroLite),
-
从头开始在MicroLite上预训练模型(表示为MicroLite)。
分类任务使用Swin-T,这是Swin Transformer的tiny版本。Swin-T包含两种架构:原始Swin-T和中间网络,分别包含[2,2,6,2]和[2,2,2,2]个Transformer块。图2显示了Swin-T的原始架构。我们推测中间网络可能足以用于显微镜分析任务,因为早期层学习角边缘和形状,中间层学习纹理或模式,而原始模型中的更深层学习高级特征,如眼点和尾附器。原始和中间Swin-T模型从头开始在MicroLite上预训练,模型权重随机初始化。这两个模型也在ImageNet上预训练并在MicroLite上微调。
预训练步骤使用AdamW优化器进行30个epoch,采用余弦衰减学习率调度器,线性预热5个epoch,批量大小为128。初始学习率为10^-3,权重衰减为0.05。微调步骤同样使用AdamW优化器进行30个epoch,批量大小为128,但学习率降低到10^-5,权重衰减降低到10^-8。模型训练直到验证分数没有改进,使用早期停止标准,耐心为5个epoch。训练数据使用allumentations库进行增强,包括对比度和亮度的随机变化、垂直和水平翻转、光度失真和添加噪声。
对于下游分割任务,为每个任务训练了几个模型,包括Swin-Unet、HiFormer和TrasDeeplapv3+。对这些模型在ImageNet和显微镜图像上预训练的结果进行了比较分析。
结合CNN与Transformer(CS-UNet)
CNN由于其固有的局部性,无法捕捉长程空间关系。Transformer被引入以克服这一限制。然而,Transformer在捕捉低级特征方面存在局限性。研究表明,局部和全局信息对于具有挑战性上下文中的密集预测任务(如分割)是必不可少的。一些研究人员引入了混合模型,有效地桥接CNN和Transformer用于图像分割。初始化混合模型中的CNN和Transformer的权重将显著提高性能。因此,我们引入了一种称为CS-UNet的混合UNet,这是一种使用CNN和Transformer的U形分割模型。如图3所示,该方法包括编码器、瓶颈、解码器和跳跃连接。
图3. 提出的CS-UNet架构概览,包括编码器、瓶颈和解码器,其中编码器部分包括CNN和Swin-T。
编码器结合了CNN编码器和Swin-T编码器,其中CNN用于提取低级特征,Swin-T用于提取全局上下文特征。Swin-T编码器对输入图像进行非重叠块分割,应用自注意力机制捕捉全局依赖关系。编码器在不同尺度上捕捉整个图像的长程依赖关系和上下文信息。受TFCN(Transformers for Fully Convolutional dense Net)[38]和轻量级Swin-Unet[37]的启发,两个连续的Swin-T块中的多层感知器(MLP)被替换为残差多层感知器(ResMLP)。如图4所示,ResMLP用于减少传输过程中的特征损失并增加编码器提取的上下文信息。ResMLP如图5所示,由两个GELU[39]非线性层、三个线性层和两个dropout层组成。CNN编码器处理输入图像的一系列卷积层,逐渐减少空间维度同时提取分层特征。编码器在早期层捕捉低级特征,在更深层捕捉高级语义特征。
为了融合来自两个编码器的信息,跳跃连接将CNN编码器和Swin-T编码器的特征图与相应的解码器层连接。为了确保CNN和Swin-T编码器的特征维度兼容,需要在融合之前对维度进行归一化。这是通过将CNN块获得的特征通过线性嵌入层来实现的,该层将特征图从形状(B,C,H,W)展平并重塑为(B,C,H×W),其中B、C、H、W分别是批量大小、通道数、高度和宽度。展平的特征图转置以交换最后两个维度,结果形状为(B,H×W,C),然后与Swin-T编码器提取的特征融合。通过融合来自不同编码器路径的信息,跳跃连接使解码器能够从CNN编码器捕捉的局部空间细节和Swin-T编码器捕捉的全局上下文中受益。
解码器类似于Swin-Unet[19]的解码器,采用块扩展层通过重塑相邻维度的特征图来上采样提取的深层特征,从而有效地实现2倍上采样。此外,它将特征维度减少到原始维度的一半。这使得解码器能够在减少特征维度的同时重建具有增加空间分辨率的输出。最终的块扩展层进一步进行4倍上采样以恢复特征图的分辨率以匹配输入分辨率(W×H)。最后,对这些上采样的特征应用线性投影层以生成像素级分割预测。
编码器部分可以使用不同的CNN系列,如EfficientNet、ResNet、MobileNet、DenseNet、VGG和Inception。我们使用MicroNet初始化CNN权重,使用MicroLite初始化Transformer权重。
4 结果
预训练的Swin-T模型用于在74个不同类别中分类显微镜图像。Swin-T模型要么在ImageNet(特别是imageNet1K数据集)上预训练并在MicroLite上微调,要么从头开始在MicroLite上训练。训练在验证准确性不再提高后停止。模型准确性使用top-1和top-5准确性进行评估。top-1准确性测量测试样本中正确标签的预测百分比,而top-5准确性测量前五个预测中正确标签的百分比。
图4. 用于改进Swin Transformer块的ResMLP [37]。
图5. ResMLP模块块 [37]。
如表1所示,从头开始训练的Swin-T模型收敛时间更长。具体来说,原始Swin-T模型需要23个epoch,而中间版本需要19个epoch。相比之下,在ImageNet上预训练并在MicroLite上微调的模型收敛时间更短,原始Swin-T模型需要13个epoch,中间版本需要12个epoch。在ImageNet上预训练并在MicroLite上微调的模型在top-1和top-5准确性方面表现更好。
微观结构图像分割
为了评估Swin-T模型提取特征表示的能力,预训练模型用于初始化分割任务的模型。为了与NASA研究[1]进行比较,我们使用了从两种材料(镍基高温合金(Super)和环境屏障涂层(EBC))派生的相同7个显微镜数据集。EBC数据集有两个类别:氧化层和背景(非氧化)层,Super数据集有三个类别:基体、次生和三生。每个数据集分割中的图像数量如表2所示。Super-1和EBC-1包含各自材料的完整标记数据集。Super-2和EBC-2在训练集中只有4张图像,以评估模型在少样本学习中的性能。Super-3和EBC-3在训练集中只有1张图像,以评估一次性学习中的性能。Super-4包含在不同成像和样本条件下获取的测试图像。本研究中的所有分割模型均使用PyTorch[40]进行训练。
表1. 预训练模型的分类准确性,其中“ImageNet → MicroLite”表示模型在ImageNet上预训练,然后在MicroLite上微调。
表2. 每个实验数据集在训练、验证和测试集中的显微镜图像数量。
EBC和Super数据集以与NASA研究[1]类似的方式进行增强,包括随机裁剪到512×512像素、对比度、亮度和伽马的随机变化,以及添加模糊或锐化。EBC数据集进行水平翻转,Super数据集随机进行垂直和水平翻转和旋转。训练使用Adam优化器,初始学习率为2×10^-4,直到验证准确性在30个epoch内没有改进。之后,训练继续使用10^-5的学习率,直到在额外30个epoch内没有验证改进时触发早期停止。由于数据集不平衡,损失函数设置为加权平衡交叉熵(BCE)和dice损失的加权和,BCE权重为70%[1]。
CS-UNet架构是一个灵活的模型,可以使用不同的CNN系列和预训练模型进行训练。表3显示了用于训练CS-UNet模型的各种预训练权重组合。第二列显示用于初始化Swin-T编码器的预训练权重,第三列显示用于初始化CNN编码器的预训练权重。在最后一列中,我们使用术语_Microscopy_来指代CNN编码器在MicroNet上训练、Transformer编码器在MicroLite上训练的情况。其他预训练权重的组合也可以用于训练CS-Unet模型。例如,Swin-T编码器可以使用MicroLite权重初始化,CNN编码器可以使用ImageNet→MicroNet权重初始化。CS-UNet架构的灵活性允许研究人员尝试不同的预训练权重组合,以找到最适合其特定任务的组合。
表4比较了在MicroNet上预训练的UNet++/UNet([1]中的图3-5报告的数字)、在MicroLite上预训练的Transformer模型(包括Swin-UNet、TransDeepLabV3+和HiFormer)和在MicroNet和MicroLite上预训练的CS-UNet的顶级性能。每个实验的最高准确性以粗体显示。CS-UNet在大多数实验中表现最好,除了EBC-2和EBC-3。对于具有充足训练数据的实验,如Super-1和EBC-1,UNet++/UNet、Transformer和CS-UNet之间的差异很小。对于少样本学习实验,如Super-2和EBC-2,CS-UNet的准确性提升有限。对于一次性学习实验,结果不一,CS-UNet在Super-3中略有改进,而在EBC-3中显著提升。对于分布外学习,CS-UNet显示出显著优于UNet或Transformer的改进。
表3. CS-UNet中的编码器组合,其中第二列显示Swin-T编码器的预训练模型,第三列显示CNN编码器的预训练模型。最后一列显示两种类型编码器的模型,其中“Microscopy”可以是MicroNet或MicroLite。
表4. 在MicroNet上预训练的UNet++/UNet([1]中的图3-5报告的数字)、在MicroLite上预训练的Transformer分割模型和在MicroNet和MicroLite上预训练的CS-UNet的顶级性能(IoU)比较。
总体而言,表中的结果表明CS-UNet是一种有前途的图像分割方法。CS-UNet在所有实验中与UNet++/UNet相似或显著更好,并且在大多数实验中优于Transformer。请注意,MicroLite的大小约为MicroNet的一半。尽管如此,Transformer + MicroLite的性能与UNet++/UNet + MicroNet相当或更好。附录A显示了表4中表现最好的Transformer模型的配置。表现最好的CS-UNet模型的配置将在下一节中显示,我们将比较CS-UNet在显微镜图像和ImageNet上预训练时的性能。
镍基高温合金(Super)分割
图6和图7比较了CS-UNet在Super数据集上的最佳性能,当其在显微镜图像和ImageNet上预训练时。Super-1和Super-2的IoU分数对于两种预训练模型相似,而Super-3的IoU分数对于显微镜模型为93.5%,对于ImageNet模型为87.01%。这一结果与NASA论文中的结果不同,其中Super-2的性能也有所提高。似乎随着CS-UNet能力的增强,领域内数据集的优势在一次学习中比少样本学习中更为显著。ImageNet模型未能识别许多暗对比图像中的三生析出物,如橙色三角形所示。ImageNet模型还过度分割并合并了一些次生析出物,如绿色箭头所示。
图6. (a) 是Super数据集的训练和测试图像示例[1],其中Super-3的训练图像用红色轮廓标出。(b) 是使用Super-1、Super-2和Super-3数据集训练CS-UNet时测试图像的最佳分割掩码,左侧为在ImageNet上预训练,右侧为在显微镜图像上预训练。每个实验的最佳模型CNN/Transformer编码器和IoU分数在分割掩码上方。绿色三角形指向ImageNet(而非显微镜)模型错误分割次生析出物的区域。黄色三角形指向ImageNet(而非显微镜)模型错误分割三生析出物的区域。
图7. Super-4的测试图像具有与训练图像不同的成像条件。左侧列显示Super-4的测试图像[1]。准确性掩码颜色与图6相同。中间列显示最佳ImageNet模型的IoU准确性掩码。右侧列显示最佳显微镜模型的相同内容。每行显示同一测试图像及其准确性掩码。绿色箭头显示模型过度分割的示例,以及显微镜模型准确分割次生析出物的地方。黄色表示错误地将三生析出物识别为次生析出物。栗色表示模型错误地将次生析出物识别为三生析出物的地方。青色表示模型过度分割次生析出物的地方。
对于Super-4,包含不同成像条件的图像,显微镜模型将ImageNet模型的性能从78.89%提高到82.13%,这与NASA论文的发现一致。如图7所示,Super-4的测试图像来自与训练图像不同的图像分布(图6)。第一行显示来自不同合金的显微照片。第(2和3)行显示具有不同蚀刻条件的显微照片,最后一行显示成像质量较差的显微照片[1]。显微镜模型在分离次生析出物时更准确且过度分割较少,如绿色箭头所示。显微镜模型在分割图像的次生和三生析出物方面表现更好。
环境屏障涂层(EBC)分割
如图8所示,EBC数据集的结果与NASA论文的发现一致,其中EBC-1和EBC-2的IoU分数对于显微镜模型和ImageNet模型相似。两种模型的EBC-1和EBC-2分割输出能够区分基底和热生长氧化层,并且在简单的形态学操作后可用于氧化层厚度测量。
图8. (a) 显示EBC数据集中的训练和测试图像示例,其中左侧列代表训练集,右侧列代表测试集。EBC-3的单个训练图像用红色轮廓标出。(b) 显示每个EBC数据集的最佳ImageNet和显微镜模型的分割结果。每个实验的最佳模型CNN/Transformer编码器和IoU分数在分割掩码上方。
EBC-3有一张训练图像(在图8中以红色轮廓显示)。最佳显微镜模型的IoU分数为70.46%,远高于最佳ImageNet模型的61.74%的IoU分数。这也证实了NASA论文的结果,尽管我们的改进更大。ImageNet模型无法区分基底和热生长氧化层,这使得准确测量氧化层厚度变得不可能。
5 讨论
我们已经展示了CS-UNet中的CNN和Transformer编码器在分割性能上优于单独使用CNN编码器的UNet。我们还展示了在显微镜图像上预训练的CS-UNet比在ImageNet上预训练的性能更好,尽管改进的程度与UNet上的观察结果不同。然而,这些比较基于表现最好的模型。根据CNN编码器、Transformer架构和预训练模型的选择,分割性能可能会有显著差异。
在本节中,我们比较了预训练对基于CNN、基于Transformer和混合分割算法平均性能的影响。之后,我们比较了三种类型分割算法的平均性能。我们的结果表明,在显微镜图像上预训练通常对性能有积极影响。我们的混合算法CS-UNet在所有实验中优于UNet,并且在大多数情况下与基于Transformer的算法具有相似或更好的性能。
基于CNN的图像分割
我们首先检查了UNet[15]在三种编码器预训练下的性能。我们包括这一结果,因为CS-UNet的配置仅使用了NASA论文[1]中的35个CNN编码器中的19个。如附录B所示,这19个编码器在至少一个分割任务中具有top-5的平均IoU分数。这一选择减少了公平比较所需的实验数量。
表5. UNet在不同预训练权重初始化时的平均性能。每个条目显示预训练模型的IoU值的均值和标准差。每个测试的最高分数以粗体显示。
当UNet在ImageNet或MicroNet上预训练时的平均性能如表5所示,其中ImageNet→MicroNet模型(即CNN编码器使用ImageNet模型初始化并在MicroNet上微调)在大多数情况下取得了最佳结果。表现最好的CNN编码器的配置如附录C所示,这也表明在MicroNet上预训练在大多数情况下提供了更好的结果。
不出所料,我们的结果与NASA论文[1]的结果基本一致,因为我们使用了他们在MicroNet上的预训练模型。具体来说,使用MicroNet进行预训练提高了一次性学习和分布外学习的性能。由于我们选择了在至少一个实验中具有top-5性能的CNN编码器,IoU分数高于NASA论文中显示的平均分数。事实上,Super-2(少样本学习)的性能在不同预训练模型下基本相同。
基于Transformer的图像分割
接下来,我们在表6中展示了基于Transformer的分割算法(Swin-UNet、HiFormer和TransDeepLabv3+)使用不同预训练配置和Swin-T架构的平均性能。我们通过使用原始或中间Swin-T架构以及使用ImageNet或显微镜预训练模型初始化它们的权重来比较算法。我们的结果表明,算法在使用MicroLite预训练模型时表现良好,并且原始Swin-T架构在一次学习和分布外学习中略微更好。总体而言,使用显微镜图像进行预训练为基于Transformer的分割算法提供了比自然图像预训练更好的结果。
表6. 基于Transformer的分割算法在不同预训练权重初始化时的平均性能。每个条目显示预训练模型的IoU值的均值和标准差。每个数据集的最高分数以粗体显示。
混合图像分割
我们还比较了表7中使用原始或中间Swin-T架构以及使用ImageNet或显微镜模型初始化权重时的混合分割算法CS-UNet的性能。由于CS-UNet使用CNN和Transformer编码器,结果不一,其中在显微镜图像上预训练并不在所有情况下提供更好的性能。CNN编码器的较弱性能减少了当它们在显微镜图像上预训练时Transformer编码器的优势。然而,当我们考虑所有实验的平均IoU分数时,使用显微镜图像进行预训练仍然提供了一些好处。
表7. CS-UNet在每个实验中使用不同预训练权重初始化时的平均性能。表中每个条目显示特定预训练模型的评估IoU指标的均值和标准差。每个数据集的最高准确性分数以粗体显示。
分割网络的比较
最后,我们在表8中比较了三种类型分割算法(UNet、CS-UNet、Swin-Unet、HiFormer和TransDeepLabv3+)在不同预训练模型和Swin-T架构下的平均性能。结果显示,CS-UNet在所有实验中的平均表现优于UNet。尽管基于Transformer的分割算法在一次学习或分布外学习中可能更优越,但它们的性能并不总是始终优于UNet。因此,我们的混合算法CS-UNet似乎是无论预训练模型如何都更为鲁棒的解决方案。
表8. CNN、Transformer和CS-UNet在所有数据集上的平均性能。表中每个条目显示特定方法的IoU评估指标的均值和标准差。
6 结论
在本文中,我们证明了在领域内数据集上训练CNN和Transformer编码器可以为图像分割任务提供更好的特征表示。我们引入了CS-UNet——一种结合CNN和Swin Transformer编码器用于图像分割的混合网络。CS-UNet的编码器从输入图像中提取低级和高级特征,而解码器部分使用Swin Transformer进行图像分割。我们展示了CS-UNet在我们的实验中通常优于UNet,并且与基于Transformer的算法具有竞争性能。CS-UNet的混合性质使其能够有效分割具有长程依赖和高空间分辨率的图像。
我们的预训练数据集规模有限。更大的数据集将通过预训练产生更好的微观结构特征表示,从而提高下游分析任务的准确性。尽管CS-UNet可以保留空间信息并处理长程依赖关系,但它比基于CNN和Transformer的算法计算密集。作为未来的工作,我们将探索在大型显微镜数据集上训练其他Transformer架构,如Focal Transformer和FocalNet。这些架构可能在图像分割方面提供进一步的进展,扩展可用选项和材料分析任务的潜在改进。
附录A Transformer性能
表9显示了基于Transformer的分割算法表现最好的配置。请注意,对于使用MicroLite进行预训练,直接在MicroLite上训练在所有情况下都取得了最佳性能,除了Super-4,其中ImageNet→MicroLite具有最佳性能。
附录B 选择的CNN编码器
表10显示了用于评估CS-UNet和UNet性能的19个CNN编码器。
附录C UNet性能
表11显示了在ImageNet或MicroNet上预训练时为每个数据集产生最佳性能的UNet的CNN编码器。EfficientNet和Se_ResNet系列比旧的编码器系列(如VGG)表现更好。
数据可用性
我们实验的预训练Swin-T模型可在我们的GitHub仓库中获取:https://github.com/Kalrfou/SwinT-pretrained-microscopy-models。
本工作还使用了来自https://github.com/nasa/pretrained-microscopy-models的预训练CNN模型。
作者贡献
Alrfou开发了CS-UNet方法,构思并设计了研究,开发了软件,评估了结果,提供了数据集,并参与了正式分析和原始稿件的撰写。Zhao和Kordijazi评估了结果,参与了正式分析和原始稿件的撰写,并校对和审查了最终稿件。