基于Swin Transformer 的巖石巖性智能識別研究

韓鑫豪，何月順，陳 杰，熊凌龍，鐘海龍，杜 萍，田 鳴

（1.東華理工大學信息工程學院，江西南昌 330013；2.江西省放射性地學大數據技術工程實驗室，江西南昌 330013；3.鄭州市公安局網監支隊，河南鄭州 450000）

0 引 言

巖石識別是地質調查的基礎性工作。在野外地質調查中，地質工作者會根據巖石的顏色、結構構造、礦物成分等辨識巖石的巖性。隨著計算機視覺和深度學習技術的飛速發展，巖石紋理圖像的自動識別和分類已經成為地質學中一個熱門的研究方向。近年來，許多國際和國內的研究團隊都投入大量的精力進行此類研究，以期獲得更高的識別精度和更穩健的分類效果[1]。

自AlexNet[2]在ImageNet[3]上取得重大突破后，卷積神經網絡（Convolutional Neural Network, CNN）[4]便一直引領著計算機視覺領域的研究。隨著深度學習的發展，各種卷積神經網絡在巖石識別分類方面取得顯著進展。卷積運算可以有效地替代人工提取特征的方法，從而更準確地獲取圖像紋理與色彩中的巖石圖像信息，精準識別巖石類型。文獻[5]基于Iception-v3 深度卷積神經網絡模型，建立了基于巖石圖像的遷移學習模型，雖然精確度不理想，但是提高了數據收斂的速度；文獻[6]采用卷積神經網絡提出了一種基于巖石圖像深度學習的巖性智能識別方法，并均取得了非常好的識別效果，結果表明基于CNN 的網絡在圖像識別上提取底層特征方面有較大的優勢，但是仍然無法避免缺失全局特征的問題。對于卷積神經網絡在全局特征缺失的問題上，文獻[7]在RetinaNet 的基礎上對特征融合進行改進，并增加空間和通道注意力，有效提升了對小目標的識別準確度。在上述文獻的基礎上，文獻[8]引入實例級去噪模塊，實現了更好的特征提取，進一步提高了圖像識別準確度。文獻[9]將水平感興趣區域轉換為旋轉感興趣區域（Rotation Region-of-Interest, RRoI），使用RRoI 池化進一步修正旋轉檢測框的坐標信息，有效提升了模型預測結果的準確率。文獻[10]提出的基于巖石目標檢測的巖性智能識別技術，其核心是采用基于Faster R-CNN 和YOLO4 的巖石檢測網絡在圖像中找出所有巖石目標，并確定它們的類別和位置。這樣不僅可以確定巖石的位置信息，同時也能進一步提高巖石巖性識別的準確率。

盡管這些研究都取得了很好的成果，但目前所使用的模型和技術在處理復雜、多變的巖石紋理圖像時仍然存在一些局限性。例如，卷積神經網絡雖然在許多視覺任務中展現出強大的性能，但其固定的感受野和局部的處理方式可能限制了其在處理巖石紋理圖像時的表現。目前，文獻[11]作為一個新型的深度學習網絡，在許多計算機視覺任務上都展現出了超過傳統卷積神經網絡的性能。自適應的感受野和全局的處理方式使其更適合處理紋理豐富、結構復雜的巖石圖像。針對巖石圖像的特點，研究者們嘗試對Transformer 進行改進。文獻[12]使用Transformer 進行圖像識別，并使用深度可分離卷積降低了計算復雜度。文獻[13]采用Swin Transformer 進行圖像分類，并且使用圖像增強技術提升模型的魯棒性。文獻[14]使用Transformer 編碼器在基于標記的緊湊時空對上下文進行建模，并將學習到豐富的上下文信息標記反饋到像素空間進而獲得更大的感受野。文獻[15]綜合CNN 和Transformer 中自注意力的優點，更加注重全局和局部信息的聚合，提升了巖石圖像識別分類性能。受文獻[16]影響，文獻[17]提出AO2-DETR 框架，避免了大量復雜的前/后處理。文獻[18]使用結構化的Transformer 提高特征收集范圍，并設計空間前饋神經網絡解決了Transformer 檢測中局部空間感受野不足的問題。文獻[19]使用預訓練好的CNN 提取特征，并利用Transformer 構建針對圖像的特征金字塔。文獻[20]在特征金字塔中嵌入Transformer 塊，以增強特征融合能力。

