整合卷积与高效自注意力机制的图像分类模型

田鑫驰,王亚刚,尹 钟,陈 浩

(上海理工大学 光电信息与计算机工程学院,上海200093)

0 引 言

神经网络是过去几年深度学习领域最重要的研究课题之一,Krizhevsky等人[1]提出的AlexNet使得研究者们重新认识了卷积神经网络(Convolutional Neural Network,CNN).后续在此基础上出现了很多性能更加优秀的网络,如Simonyan等人[2]提出的VGG首次尝试更深的网络结构,旨在叠加更多的卷积层以获得更优秀的特征提取性能.He等人[3]认为简单加深模型深度会造成网络性能衰减,提出了一种瓶颈残差结构来减少特征退化问题.卷积运算依赖于固定的感受野,只能提取到感受野附近的特征信息.依靠多个卷积层可以一定程度上缓解全局感受野的缺失,但大量堆叠卷积或池化层会造成特征退化,使得网络丢失学习到的信息,导致模型性能下降.Woo等人[4]尝试采用更大的卷积核来扩大感受野,却带来了庞大的参数量和计算资源浪费.由于传统的卷积神经网络在全局特征表示方面存在着巨大的困难,导致在准确率性能方面无法实现更大地突破.

2017年谷歌提出的自注意力机制[5]首次被应用于自然语言处理领域并迅速成为其主流架构.随后Dosovitskiy等人[6]提出的Vision Transformer(ViT)首次将自注意力机制应用在计算机视觉,并在许多领域都表现出了显著的优势,例如目标检测[7]、语义分割[8]和图像去雨[9]等.然而ViT还存在着许多尚待解决的问题.一方面由于自注意力机制缺乏局部表征的特性,导致Transformer模型在处理底层信息(如图像的边缘和角落部分)会耗费大量的运算资源.另一方面,ViT缺乏卷积特有的图像归纳偏置属性导致模型参数量和计算复杂度过于庞大.Dai等人[10]已经证明图像归纳偏置对于计算机视觉任务来说是极其重要的.尤其在数据量有限的情况下,ViT的模型性能会出现较大的衰退,原因在于多头注意力中头之间过分灵活的信息交互能力为模型带来了极高的过拟合风险.

基于上述讨论,本文提出了一种高效的图像分类模型.CSNet(Convolution with Self-Attention Network)沿用多阶段分辨率下采样的结构,在模型浅层阶段使用一种特殊的卷积特征提取模块(Convolution Block).它在继承大核卷积大范围感受野的同时,还减少了不必要的计算浪费.在深层阶段利用附带深度可分离注意力机制的全局特征构建模块(Attention Block),将深度可分离卷积融合在自注意力机制中,使得多头注意力的全局依赖性由输入信息和卷积感受野及其邻域的权重共同决定.在数据量有限的情况下有效降低了过拟合的风险.通过以分阶段的方式逐步堆叠Convolution Block和Attention Block,可以很好地整合两种不同类型的特征构建模块来促进表征学习,为图像分类领域提供了一个强大的识别网络.

1 相关工作

1.1 自注意力机制与Vision Transformer

(1)

Vision Transformer(ViT)是第1个将自注意力机制应用在计算机视觉领域的模型.图1所示为ViT的工作原理.首先将输入的一张图片划分为一连串尺寸相同且不重复的二维标记,然后将这些标记与附带一组表征区域位置信息的位置编码一起送入Transformer Block中对全局特征信息进行建模.最后送入分类头进行分类.其中自注意力机制是Transformer Block收集全局信息的主要依据.ViT拥有相较于CNN更加巨大的参数量和计算损耗,且只能应对固定尺寸的输入信息,不适用于一些更加密集的视觉预测任务.Li等人[12]认为ViT中的自注意力机制在编码低层次特征信息方面效率不高.即大部分的计算量都用来提取输入图片低层次的信息,这阻碍了它们在高效学习表征方面的巨大潜力.

图1 Vision Transformer结构Fig.1 Structure of Vision Transformer

1.2 卷积神经网络与自注意力机制的结合

