基于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 模型的计算复杂度相对较高，这可能会限制其在资源受限的环境下的应用；其次，尽管模型在当前的数据集上表现良好，但在更多样化和复杂的岩石图像数据集上的适应性和鲁棒性仍需进一步验证和优化。

未来研究的发展趋势应当聚焦于以下几个方向：一是继续优化模型结构，以降低计算复杂度和提高运算效率，使模型更适合于不同的应用场景；二是扩展和丰富训练和测试用的岩石图像数据集，以增强模型的泛化能力和鲁棒性；三是探索更高效的数据增强和迁移学习策略，以进一步提升模型在实际应用中的性能。通过对这些方向的深入研究，有望进一步推动岩石岩性智能识别技术的发展，为地质学领域带来更多的创新和应用价值。

总体而言，本研究在岩石岩性智能识别领域取得了重要进展，为后续的研究和应用提供了宝贵的经验和参考。期待未来在这一领域能够取得更多突破性的成果，为地质调查和研究工作提供更加强大的技术支持。

注：本文通讯作者为何月顺。