CSUNet | 完美缝合Transformer和CNN，性能达到UNet家族的巅峰！

点击下方卡片，关注「集智书童」公众号

迁移学习通过使用在更大数据集上预训练的参数来初始化深度学习模型，从而提高它们的性能。直观来说，在域内数据集上进行预训练时，迁移学习更加有效。NASA最近的一项研究表明，使用在显微镜图像上进行预训练的CNN编码器对显微镜图像的微观结构分割效果更好，而不是使用在自然图像上进行预训练的CNN编码器。然而，CNN模型只能捕捉图像中的局部空间关系。
近年来，像Transformer这样的注意力网络越来越多地用于图像分析，以捕捉像素之间的远距离关系。在本研究中将在显微镜图像上进行预训练的Transformer和CNN模型的分割性能与在自然图像上进行预训练的模型进行了比较。结果部分验证了NASA的研究结论，即在显微镜图像上进行预训练的模型对于不同成像和样本条件下的图像（属于分布外的图像）的分割性能显著提高。
然而，对于One-Shot学习和Few-Shot学习，使用Transformer的性能提升较为有限。作者还发现，在图像分割任务中，预训练的Transformer和CNN编码器的组合始终优于仅使用预训练的CNN编码器。作者的数据集（约有50,000张图像）将NASA数据集的公开部分与作者收集的额外图像相结合。即使使用了更少的训练数据，作者的预训练模型在图像分割方面性能显著更好。
结果表明，Transformer和CNN相互补充，当它们在显微镜图像上进行预训练时，对下游任务更有益。

1、简介

显微镜成像提供了关于物质的真实信息，但要获取有关形态、大小和分布的定量信息需要对显微图进行手动测量，这不仅耗时且劳动密集，还容易产生偏见。材料结构和现象的长度和时间尺度在各个组分之间差异显著，增加了复杂性。因此，建立过程、结构和性能之间的联系是一个具有挑战性的问题。

深度学习（DL）由于其自动提取重要信息的能力而被广泛应用于复杂系统。研究人员已经将DL算法应用于图像分析，以识别结构并确定微结构与性能之间的关系。DL已被证明可以用于材料设计，与基于物理的方法相辅相成。然而，DL需要大量的训练数据，而有限数量的显微镜图像往往会降低其效果。

为了使DL适用于较小的数据集，开发了学习技术，如迁移学习、多保真建模和主动学习。迁移学习使用在更大数据集上预训练的模型的参数来初始化在较小数据集上进行下游任务训练的模型。例如，可以使用在自然图像上预训练的卷积神经网络（CNN）来初始化用于图像分割的神经网络，以提高其精度并减少训练时间。

然而，使用自然图像（例如ImageNet）进行预训练并不理想，因为在自然图像上进行预训练的模型会识别显微图像中不存在的高级特征。最近Stuckner等人的研究展示了使用一个名为MicroNet的显微镜图像数据集进行CNN预训练的优势，该数据集包含超过110,000张图像。他们评估了使用MicroNet进行预训练的CNN编码器对镍基高温合金（Super）和环境障碍涂层（EBCs）的显微镜图像分割的准确性。使用MicroNet进行预训练显著提高了一次性学习和少量样本学习以及不同成分、蚀刻和成像条件的分布外图像的准确性，以IoU（交并比）进行度量。

近年来，名为Transformer的基于注意力机制的神经网络在计算机视觉中被广泛采用。CNN从图像的局部区域中提取特征，使用卷积滤波器捕捉像素之间的空间关系，而Transformer将图像分成多个块（patches）并将它们馈送到基于Transformer的编码器中，以捕捉图像之间像素的远距离关系。因此，CNN和Transformer的组合可能在迁移学习中比单独使用任何一个模型更加有效。

在本文中，作者评估了使用CNN和Transformer编码器组合进行显微镜图像分割的迁移学习。迁移学习方法如图1所示，包含了用于图像分割的编码器-解码器架构。每个编码器将输入图像转换为潜在表示向量，以提取语义信息。解码器将提取的信息映射回输入图像中的每个像素，生成图像的像素级分类。

作者使用了Transformer的一种常见版本，即Swin-Transformer，特别是其精简版本Swin-T，以提高效率。作者的预训练数据集包含约50,000张显微镜图像，分为74类，作者将其称为MicroLite数据集。Swin-T模型可以在微型Lite数据集上进行微调之前，使用在ImageNet上预训练