卷积运算有着独特的跨度不变和权重共享的特性,使得卷积神经网络可以应对任意大小的输入信息.卷积操作擅长减少局部冗余,避免不必要的计算量,通过聚合每个感受野像素及其邻域的上下文信息来构建特征图.之后有很多工作尝试将卷积加入自注意力机制中,因为卷积所拥有的图像归纳偏置可以很好地弥补自注意力所缺失的捕获局部特征信息的能力.CvT[13]替换了Self-Attention中原有的线性映射层,输入序列通过3个卷积层分别获得矩阵Q、K、V,并通过跨步卷积来减少获取标记的计算复杂度.MobileViT[14]将ViT中的Transformer Block与特殊设计的纯卷积模块组合起来,在保持高准确率的同时拥有较低的参数量和模型规模,使其可以在移动设备中很好地发挥优势.CPVT[15]和CSwin Transformer[16]则采用了基于卷积的位置编码,降低了模型计算的复杂度,尤其在下游任务中变现优异.CoAtNet[10]在模型前两个阶段采用MBConv[17]来获取图片的底层信息,并且详细讨论了卷积模块与Transformer Block最优的组合方式,表现出极强的模型性能.本文延续了这种卷积模块与Transformer Block互补的整体架构,在模型前期阶段选择利用卷积来构建特征信息.

图2 卷积分解操作Fig.2 Convolutional decomposition

本文将上述卷积分解的思想运用在CSNet前期的卷积特征提取模块中以增大卷积运算的感受野.通过分解后的不同类型卷积层感受野及其邻域的上下文信息聚合来弥补全局感受野的缺失.

2 CSNet

本节将介绍CSNet的详细架构.CSNet延续PVT[19]多阶段分辨率下采样的整体结构,在4个特征提取阶段中分别使用Convolution Block和Attention Block相结合的形式.首先描述了CSNet中的各种特征提取模块的详细架构,其次解释了一些关键模块,如Convolution Block中的卷积等效分解.最后讨论了CSNet与现有的CNN和Transformer模型之间的区别,显示了其在准确性和模型规模上较先前工作的优势.

2.1 Convolution Block

图3 卷积特征提取模块Fig.3 Convolution block

Params=K×K×Cin×Cout
FLOPs=K×K×Cin×Cout×H×W

(2)

而进行上述卷积分解操作后的参数量和计算量分别为:

(3)

其中Cin与Cout分别为图像的输入与输出通道数.CSNet中的Convolution Block将一个21×21的卷积分解为一个7×7的深度可分离卷积、扩张率为3的7×7深度扩张可分离卷积,最后经过一个1×1的点卷积来还原通道数以实现残差连接.由公式(2)和公式(3)可以看出,当卷积核与扩张率越大时,参数量与计算量的减少更加明显,也说明了此种对大核卷积分解方法的高效性.

与ViT中的Transformer Block相比,Convolution Block采用一个名为ConMLP的模块来代替原有的多层感知机(MLP)结构.两者的区别是将MLP中的线性映射层替换为卷积核大小和步距都为1的卷积层.该结构设计的思路源于在ViT的Transformer Block中,输入序列X∈N×C经过多头注意力层后进入多层感知机(全连接层)结构,无需调整输入形状即可连接到下一个特征构建模块.ConMLP若沿用原始MLP结构,在经过分解卷积模块后需要将输入序列X∈H×W×C先展平为X∈N×C,经过全连接层后为了能与下一个卷积特征提取模块有效连接,不得不再将输入形状还原回X∈H×W×C.当模型规模较大,即Convolution Block重复次数和卷积层输出通道数变大时,使用MLP结构会消耗大量的运算资源,使得模型变得难以训练.在一个特征提取模块内使用相同性质的结构可以有效减少计算复杂度.故ConMLP中用1×1的卷积层来代替原始MLP中的全连接层.

Convolution Block同样使用先提高输出通道维度再还原的结构,在每个模块前进行批量归一化(Batch Normalization,BN)来加快模型收敛.在BN后使用GELU激活函数,利用残差连接以减少特征退化.Convolution Block放置在靠前阶段中利用较大感受野的静态卷积更好地捕获局部信息,以减少自注意力机制花费过多的计算资源去构建图像底层特征信息,有效地提升了模型的运算效率.

2.2 DWMSA与Attention Block

Zhou等人[20]认为在多头自注意力机制(Mutil-head Self-Attention,MSA)中随着Transformer Block的增多,一部分头在堆叠多次特征构建模块后会学习到重复的特征信息,造成计算资源浪费和模型性能衰退.注意力退化的根本原因在于多头注意力机制的输出动态地取决于输入表示,即学习到的权重信息无法去影响模型的学习偏好.MSA中过于灵活的头为模型带来动态的全局感受野和捕获不同空间位置复杂信息能力的同时,也伴随着很高的过拟合风险.尤其在数据量并不充足时,这种现象会更加频繁地出现.卷积的输出仅取决于感受野及其邻域的权重信息,这种特性使得卷积神经网络在有限的数据集下的表现要比Transformer模型更好.故本文考虑将卷积的特性引入到多头注意力机制中来缓解注意力崩溃的问题.