雖然上述方法在巖石圖像識別領域都取得了一定的提升，但是Transformer 在巖石圖像識別中的直接應用仍存在很多問題。例如受困于之前研究中的算法模型和數據集等原因，對于復雜的非線性問題難以進行處理；又比如在脫離樣本數據之后，無法對其他環境中的巖石圖像樣本進行準確的巖性識別，缺乏泛化性能。

鑒于此，本文引入全新的改進Swin Transformer[21]網絡來進一步研究巖石紋理圖像的自動識別和分類，期望能夠進一步提高識別的精度和穩健性。本文旨在探討Swin Transformer 網絡在巖石紋理圖像識別和分類上的應用，并與傳統的模型進行比較，驗證其在此類任務上的優越性。

1 基于巖石紋理增強的Swin Transfomer 網絡

1.1 AugMix 數據增強算法

在本文架構的優化Swin Transformer 網絡中的數據預處理階段中，使用了AugMix[22]數據增強算法對巖石圖像樣本進行數據增強以達到數據預處理的目的。在數據預處理的方法中，數據增強技術是目前可以將擴增的圖像樣本盡可能接近真實分布的一種重要方法。

為了提高對巖石圖像樣本識別的泛用性，并且提升對于巖石紋理的特征提取，本文引入了AugMix 數據增強算法對數據集進行處理。AugMix 首先通過多樣化增強對于給定的輸入圖像應用多種不同的增強操作（如旋轉、平移、剪切等）以生成多個增強版本的圖像；之后再將這些增強的圖像使用混合權重線性組合以產生最終的增強輸出，這種混合過程有助于創建更多樣化和具有挑戰性的樣本，因為結果圖像將包含來自不同增強的組合特征；最后AugMix 引入了一致性損失，訓練模型確保對混合增強樣本的預測與其對原始圖像和單獨增強版本的預測保持一致。在AugMix 中，一致性損失使用Jensen-Shannon Divergence（JSD）度量。JSD 被用來衡量模型對于原始圖像、增強圖像以及混合圖像之間預測的一致性。具體來說，假設模型的預測概率分布分別為p（原始圖像）、q1,q2,…,qn（n個增強圖像），JSD 定義為：

式中：KL 是Kullback-Leibler 散度；每一個mi是原始圖像的預測p和增強圖像qi的預測之間的平均值。

在AugMix 的上下文中，模型在每次迭代中使用JSD最小化原始圖像和它的增強版本之間的預測差異，這樣做是為了確保模型對于不同的數據擾動或增強具有魯棒性，從而提高其泛化能力。通過AugMix 在ImageNet-2012、CIFAR-10 和CIFAR-100 等數據集進行實驗表明，AugMix 數據增強算法可以改進網絡架構的泛化能力，并提升網絡模型對圖像樣本識別的魯棒性。

1.2 遷移學習技術和Swin Transfomer 網絡

本文提出的優化Swin Transformer 中的另一關鍵技術是：采用遷移學習在增強巖石的數據集上訓練Swin Transformer 模型遷移學習技術，使模型在訓練開始前能擁有更高的初始性能，訓練過程中模型性能提升的速率更快，訓練所得模型擁有更好的泛化能力與魯棒性[23]。

鑒于傳統Transformer 中逐像素計算的自注意力機制計算量龐大，且其中的位置編碼操作無法捕獲到圖像中的局部相關性和整體結構信息，Swin Transformer 提出了分層結構和移位窗口機制，有效減少了計算量并提升了檢測的準確度。Swin Transformer 的具體結構如圖1 所示。

圖1 Swin Transformer 具體結構圖