的模型的权重进行初始化。作者使用了在MicroNet上进行预训练的CNN模型来初始化图1中蓝色编码器的权重，而Swin-T模型用于初始化图1中橙色编码器和解码器的权重。在连接到解码器之前，CNN和Swin-T编码器的输出被融合在一起。

为了评估迁移学习的分割性能，作者比较了使用仅在ImageNet上进行预训练的模型和使用在显微镜图像上进行预训练的模型进行的7个数据集（Super和EBC的子集）上的图像分割的IoU分数。作者的结果表明，尽管一次性学习和少量样本学习的分割准确性得到了改善，但与NASA论文中展示的效果相比，提高并不那么显著。对于分布外图像，使用在显微镜图像上进行预训练的模型仍然显著优于仅在ImageNet上进行预训练的模型。作者还比较了使用CNN、Swin-T以及它们的组合进行分割的性能。作者的结果表明，在大多数情况下，组合优于仅使用CNN，而在某些情况下也优于仅使用Swin-T。

2、本文方法

作者的目标是证明基于Transformer的显微镜图像预训练模型对于图像分割等下游任务是有益的，并且它们比基于CNN的预训练模型更加健壮。为此，作者完成了以下任务。

收集了一个包含约50,000张图像的显微镜数据集（MicroLite）经过预处理，
在MicroLite上预训练了Transformer编码器，并使用它们初始化了几个基于Transformer的分割算法（Swin-Unet、TransDeeplabv3+ 和 HiFormer）以及基于CNN和Transformer编码器的混合分割神经网络（CS-UNet）。
为了证明CS-UNet的优势，作者将CNN-based分割算法的最佳性能与Transformer-based分割算法和CS-UNet进行了比较。这些算法使用NASA团队的7个测试集进行比较，其中CNN编码器在MicroNet上进行了预训练，而Transformer编码器在MicroLite上进行了预训练。
为了评估在领域内数据上进行预训练的优势，作者比较了CS-UNet在预训练于ImageNet和MicroLite时的最佳性能。
为了检查在领域内数据上进行预训练的详细效果，作者比较了具有不同预训练设置的CNN-based分割算法的平均性能。同样，作者还比较了具有不同预训练设置和Transformer架构的Transformer-based和混合分割算法的平均性能。
最后，为了说明作者混合策略的健壮性，作者比较了在所有配置上平均的三种类型分割算法的性能。

2.1 数据集预处理

作者的MicroLite数据集中的图像来自多个来源，包括使用光学显微镜、扫描电子显微镜（SEM）、透射电子显微镜（TEM）和X射线等不同测量技术获得的不同材料和化合物的图像。MicroLite汇集了Aversa数据集、超高碳钢显微图、来自Materials Data Repository的SEM图像以及一些最近出版物的作者的图像中的图像。

Aversa数据集包括10个类别的超过25,000张SEM显微镜图像，每个类别包含不同尺度（包括1、2、10、20um和100、200nm）和对比度的图像。为了正确分类这些图像，作者使用了预训练的VGG-16模型从这些图像中提取特征映射，并使用K均值算法对特征映射进行聚类，以使具有相似特征映射的图像分组到相同的类别中。在预处理步骤之后，作者得到了53个类别。Aversa数据集的作者手动将一小部分图像（1038张）分类到一个分层数据集中，其中这10个类别进一步分为27个子类别。作者的这1038张图像的分类与手动分配的子类别基本一致。需要注意的是，由于作者处理了整个Aversa数据集，所以作者有更多的类别。

总之，MicroLite包括74个类别中标记的50,000张显微镜图像，这些图像经过以下预处理步骤获得。

从图像中删除比例尺等工件。
将图像分割成512×512像素的图块，根据原始图像的大小是否有重叠来进行分割。
进行数据增强以增加数据集的大小。
聚合原始图像、图像块和增强图像，形成最终的数据集。

2.2. 预训练

作者训练了Swin Transformer模型来学习显微镜图像的特征表示，以便将其迁移到分割等任务中。作者评估了两种类型的训练。

对在ImageNet上预训练的模型进行微调，然后使用MicroLite进行微调（表示为ImageNet → MicroLite）。
从头开始使用MicroLite对模型进行预训练（表示为MicroLite）。

分类任务使用了Swin-T，这是Swin Transformer的迷你版本。Swin-T包含两种类型的架构：原始的Swin-T，具有[2,2,6,2]的Transformer块，以及中间网络，具有[2,2,2,2]的Transformer块。