对于卷积操作Attention Block主要考虑深度可分离卷积,它可以在保持准确率的前提下有效降低普通卷积的计算量.更为重要的是,深度可分离卷积的输出取决于感受野及其邻域附近的权重,可以帮助MSA摆脱全局注意力依赖.图4所示为本文提出的深度可分离注意力机制(DepthWise Mutil-head Self-Attention,DWMSA).具体来说,对于一个二维输入标记X∈N×C,沿空间维度还原成原始输入图像X∈H×W×C后,送入深度可分离卷积提取卷积感受野及其邻域特征信息,同时加入注意力机制所或缺的卷积图像归纳偏置.深度可分离卷积的加入使得多头注意力机制拥有了捕获底层特征相关的位置信息的能力.在特征构建过程进行到更深层次时,DWMSA对于位置信息的获取将更加显著.如果在获取输入标记前没有加入深度可分离卷积,整个注意力特征构建模块的输出将全部动态地取决于输入信息,无法区分局部附近特征关系的差异.如语义分割任务[21](天空和大海中的相同蓝色斑块会被分割为同一类别).为了防止原始二维标记内部信息被破坏,在深度可分离卷积运算之后将三维图像重新展平为X∈N×C,使用层归一化(Layer Normalization)处理.最后通过线性映射分别得到矩阵Q、K、V.为了缓解每个头之间学习到相同的特征信息,减少注意力崩溃的问题,在Softmax之后使用一个1×1的普通卷积来加强各个头之间的信息联系与交互.最后与V相乘获取注意力图.深度可分离注意力机制的表达式可以为:

图4 深度可分离注意力机制Fig.4 DepthWise mutil-head self-attention(DWMSA)

(4)

上式中BN表示Batch Normalization,Conv1×1表示卷积核大小为1的普通卷积运算.

图5(a)图和图5(b)图分别为ViT中的Transformer Block和CSNet中的Attention Block.两种特征构建模块唯一的区别是将Transformer Block中的多头注意力机制(MSA)替换为更加高效的深度可分离注意力机制(DWMSA).仍然保留线性MLP来扩展通道维度,在每个模块前使用Layer Normalization且利用残差连接来防止特征退化.将Attention Block放置在CSNet的深层阶段中可以最大化发挥深度可分离注意力机制的全局特征信息提取的能力,还可以通过深度可分离卷积捕获局部空间信息的特性为自注意力机制加入所需的卷积特性以提高模型的泛化能力和准确率.最终Attention Block的输出表达式为:

图5 注意力特征提取模块Fig.5 Attention block

y=x′+MLP(LN(x′)) x′=x+DWMSA(LN(x))

(5)

其中LN代表Norm模块,即Layer Normalization操作.

2.3 CSNet整体结构

图6所示为CSNet的整体结构.图中L表示对应特征提取模块重复的次数,FC表示最后用来分类的线性映射层.模型整体沿用多阶段分辨率下采样的架构,共分为4个不同类型的阶段(Stage)来分别构建特征信息.前两个阶段(模型浅层阶段)采用Convolution Block,后两个阶段(模型深层阶段)采用Attention Block进行特征提取.下采样模块采用一个3×3的普通卷积层将分辨率下采样的同时提高输出通道数.除第1个下采样模块的下采样率(即步距S)为4以外,其余阶段S均设置为2.Convolution Block中的卷积分解模块和Attention Block中的DWMSA都采用预激活结构[22].与ViT模型相比,CSNet放弃了使用原始位置编码.原因在于Attention Block中的DWMSA利用深度可分离卷积的特性在感受野及其邻域附近已经提取到了部分权重信息,其中包含了原始图像的位置信息.同样也没有选择增加一个额外的分类标记,而是利用类似卷积神经网络的结构将全局平均池化层应用于最后阶段的输出,以获得简单的输出表示.

图6 CSNet结构Fig.6 Structure of CSNet

为了适应不同规模的应用场景,根据每个Stage的模块重复次数和输出通道数的不同实例化了3种不同模型容量的CSNet,分别名为CSNet-Tiny、CSNet-Small和CSNet-Base.表1所示为3种模型变体的详细架构.其中C表示每个Stage的输出通道数.H表示多头注意力机制中头的个数.R表示不同特征提取模块中ConMLP和MLP输出通道维度的扩展倍数.对于不同规模的CSNet模型,特征构建阶段统一使用C-C-A-A的模块组合方式,其中C和A分别表示本文提出的Convolution Block和Attention Block.不同模块组合方式的性能对比将在3.4.2节中进行详细说明.