將輸入大小為H×W×3 的圖像通過標記分割塊分割成相同大小的塊以送入后續處理。一共被分為4 個階段，每個階段中都包含兩部分，除了第一個階段由一個線性輸入層和一個Swin Transformer 塊構成之外，其余三個均由一個標記合并塊和一個Swin Transformer 塊構成。其中，標記合并塊類似于池化操作，但是不會造成信息的損失。經過每個階段處理后分辨率都變為原先的一半，而通道數則變為之前的兩倍。

圖2 為Swin Transformer 塊的詳細結構，可以看到與Transformer 塊結構類似，不同之處在于將原始的多頭自注意力（Multi-Head Self Attention, MSA）換成了窗口多頭自注意力（Window Multi-Head Self Attention, W-MSA）和移動窗口多頭自注意力（Shift Window Multi-Head Self Attention, SW-MSA）?？紤]到MSA 在全局內逐像素計算的元余性，W-MSA 僅在一個小窗口內進行Transformer 操作。假設每個窗口包含H×W個小塊，那么二者的計算復雜度分別如公式（1）和公式（2）所示，可以看出W-MSA 大大減少了計算復雜度。此外，考慮到窗口的設置導致了不同窗口間的信息無法獲取的問題，進一步提出SW-MSA。首先使用移動窗口實現跨窗口的特征提取，然后使用循環移位操作合并窗口以進行批處理，并利用掩碼操作消除不相關部分的關聯性，實現了高效的檢測性能。

圖2 Swin Transformer 塊的結構圖

1.3 改進的Swin Transformer 主干網絡

1.3.1 基于融合AugMix 算法和Swin Transfomer 網絡的巖石巖性識別網絡結構

本文提出了一種以Swin Transformer 為骨干網絡的巖石圖像分類網絡，該模型的具體結構組成如圖3 所示。首先，利用AugMix 算法對已有的數據集進行數據增強；然后，將ImageNet 上預訓練的基于巖石紋理增強的Swin Transformer 模型在增強數據上進行微調；最后，通過Layer Norm 層、平均池化層、全連接層和Softmax 層對巖石樣本進行分類。

圖3 融合AugMix 算法和Swin Transfomer 網絡的巖石巖性識別網絡結構

1.3.2 局部增強Swin Transfomer 主干網絡

對于紋理復雜、特征難以捕捉的巖石圖像分類，Swin Transformer 中仍然沒有很好地對大量空間上下文信息進行編碼。為了解決這個問題，本文改進Swin Transformer 并提出了局部增強Swin Transformer 主干網絡LEST，以更好地提取巖石圖像特征。該主干網絡的具體設計如圖4 所示?？梢钥吹?，該主干網絡一共包含4 個階段，每個階段由一個標記合并塊（第一個階段為線性嵌入塊）、n個空間局部感知塊和n個Swin Transformer 塊組成。其中，n表示對應階段的塊數。Swin Transformer 中一共提供了4 種不同大小的模型，考慮到巖石圖像的尺寸及模型計算量問題，本文中僅使用Swin-T 進行改進，即對應每個階段n的數值分別為2、2、6、2。

圖4 局部增強Swin Transfomer 主干網絡

1.3.3 空間局部感知模塊

為了加強網絡對巖石圖像中局部相關性和結構信息的提取能力，本文融合空洞卷積和殘差連接提出了一種空間局部感知塊。本文將這一模塊插入到每一個Swin Transformer 塊之前，并與Swin Transformer 塊一起在每個階段中重復對應的次數，其具體結構如圖5 所示?？紤]到Swin Transformer 中的數據輸入格式與傳統卷積神經網絡的不同，首先對原始輸入數據格式進行調整，假設原始數據輸入為(B,H×W,C)，調整后格式為(B,C,H,W)，將這一特征作為殘差連接的一個分支，另一個分支則首先進行一次3×3 的空洞卷積；然后送入一個GeLU 函數激活；最后將兩個分支相加，以擴大感受野并提取到巖石圖像中更多的空間局部信息。與傳統卷積操作相比，空洞卷積可以在不損失圖像信息的情況下，擴大感受野范圍（普通3×3 卷積的感受野為3×3，相同內核大小下擴張率為2 的空洞卷積感受野為5×5），從而更好地在不同尺度上對更大范圍的上下文信息進行編碼。