图2显示了Swin-T的原始架构。作者推测，中间网络对于显微镜分析任务可能已经足够，因为较早的层学习角落的边缘和形状，中间层学习纹理或模式，原始模型中较深的网络层学习高级特征，如眼斑和尾附件。原始和中间的Swin-T模型都是从头开始在MicroLite上进行预训练的，其中模型权重是随机初始化的。这两个模型还在ImageNet上进行了预训练，并在MicroLite上进行了微调。

预训练步骤使用AdamW优化器进行30个Epoch，采用余弦衰减学习率调度器，线性热身5个Epoch，Batch-Size为128。初始学习率为，权重衰减为0.05。微调步骤也使用AdamW优化器进行30个Epoch，Batch-Size为128，但学习率降低到，权重衰减降低到。模型一直训练，直到验证分数不再改善，采用了5个Epoch的早停准则。训练数据已使用albumentations库进行增强，包括对比度和亮度的随机变化、垂直和水平翻转、光度扭曲和添加噪声等。

对于下游分割任务，针对每个任务训练了多个模型，包括Swin-Unet、HiFormer和TrasDeeplapv3+。对使用ImageNet和显微镜图像进行预训练的这些模型的结果进行了比较分析。

由于CNN具有固有的局部性，它不能捕获长程空间关系。Transformer被引入以克服这一局限性。然而，Transformer在捕获低级特征方面存在局限性。研究表明，对于像在复杂背景下进行分割这样的密集预测任务，既需要局部信息又需要全局信息。

一些研究人员引入了混合模型，有效地将CNN和Transformer用于图像分割。在混合模型中初始化CNN和Transformer的权重将显著提高性能。因此，作者引入了一种名为CS-UNet的混合UNet，它是一种使用CNN和Transformer的U形分割模型。如图3所示，该方法包括编码器、瓶颈、解码器和跳跃连接。

编码器组合了CNN编码器和Swin-T编码器，其中CNN用于提取低级特征，而Swin-T用于提取全局上下文特征。Swin-T编码器对输入图像分成不重叠的块，应用自注意机制以捕获全局依赖关系。编码器从不同尺度捕获整个图像的长程依赖关系和上下文信息。受到TFCN（Transformers for Fully Convolutional dense Net）和Lightweight Swin-Unet的启发，两个连续的Swin-T块中的多层感知器（MLP）被Residual Multi-Layer Perceptron（ResMLP）替换。

如图4所示，ResMLP用于减少传输过程中的特征损失，并增加编码器提取的上下文信息。

ResMLP如图5所示，由2个GELU非线性层、3个线性层和2个dropout层组成。CNN编码器通过一系列卷积层处理输入图像，逐渐减小空间维度，同时提取分层特征。在此过程中，编码器在早期层捕获低级特征，在较深层次上捕获高级语义特征。

为了融合来自两个编码器的信息，跳跃连接将CNN编码器和Swin-T编码器的特征图与相应的解码器层进行连接。为确保CNN和Swin-T编码器的特征维度之间的兼容性，需要在融合它们之前对维度进行归一化。这是通过将从CNN块获得的特征经过线性嵌入层来实现的，该层将特征图从（B，C，H，W）Reshape状为（B，C，H×W），其中B，C，H，W分别是Batch-Size，通道数，高度和特征图的宽度。平铺的特征图被转置以交换最后两个维度，结果是形状为（B，H×W，C）的形状，然后与从Swin-T编码器提取的特征融合。

通过融合来自不同编码器路径的信息，跳跃连接使解码器能够既受益于CNN编码器捕获的局部空间细节，又受益于Swin-T编码器捕获的全局上下文。

解码器类似于Swin-Unet的解码器，它使用了扩展图块层来通过重塑相邻维度的特征图来上采样提取的深层特征，从而有效地实现了2×的上采样。此外，它将特征维度降低到原始维度的一半。这使解码器能够以更高的空间分辨率重建输出，同时降低特征维度以进行高效处理。

最终的扩展图块层进一步进行4×的上采样，将特征图的分辨率恢复到与输入分辨率（W×H）匹配。然后，应用线性投影层对这些上采样特征进行操作，以生成像素级的分割预测。

编码器部分可以使用不同的CNN系列，如EfficientNet、ResNet、MobileNet、DenseNet、VGG和Inception。作者使用MicroNet初始化CNN权重和MicroLite初始化Transformer权重。

3、结果

