[논문 리뷰] Transfer Learning for Microstructure Segmentation with CS-UNet: A Hybrid Algorithm with Transformer and CNN Encoders

Khaled Alrfou, Tian Zhao외 1명|arXiv (Cornell University)|2023. 08. 26.

Mineral Processing and Grinding인용 수 5

한 줄 요약

이 논문은 전자현미경 영상에서 미세구조 분할을 위해 CNN과 스위닝 트랜스포머 인코더를 융합한 하이브리드 딥러닝 모델인 CS-UNet을 제안한다. 자체 제작한 전자현미경 데이터셋(MicroLite)에서의 사전 훈련은 특히 소수의 샘플(한 개 또는 소수의 샘플)과 분포 외 데이터에 대해 분할 성능을 크게 향상시킨다. 이는 도메인 특화 데이터에서 사전 훈련된 CNN과 트랜스포머 인코더의 융합이 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과 트랜스포머 인코더의 융합 효과를 조사하는 것.
• 도메인 특화 전자현미경 데이터에서 사전 훈련된 하이브리드 모델이 저자료 환경(1-shot 및 소수의 샘플 학습)에서 뛰어난 성능을 내는지 보여주는 것.
• 재료 과학 공동체가 후속 분석 작업을 가속화할 수 있도록 공개적으로 이용 가능한 사전 훈련된 모델을 제공하는 것.

제안 방법

• CS-UNet은 U-Net 기반 아키텍처로, 특징 추출을 위한 이중 인코더(예: EfficientNet, ResNeXt 등 CNN 및 Swin-T)를 포함한다.
• MicroNet(11만 개의 전자현미경 영상)에서 사전 훈련된 가중치로 CNN 인코더를 초기화하고, Swin-T 인코더는 ImageNet 가중치로 초기화한 후 MicroLite 데이터셋에서 미세조정을 수행한다.
• 디코더로 전달하기 전에 CNN과 Swin-T 인코더 간의 특징 융합을 수행한다.
• 니켈 기반 니켈 합금 및 환경 보호 코팅을 포함하는 74개 클래스에 대해 약 5만 개의 전자현미경 영상으로 구성된 커스터마이징된 데이터셋(MicroLite)에서 모델을 훈련한다.
• 일시적 및 소수의 샘플 학습 설정 하에서 일곱 개의 테스트 서브셋(Super-1부터 Super-4, EBC-1부터 EBC-3까지)에서 교차 검증률(IoU)을 사용해 성능을 평가한다.
• ImageNet과 MicroLite에서 사전 훈련된 모델을 비교하고, 개별 및 하이브리드 인코더 구성 모두를 평가한다.

실험 결과

연구 질문

• RQ1 전자현미경 영상에 대해 도메인 특화 데이터에서의 사전 훈련이 ImageNet 사전 훈련보다 분할 성능 향상에 기여하는가?
• RQ2 소수의 샘플 및 1-shot 학습 환경에서 하이브리드 CNN-Transformer 인코더의 성능은 CNN 전용 인코더보다 어떻게 다른가?
• RQ3 전자현미경 데이터에서 사전 훈련된 상태에서 CNN과 스위닝 트랜스포머 인코더의 조합이 개별 인코더보다 더 높은 분할 정확도를 달성할 수 있는가?
• RQ4 도메인 특화 사전 훈련이 다양한 촬영 조건 및 샘플 조건을 가진 분포 외 이미지에 미치는 영향은 무엇인가?
• RQ5 제한된 전자현미경 데이터셋에서 미세조정된 Swin-T 기반 모델의 성능은 CNN 기반 모델과 어떻게 비교되는가?

주요 결과

• 5만 개의 전자현미경 영상으로 구성된 MicroLite 데이터셋에서의 사전 훈련은 특히 1-shot 및 소수의 샘플 학습 설정에서 분포 외 이미지의 분할 성능을 크게 향상시킨다.
• CNN과 Swin-T 인코더를 모두 포함한 하이브리드 CS-UNet 모델은 CNN 전용 모델보다 높은 IoU 점수를 기록하며, 두 인코더 모두 전자현미경 데이터에서 사전 훈련된 경우 최고의 성능를 보였다.
• EBC-1 데이터셋에서 최고의 IoU 점수는 96.59%였으며, 이는 MicroLite에서 사전 훈련된 중간 Swin-T 인코더를 사용한 TransDeepLabV3+ 모델이 ImageNet 전용 사전 훈련을 초월했다.
• Super-3 데이터셋에서 MicroLite에서 사전 훈련한 결과 IoU는 ImageNet 사전 훈련 시 70.74%에서 MicroLite 사전 훈련 시 92.23%로 향상되었다.
• Transformer에 비해 CNN의 경우 도메인 특화 사전 훈련에서 성능 향상이 더 뚜렷했으며, 이는 저자료 환경에서 CNN이 도메인 특화 사전 훈련에서 더 큰 이점을 얻는다는 것을 시사한다.
• 약 5만 개의 훈련 이미지로 제한된 데이터에도 불구하고 사전 훈련된 모델은 최신 기술 수준의 성능를 달성했으며, 이는 재료의 미세구조 분할에서 효과적인 전이 학습을 위해 도메인 특화 사전 훈련이 필수적임을 보여준다.