表1 不同规模的CSNet Table 1 Different sizes of CSNet

3 实验与分析

在这一节中,实验分别在ImageNet-1K和CIFAR-10数据集上评估了CSNet模型的图像分类性能.随后在提出的Convolution Block中尝试不同卷积核大小和扩张率的搭配来寻找较为合适的卷积分解运算.最后综合模型的参数量和准确率,找到了不同模块之间最佳的组合方式.

3.1 实验环境及模型训练

实验全部在PyTorch框架下进行训练.对于ImageNet-1K数据集来说,训练过程使用动量大小为0.9的AdamW优化器,总共训练300个epochs,使用余弦衰减学习率策略[23]和20个epochs的线性预热阶段(Linear warm-up[24]).模型使用交叉熵损失函数(Cross Entropy Loss).交叉熵损失函数的表达式为:

(6)

训练使用总批量大小(Batch Size)为1024(训练过程使用4个GPU并行,每个GPU上有256幅图像).特别的,训练CSNet-Base时Batch Size调整为512.所有训练模型初始的学习率设置为0.002,权重衰减为0.05.数据增强策略使用简单的随机水平翻转,随机擦除.使用标签平滑策略[25]对网络进行正则化.整个训练过程默认使用分辨率为224×224的图像.对于CIFAR-10数据集的训练,初始学习率设置为0.02,权重衰减为0.0001.采用动量为0.9的AdamW优化器,批量大小为128,总共训练150个epochs.

3.2 数据集介绍

CIFAR-10是一个用于图像分类的小型数据集.该数据集共有6万张分辨率大小为32×32的彩色图像.一共有10个类别,如猫、汽车和飞机等.其中训练集包含5万张图像,测试集包含1万张图像.

ImageNet-1K数据集共有1000个分类类别,其中包含约130万张训练图像和5万张验证图像.作为图像分类领域最为广泛使用的公共数据集,在ImageNet-1K上测试出的实验结果通常被作为模型普遍的性能指标.在训练过程中将ImageNet-1K数据集中所有图像的输入分辨率调整为224×224.

3.3 模型性能对比实验

在这一部分中,实验将CSNet的模型性能与许多先进的分类模型分别在ImageNet-1K和CIFAR-10数据集上进行了对比.

在ImageNet-1K上的对比结果如表2所示.对比的模型包含了卷积模型、Transformer模型和CNN与Transformer相结合的模型.在模型规模较小的情况下,CSNet-Tiny比ResNet-18[3]、PoolFormer-S12[26]和T2T-ViT-12[27]的准确率分别高出8.0%,、0.6%和1.3%,同时参数量和FLOPs都有不同程度的减少.在模型规模较大时,CSNet相较于卷积模型有了明显的优势.如CSNet-Small比ResNet-101的准确率高出4.4%,而相应的参数量和FLOPs则分别下降了12.1%和12.7%.CSNet-Small与RegNetY-8G[28]相比则在FLOPs方面上表现出优势.对于Transformer模型来说,CSNet也表现出了显著的性能优势.CSNet-Small的准确率在比PVT-M[18]提高了0.6%的情况下,依然能保持较小的参数量和计算复杂度.当需要使用更大规模的模型应对复杂场景时,CSNet-Base也能展示出较Transformer模型更加优秀的性能表现.如CSNet-Base的准确率比DeiT-B[29]和T2T-ViT-24分别提高了0.8%和0.3%,参数量分别下降了36.2%和13.8%,FLOPs相对应的分别下降了70.6%和32.0%.图7展示了CSNet与其他分类网络的性能对比图.可以清晰直观地看出CSNet相较于这些网络在参数量及FLOPs性能指标上都有明显的优势.

表2 ImageNet-1K上的模型对比试验Table 2 Model comparison experiments on ImageNet-1K

图7 模型性能对比图Fig.7 Model performance comparison

上一组对比实验已经验证了在大型数据集上不同规模的CSNet能够有良好的表现.为了进一步研究CSNet在各种不同应用情况下的适应性,在小型数据集CIFAR-10上测试了CSNet的模型表现,对比结果如表3所示.可以看出CSNet-Small和CSNet-Base在CIFAR-10上的准确率分别达到了98.6%和98.9%,性能优于ResNet、ViT和T2T-ViT模型.实验证明了CSNet拥有比传统CNN模型更好的模型性能和比Transformer模型更好的泛化性,能够在小型数据量的下游任务中取得良好的效果.