圖5 空間局部感知塊

1.3.4 Dropout 函數

為了使模型對單個神經元的依賴性減小，進而增強模型的泛化能力。在Swin Transformer 算法中的Swin Transformer 塊中對多層感知機（MLP）添加Dropout 層。Swin Transformer 塊的結構圖如圖3 所示。使用Dropout可以防止模型在訓練數據上過度擬合，從而在測試數據上達到更好的性能。輸入首先經過全連接層（Linear），然后通過ReLU 激活函數進行非線性變換，接著經過Dropout 層隨機丟棄一部分神經元以防止過擬合，最后輸入到下一個全連接層（Linear）。整體結構如圖6所示。

圖6 MLP 層改進結構圖

2 實驗結果與討論

2.1 巖石樣本數據集

本文使用了1 個采集于東華理工地質博物館的巖石樣本圖像數據集。對該數據集進行數據預處理：主要采用反轉圖片、裁剪巖石樣本中心區域等方法，所得巖石圖像樣本共4 000 張。使用AIChallenger 平臺將圖片數據的屬性存儲在JSON 文件中。本文采用Python 腳本處理JSON 文件和數據集，將巖石圖像從數據集中篩選出來。圖7 展示了部分巖石種類圖片。

圖7 巖石圖像數據樣例

本實驗所使用的巖石圖像樣本及其種類見表1。由于巖石圖像樣本在采集過程中數據集的樣本數目存在不平衡的情況，因此，本文在本網絡的訓練過程中使用單樣本數據增強方法解決類別不平衡造成的模型退化問題，充分發揮分類模型的性能；使用在線數據增強的方式（如水平翻轉、垂直翻轉和旋轉等）提升各類巖石圖像的多樣性，最終使各個數量的巖石樣本達到一致，減小巖石圖像樣本類別不平衡對模型性能的影響。

表1 巖石樣本詳情

2.2 評價標準

為了科學分析本文新提出的模型在巖石分類時的各項具體工作性能，使用經典的性能評價指標[24]，如準確率、精確度、召回率、F1分數進行評價，公式如下：

式中：TP 為真陽性樣本數量，表示準確預測為正確樣本的數量，單位為張；TN 為真陰性樣本數量，表示被預測為正確樣本，實際值也是正確樣本的數量，單位為張；FP 為假陽性樣本數量，表示被錯誤預測為錯誤樣本的數量，單位為張；FN 為假陰性樣本數量，表示被錯誤預測為正確樣本的數量，單位為張。

2.3 實驗環境

實驗使用了1 張GeForce RTX 3060 顯卡Pytorch 深度學習框架、CUDA 11.3 和CuDNN 8.0 來構建網絡模型所需要的訓練環境。為了滿足Swin Transformer 網絡輸入圖像的尺寸要求，將數據集中所有圖像的大小調整為224×224，以方便運算；將巖石圖像數據集按照60%、10%和30%的比例劃分為訓練集、驗證集和測試集以完成模型評估實驗。其中，訓練集和驗證集用于模型訓練以及模型權重參數選擇，測試集用于評估模型的性能。將AugMix 數據增強技術嵌入模型訓練過程，在訓練時完成數據增強，然后將增強數據集直接輸入網絡模型進行訓練。同時，采用遷移學習技術將模型在ImageNet 數據集上預訓練的權重參數用于當前任務，提升模型訓練的效率。

模型訓練期間使用的具體超參數見表2。為保證實驗結果的嚴謹性，本文使用統一的超參數進行網絡模型的訓練，超參數根據文獻和多組實驗確定。在訓練過程中，保存準確率最優的模型，然后用測試數據集計算模型的準確率、精確度、召回率和F1分數，對模型性能進行評估。

表2 模型訓練超參數

2.4 各種模型性能比較