预训练的Swin-T模型用于对74个不同类别的显微镜图像进行分类。Swin-T模型要么是在ImageNet上进行预训练的（具体来说是imageNet1K数据集），然后在MicroLite上进行微调，要么是使用随机参数进行MicroLite训练。当验证准确度在经过5个Epoch后不再改善时，训练会停止。模型准确度使用top-1和top-5准确度来评估。top-1准确度测量了正确标签被预测的测试样本的百分比，而top-5准确度测量了前五个预测中正确标签的百分比。

如表1所示，从头开始训练的Swin-T模型需要更长时间才能收敛。具体而言，原始Swin-T模型需要23个Epoch，而中间版本需要19个Epoch。相比之下，经过ImageNet预训练然后在MicroLite上微调的Swin-T模型收敛更快。原始Swin-T模型仅需要13个Epoch，而中间版本需要12个Epoch。

平均而言，使用ImageNet权重初始化的模型的收敛速度约快40.16%，比随机初始化的模型要快得多。在ImageNet预训练后在MicroLite上微调的原始Swin-T模型达到了84.63%的top-1准确度。总体而言，经过ImageNet预训练并在MicroLite上微调的Swin-T模型具有更高的准确性和更快的收敛速度。

4.1. 显微图像分割

为了评估Swin-T模型能够如何提取特征表示，预训练模型被用于初始化分割任务的模型。为了与NASA的研究进行比较，作者使用了从两种材料（镍基高温合金（Super）和环境屏障涂层（EBC））派生的相同的7个显微镜数据集。EBC数据集有两个类别：氧化物层和背景（非氧化物）层，Super数据集有3个类别：基体、次生和三级。

每个数据集分割中的图像数量在表2中显示。Super-1和EBC-1包含各自材料的完整数据集。Super-2和EBC-2只包含训练集中的4张图像，以评估模型在少量样本情况下的性能。Super-3和EBC-3只包含训练集中的1张图像，以评估一次学习期间的性能。Super-4包含在不同成像和样本条件下拍摄的测试图像。

EBC和Super数据集的增强方式类似于NASA的研究，包括:

将图像随机裁剪为512×512像素、随机更改对比度、亮度和伽马，并添加模糊或锐化。
EBC数据集进行了水平翻转，Super数据集进行了随机的垂直和水平翻转以及旋转。
训练使用Adam优化器，初始学习率为，直到验证准确性在30个Epoch内没有改善为止。之后，训练继续使用学习率为，在没有任何验证改善的情况下触发了额外的30个Epoch的早停止。
由于数据集不平衡，损失函数设置为平衡交叉熵（BCE）和dice损失的加权和，其中BCE占70%的权重。

CS-UNet架构是一种灵活的模型，可以使用不同的CNN系列进行训练，并初始化不同的预训练模型。表3显示了用于训练CS-UNet模型的不同预训练权重的各种组合。第二列显示了初始化Swin-T编码器的预训练权重，第三列显示了初始化CNN编码器的预训练权重。在最后一列，作者使用术语“显微镜学”来指代CNN编码器使用MicroNet进行训练，Transformer编码器使用MicroLite进行训练的情况。还可以使用其他组合的预训练权重来训练CS-Unet模型。例如，Swin-T编码器可以使用MicroLite权重进行初始化，而CNN编码器可以使用ImageNet→MicroNet权重进行初始化。CS-UNet架构的灵活性使研究人员可以尝试不同的预训练权重组合，以找到适合其特定任务的最佳组合。

表4比较了在MicroNet上预训练的UNet++/UNet，在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表现出显著的改进。

总体而言，表4的结果表明，CS-UNet是图像分割任务的一种有前途的方法。在所有实验中，CS-UNet与UNet++/UNet相似或明显更好，并且大多数实验中优于Transformer。值得注意的是，MicroLite的大小约为MicroNet的一半。尽管如此，Transformer + MicroLite的性能与UNet++/UNet + MicroNet的性能相媲美或更好。附录A显示了表4中显示的性能最佳的Transformer模型的配置。最佳性能的CS-UNet模型的配置显示在下一节中，作者将比较在显微镜图像和ImageNet上预训练时CS-UNet的性能。

图6和图7比较了CS-UNet在显微图像和ImageNet上预训练时在Super数据集上的最佳性能。对于Super-1和Super-2，两种预训练模型的IoU分数相似，而对于Super-3，显微模型的IoU分数显著高，达到了93.5%，而ImageNet模型的IoU分数仅为87.01%。这一结果与NASA的研究不同，其中Super-2的性能也有所提高。似乎对于CS-UNet的增强能力而言，在领域内数据集的好处在一次学习中比在少量学习中更显著。ImageNet模型未能识别出许多暗对比度图像中的三级析出物，如橙色三角所示。ImageNet模型还对一些次生析出物进行了过度分割和合并，如绿色箭头所示。

