作者丨AI視界引擎
來源丨AI視界引擎
編輯丨極市平台
論文連結：https://arxiv.org/pdf/2308.13917v1.pdf
遷移學習通過使用在更大資料集上預訓練的參數來初始化深度學習模型，進而提高它們的性能。直覺來說，在域内資料集上進行預訓練時，遷移學習更加有效。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和顯微鏡圖像進行預訓練的這些模型的結果進行了比較分析。
2.3. 結合CNN和Transformer（CS-UNet）
由于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的性能。
4.2. 鎳基高溫合金（Super）分割
圖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分數時，顯微鏡圖像上的預訓練仍然提供了一些好處。
5.4. 分割網絡的比較
最後，作者在表8中比較了三種類型的分割算法（UNet、CS-UNet、Swin-Unet、HiFormer和TransDeepLabv3+）的性能，這些算法的性能是基于不同的預訓練模型和Swin-T架構的平均性能。結果表明，CS-UNet在所有實驗中平均表現都優于UNet。
盡管基于Transformer的分割算法在一次學習或超出分布學習中可能更為優越，但它們的性能并不總是一緻優于UNet。是以，作者的混合算法CS-UNet似乎是更穩健的解決方案，無論采用哪種預訓練模型。
參考
[1]. Transfer Learning for Microstructure Segmentation with CS-UNet: A Hybrid Algorithm with Transformer and CNN Encoders.