本實驗使用Pytorch 框架構建8 種模型進行實驗。本文評估了所提出的新模型在20 種巖石圖像種類識別任務中的性能，表3 展示了原始的Swin Transformer、預訓練的Swin Transformer 以及本文新提出的優化Swin Transformer和其他7種深度學習訓練模型的分類性能，結果發現，在實驗中的各個指標所呈現出的最優性能均為本文提出的優化Swin Transformer模型。

表3 模型的準確率、精確度、召回率和F1 分數%

實驗結果表明，本文新提出的優化Swin Transformer 結構對巖石圖像樣本識別的準確率為96.40%，明顯優于 VGG16、AlexNet、GoogLeNet[25]、ResNet50、MobileNetV2[26]、ViT 和MobileViT[27]這7 種模型結構的巖石圖像樣本識別準確率，且與其他模型相比，在精確度、召回率以及F1分數這3 個評價指標上，Swin Transformer 占據顯著優勢。 本文還比較了Swin Transformer 模型使用和不使用AugMix 增強算法的效果，與原始數據集相比，使用AugMix 增強算法的網絡模型在準確率、精確度、召回率和F1分數上都有所提高。這些評價指標的提升清楚地表明：使用AugMix 增強算法增強數據集，可以防止網絡模型過擬合，顯著改進網絡架構的泛化能力，而且AugMix 增強算法通過對給定輸入進行梯度懲罰提高了網絡模型的魯棒性。

圖8為模型的訓練準確率與Loss曲線圖，其中圖8a）和圖8b）展示了7 種深度學習模型與本文新模型的準確率和Loss 曲線的比較。

圖8 模型的訓練準確率與Loss 曲線圖

從圖8 中可以看出，本文所提出的新模型識別準確率和收斂速度明顯優于其他模型，網絡的波動幅度更小，比其他模型更穩定。從圖8c）和圖8d）可以看出，經過AugMix 數據增強的模型曲線波動幅度更小，Loss 損失函數曲線明顯低于未使用AugMix 數據增強的Swin Transformer。這一結果再次表明：使用AugMix 數據增強算法能大大提高模型的識別準確率和泛化能力。

3 結 語

本文通過深入探討和實驗驗證，成功提出并實施了一種基于改進Swin Transformer 的巖石識別方法，以解決常規卷積神經網絡在處理紋理多變的巖石圖像時的局限性。 通過引入空間局部感知模塊和結合Transformer 的自注意力結構，本文方法有效地提升了對局部相關性的捕捉能力，顯著增強了模型的識別精度。此外，通過添加Dropout 層和采用AugMix 算法進行數據增強，以及結合遷移學習技術進行預訓練，本文在提高模型泛化能力方面取得了顯著成效，最終實現了96.4%的識別準確率，超越了當前主流的ResNet50、GoogLeNet、VGG16 等網絡模型。

盡管本文取得了一定的成果，但仍存在一些不足和改進空間：首先，在實際應用中，改進后的Swin Transformer 模型的計算復雜度相對較高，這可能會限制其在資源受限的環境下的應用；其次，盡管模型在當前的數據集上表現良好，但在更多樣化和復雜的巖石圖像數據集上的適應性和魯棒性仍需進一步驗證和優化。

未來研究的發展趨勢應當聚焦于以下幾個方向：一是繼續優化模型結構，以降低計算復雜度和提高運算效率，使模型更適合于不同的應用場景；二是擴展和豐富訓練和測試用的巖石圖像數據集，以增強模型的泛化能力和魯棒性；三是探索更高效的數據增強和遷移學習策略，以進一步提升模型在實際應用中的性能。通過對這些方向的深入研究，有望進一步推動巖石巖性智能識別技術的發展，為地質學領域帶來更多的創新和應用價值。

總體而言，本研究在巖石巖性智能識別領域取得了重要進展，為后續的研究和應用提供了寶貴的經驗和參考。期待未來在這一領域能夠取得更多突破性的成果，為地質調查和研究工作提供更加強大的技術支持。

注：本文通訊作者為何月順。