对于包含不同成像条件图像的Super-4，显微模型将ImageNet模型的性能从78.89%提高到了82.13%，这与NASA的研究结果一致。如图7所示，Super-4的测试图像来自不同的图像分布，与训练图像不同（图6）。第一行显示了来自不同合金的显微图像。第2和第3行显示了具有不同腐蚀条件的显微图像，最后一行显示了成像较差的显微图像。显微模型对于分离次生析出物的准确性更高，过度分割较少，与ImageNet模型相比，这些差异用绿色箭头标示。显微模型对于分割图像的次生和三级析出物性能更好。

4.3. 环境屏障涂层（EBC）分割

如图8所示，EBC数据集的结果与NASA的研究结果一致，EBC-1和EBC-2的IoU分数在显微镜和ImageNet模型上都相似。对于EBC-3，有一个训练图像（在图8中用红色轮廓标出）。最佳显微模型的IoU分数为70.46%，远高于最佳ImageNet模型的61.74%的IoU。这也证实了NASA的研究结果，尽管作者的改进要大得多。ImageNet模型不能区分基板和热长氧化层，这使得准确测量氧化物厚度成为不可能。

5、讨论

作者已经展示了CS-UNet中的CNN和Transformer编码器比UNet中的CNN编码器单独提供更好的分割性能。作者还展示了在显微镜图像上进行预训练比在ImageNet上进行预训练对CS-UNet的性能有更好的影响，尽管改进程度与UNet上观察到的不同。然而，这些比较是基于性能最佳的模型。根据CNN编码器的选择、Transformer体系结构和预训练模型，分割性能可能会有很大的差异。

在本节中，作者比较了对CNN、Transformer和混合分割算法的平均性能的预训练的影响。之后，作者比较了三种类型的分割算法的平均性能。作者的结果表明，在显微镜图像上进行预训练通常对性能有积极影响。作者的混合算法CS-UNet在所有实验中优于UNet，而在大多数实验中性能与基于Transformer的算法相似或更好。

5.1. 基于CNN的图像分割

首先，作者检查了UNet [15]在3种编码器预训练方式下的性能。作者包括这个结果是因为CS-UNet的配置仅使用了NASA论文[1]中的35个CNN编码器中的19个。如表10所示，这19个编码器在至少一项分割任务中具有前5位的准确性。这个选择减少了需要进行公平比较的实验数量。

UNet在使用ImageNet或MicroNet进行预训练时的平均性能如表5所示，表明ImageNet→MicroNet模型（即CNN编码器初始化为ImageNet模型并在MicroNet上进行微调）在大多数情况下取得了最佳结果。性能最佳的CNN编码器的配置显示在表11中，在大多数情况下，MicroNet上的预训练提供了更好的结果。

毫不奇怪，作者的结果在很大程度上与NASA的研究结果一致，因为作者使用了他们在MicroNet上的预训练模型。具体来说，MicroNet上的预训练提高了一次学习和超出分布的性能。由于作者选择了在至少一次实验中具有前5名表现的CNN编码器，IoU分数比NASA论文中显示的平均分数更高。事实上，在Super-2（少样本学习）上的性能基本上与不同的预训练模型相同。

5.2. 基于Transformer的图像分割

接下来，作者在表6中展示了使用不同配置的预训练和Swin-T架构的基于Transformer的分割算法（Swin-UNet、HiFormer和TransDeepLabv3+）的平均性能。作者通过使用原始或中级Swin-T架构，并使用ImageNet或显微镜预训练模型来比较算法。作者的结果表明，这些算法在MicroLite预训练模型下表现良好，而原始Swin-T架构在一次学习和超出分布学习时稍微好一些。

总体而言，与在自然图像上进行预训练相比，在显微镜图像上进行预训练为基于Transformer的分割算法提供了更好的结果。

5.3. 混合图像分割

作者还在表7中比较了作者的混合分割算法CS-UNet的性能，当它使用原始或中级Swin-T架构以及从ImageNet或显微镜模型初始化的权重时。由于CS-UNet使用了CNN和Transformer编码器，结果各异，预训练显微图像并不总是提供更好的性能。CNN编码器性能较弱降低了在显微镜图像上进行预训练的Transformer编码器的优势。然而，当作者考虑所有实验中的平均IoU分数时，显微镜图像上的预训练仍然提供了一些好处。