[论文解读] Transfer Learning for Microstructure Segmentation with CS-UNet: A Hybrid Algorithm with Transformer and CNN Encoders

Khaled Alrfou, Tian Zhao+1 more|arXiv (Cornell University)|Aug 26, 2023

Mineral Processing and Grinding被引用 5

一句话总结

本文提出CS-UNet，一种结合CNN与Swin Transformer编码器的混合深度学习模型，用于显微镜图像中的微结构分割。在自定义显微镜数据集（MicroLite）上进行预训练可显著提升分割性能，尤其在少样本和分布外数据上表现更优，表明在领域特定数据上预训练的CNN与Transformer编码器融合模型，其性能优于仅使用CNN的模型。

ABSTRACT

Transfer learning improves the performance of deep learning models by initializing them with parameters pre-trained on larger datasets. Intuitively, transfer learning is more effective when pre-training is on the in-domain datasets. A recent study by NASA has demonstrated that the microstructure segmentation with encoder-decoder algorithms benefits more from CNN encoders pre-trained on microscopy images than from those pre-trained on natural images. However, CNN models only capture the local spatial relations in images. In recent years, attention networks such as Transformers are increasingly used in image analysis to capture the long-range relations between pixels. In this study, we compare the segmentation performance of Transformer and CNN models pre-trained on microscopy images with those pre-trained on natural images. Our result partially confirms the NASA study that the segmentation performance of out-of-distribution images (taken under different imaging and sample conditions) is significantly improved when pre-training on microscopy images. However, the performance gain for one-shot and few-shot learning is more modest with Transformers. We also find that for image segmentation, the combination of pre-trained Transformers and CNN encoders are consistently better than pre-trained CNN encoders alone. Our dataset (of about 50,000 images) combines the public portion of the NASA dataset with additional images we collected. Even with much less training data, our pre-trained models have significantly better performance for image segmentation. This result suggests that Transformers and CNN complement each other and when pre-trained on microscopy images, they are more beneficial to the downstream tasks.

研究动机与目标

• 通过使用混合CNN-Transformer编码器的迁移学习方法，提升显微镜图像中微结构的分割性能。
• 评估在显微镜特定数据上预训练是否能提升分割性能，相比ImageNet预训练。
• 研究在统一架构（CS-UNet）中结合CNN与Transformer编码器在材料微结构分析中的有效性。
• 证明在低数据场景（如单样本和少样本学习）下，基于领域特定显微镜数据预训练的混合模型可实现更优性能。
• 为材料科学社区提供公开可获取的预训练模型，以加速下游分析任务。

提出的方法

• 提出CS-UNet，一种基于U-Net的架构，采用双编码器：一个CNN（如EfficientNet、ResNeXt）和一个Swin Transformer（Swin-T）用于特征提取。
• 通过迁移学习方法，将CNN编码器初始化为在MicroNet（11万张显微镜图像）上预训练的权重，Swin-T编码器则使用ImageNet权重，随后在MicroLite数据集上进行微调。
• 在将联合表示传递给解码器进行像素级分割之前，对CNN与Swin-T编码器的特征进行融合。
• 在包含约5万张显微镜图像、74个类别的自定义数据集（MicroLite）上训练模型，涵盖镍基超合金和环境障涂层。
• 在七组测试子集（Super-1至Super-4，EBC-1至EBC-3）上，采用单样本和少样本学习设置，使用交并比（IoU）评估性能。
• 比较在ImageNet与MicroLite上预训练的模型，并评估单独与混合编码器配置的性能。

实验结果

研究问题

• RQ1 与ImageNet预训练相比，在显微镜特定数据上预训练是否能提升显微镜图像的分割性能？
• RQ2 在少样本和单样本学习场景下，混合CNN-Transformer编码器的性能是否优于仅使用CNN的编码器？
• RQ3 当在显微镜数据上预训练时，CNN与Swin Transformer编码器的组合是否能实现优于任一编码器单独使用时的分割精度？
• RQ4 领域特定预训练对成像条件和样品条件各异的分布外图像有何影响？
• RQ5 当在有限显微镜数据集上微调时，基于Swin-T的模型与基于CNN的模型性能如何比较？

主要发现

• 在包含5万张显微镜图像的MicroLite数据集上进行预训练，可显著提升分布外图像的分割性能，尤其在单样本和少样本学习设置下表现更优。
• 采用CNN与Swin-T编码器的混合CS-UNet模型，其IoU得分高于仅使用CNN的模型，当两个编码器均在显微镜数据上预训练时性能最佳。
• 在EBC-1数据集上，使用中间Swin-T编码器在MicroLite上预训练的TransDeepLabV3+模型取得了96.59%的最佳IoU得分，优于仅使用ImageNet预训练的模型。
• 在Super-3数据集上，使用MicroLite预训练使IoU从ImageNet预训练的70.74%提升至92.23%（采用TransDeepLabV3+架构）。
• 尽管Transformer模型从显微镜预训练中获得的性能增益较CNN模型更小，表明CNN在低数据场景下从领域特定预训练中获益更多。
• 尽管训练数据有限（约5万张图像），预训练模型仍达到了最先进性能，表明在材料显微结构分割中，领域特定预训练对有效迁移学习至关重要。