表3 CIFAR-10上的模型对比试验Table 3 Model comparison experiments on CIFAR-10

3.4 消融实验

3.4.1 分解不同大小的卷积层

CSNet的Convolution Block利用将一个大核卷积分解为一个深度可分离卷积、一个深度可分离扩张卷积和一个点卷积的思想来捕获长距离的关系.在这一部分中,实验选取CSNet-Tiny在ImageNet-1K上测试Convolution Block中不同卷积核大小以选择最优的大型卷积核予以分解.表4展示了4种附带不同大小扩张率(Dilation rate)的卷积核.从表中可以清楚地看到,当卷积核大小分别为21和28,扩张率分别为3和4时,准确率达到了最高的77.8%.21×21与7×7的卷积核相比,在参数量和FLOPs只有小幅提升的前提下,准确率上升了0.2%.上述现象表明大核卷积对于计算机视觉任务是十分重要的.在参数量与计算复杂度相似的情况下,更大的卷积核意味着更大的感受野,可以提取到更多区域的特征信息.28×28与21×21的卷积核相比,准确率方面没有明显的提高,但却增加了少量的参数和FLOPs.故CSNet的Convolution Block都默认分解卷积核大小为的卷积运算.

表4 分解不同大小卷积核的性能表现Table 4 Performance of decomposing different size convolutional kernels

3.4.2 CSNet中模块的不同组合

CSNet与CoAtNet[10]分阶段组合不同模块的整体架构相同.为了更好地发挥卷积提取底层信息的优势,各模块的组合应遵守Convolution Block在前期阶段处理低层次特征信息的约定.这导致出现了3种不同阶段的模型变体.从Stage 1～Stage 4分别为C-C-C-A、C-C-A-A和C-C-A-A的组合模式,其中C和A分别代表Convolution Block和Attention Block.为了能系统地研究每种模块组合的性能,实验使用CSNet-Small测试了3种不同模型变体在ImageNet-1K上的准确率.由表5可以看出,组合模式为C-C-C-A的模型变体拥有最多的参数量和计算复杂度,准确率只达到了80.9%.造成上述现象的原因是在CSNet整体结构中Stage3占据了大部分的模块重复次数.当输出通道数随着阶段数增大后,Convolution Block中的分解卷积运算占据了大部分的计算量.组合模式分别为C-C-A-A和C-A-A-A的模型变体在参数量和计算复杂度相近的情况下,准确率分别达到了81.8%和81.6%.对于C-C-A-A和C-C-A-A来说,都在前期Stage利用了Convolution Block中卷积提取局部特征信息的能力来增强模型的泛化能力.从最后的准确率可以看出,C-C-A-A比C-A-A-A组合模式的准确率高0.2%.先前的一些工作[30]已经证明了自注意力机制所拥有的全局感受野在计算机视觉领域是非常重要的.上述现象表明仅仅拥有更多的Attention Block并不意味着更好地视觉处理能力.综上所述,不同模型容量的CSNet均采用C-C-A-A的模块组合方式.

表5 不同模块组合的性能表现Table 5 Performance of different module combinations

3.5 类激活图可视化

Grad-CAM[31]可以用于可视化模型的鉴别区域.图8展示了Grad-CAM分别对ResNet-101,Swin-T[32]和CSNet-Small在ImageNet数据集上的类激活图可视化.可以看出,与ResNet相比CSNet拥有更加宽广的视觉范围,这得益于在模型浅层扩张卷积较大的感受野.与Swin相比,CSNet可以更为准确和完整地定位目标物体,尤其图片中有其他标签物体时,CSNet表现出更加显著的优势.例如第1张输入图片CSNet准确注意到了整个狗的部分,并忽略了旁边的人物.

图8 类激活图可视化Fig.8 Gradient-weighted class activation mapping

4 结束语

本文提出了一种全新高效的图像分类网络,命名为CSNet.它可以有效地分阶段结合卷积与自注意力模块以克服全局感受野缺失和局部表征不足等问题.在模型浅层阶段将大核卷积分解为3种不同类型的卷积层,在保持较少参数的同时弥补了传统卷积神经网络较小感受野所带来的局限性.在深层阶段使用本文提出的深度可分离注意力机制,利用深度可分离卷积捕获局部空间交互的能力来缓解多头注意力机制极高的过拟合风险和对于数据本身严重的依赖性.尤其在数据量和模型规模有限的情况下也能利用全局感受野来捕获特征信息以实现较高的准确率.实验表明CSNet拥有强大的特征建模能力,在图像分类任务中以适量的参数和计算复杂度取得了优异的性能表现.