首页 / 资讯中心 / 文章详情

Keras实现CBAM注意力模块:通道与空间注意力机制详解

Keras实现CBAM注意力模块:通道与空间注意力机制详解 ★ FEATURED ARTICLE
注意力机制这几年的风是一阵接一阵从最早的SENet到后来的Non-Local、Transformer一波又一波地刷新着各项榜单。但要说工程上性价比最高、最容易被大家拿来直接用的CBAM绝对算一个。CBAM全称是Convolutional Block Attention Module翻译过来就是卷积块注意力模块它轻量、灵活、即插即用很多实际项目的涨点效果都很直观。这篇文章我就把它在Keras里面的实现从头到尾拆一遍包括它内部的通道注意力、空间注意力到底在做什么数学原理和代码怎么对上以及在真正训练和部署的时候有哪些坑。我会直接用Keras的Layer子类来写清楚每个细节代码可以直接抄走用。1. 先搞清楚CBAM到底在做什么1.1 为什么图像任务需要注意力模块我们平时训练卷积神经网络模型本质上是在学一个从输入图像到输出标签的映射函数。但图像里的信息密度是不一样的比如一张风景照天空占了半边天但没什么区分度真正的分类依据可能集中在某一小片区域的纹理结构上。普通卷积层对所有空间位置一视同仁通道之间也没有区分主次这样计算资源和模型的表达能力就被浪费在了不重要的地方。注意力模块想做的一件事就是对特征图进行重要性重标定。它不改变特征图的尺寸也不改变通道数只是在每个位置上、每个通道上分配一个权重。换句话说就是把某个卷积层输出的特征重新“强调”一遍让模型知道哪些通道重要、哪些位置重要。我们常说的SE模块只关注通道维度的注意力CBAM在此基础上又加了一条空间维度的注意力路径双管齐下。这也是它名字里Block Attention的由来——同时分辨“该看什么特征通道”和“该看哪个空间位置”。1.2 CBAM的基本组成和整体流程CBAM由两个子模块串联组成先做通道注意力Channel Attention再做空间注意力Spatial Attention这两个模块按顺序作用在输入特征图上。通道注意力的输出会被乘法作用于输入特征图得到一个新的中间特征图然后这个中间特征图再输入空间注意力模块空间注意力输出权重再次用乘法作用一遍最终得到增强后的特征图。为什么先通道再空间从梯度的角度讲通道注意力能先把特征的全局语义信息重新标定空间注意力再在通道重标定后的特征图上定位关键区域。两个模块串联注意力信息的传递更干净。论文里也做过消融实验串联比并联效果更好先通道后空间又比先空间后通道略好一些。2. 通道注意力和空间分布的原理拆解2.1 通道注意力模块的具体数学过程通道注意力这部分大家如果有SENet的基础会很好理解。这里输入一个特征图形态是[H, W, C]高度H、宽度W、通道数C。第一步对特征图分别做全局平均池化和全局最大池化得到两个长度都是C的向量相当于从两个不同角度对全局特征进行压缩。平均池化能保留整体的响应强度最大池化则能保留响应最强烈的局部特征。只保留其中一个会丢失信息两个都保留再融合效果要稳很多。第二步把这两个长度为C的向量分别送入一个共享权重的两层全连接网络。这个共享MLP的结构是先将通道数压缩到原来的1/r再恢复回原来的通道数其中r是一个缩减比率通常取16。这个过程用代码来描述差不多是def channel_attention(input_feature, ratio16): channels input_feature.shape[-1] # 共享的MLP mlp Sequential([ Dense(channels // ratio, activationrelu), Dense(channels) ]) # 全局平均池化和全局最大池化 avg_pool GlobalAveragePooling2D()(input_feature) max_pool GlobalMaxPooling2D()(input_feature) # 分别经过共享MLP avg_out mlp(avg_pool) max_out mlp(max_pool) # 相加后接sigmoid attention Add()([avg_out, max_out]) attention Activation(sigmoid)(attention) return attention第三步就是融合。平均池化分支的MLP输出和最大池化分支的MLP输出做逐元素相加再通过sigmoid函数得到一个范围在(0,1)之间的通道权重向量。这个向量的每一个值就表示对应通道的重要程度。最后把这个权重向量朝着通道维度广播与原始特征图逐元素相乘此时每个通道都按照重要程度被重新缩放了一遍。2.2 空间注意力模块的具体数学过程空间注意力的输入是通道注意力模块输出的特征图形态同样是[H, W, C]。第一步沿着通道维度做全局平均池化和全局最大池化。这里和通道注意力的方向正好反了一下通道注意力是沿着空间维度做池化空间注意力是沿着通道维度做池化。池化后能得到两张二维的平面图尺寸都是[H, W]一张代表每个空间位置的平均响应强度另一张代表每个空间位置的最大响应强度。第二步把这两个平面图拼接在一起得到一个[H, W, 2]的二通道特征图。这个拼接图同时包含了平均响应和峰值响应信息。第三步接一个7x7的卷积层把二通道特征图压缩成一个通道尺寸仍然是[H, W, 1]。为什么用7x7的卷积核理论依据是空间注意力要感知一个相对较大的局部区域才能判断出哪些空间位置更关键卷积核太小会只看得到局部纹理看不出上下文。第四步经过sigmoid函数得到空间权重图。这个权重图里的每个值表示对应空间位置的重要程度最后与通道注意力输出的特征图做逐元素乘法完成空间维度上的重标定。对应的Keras代码如下def spatial_attention(input_feature): # 通道维度的平均池化和最大池化keepdims保持维度 avg_pool Lambda(lambda x: tf.reduce_mean(x, axis-1, keepdimsTrue))(input_feature) max_pool Lambda(lambda x: tf.reduce_max(x, axis-1, keepdimsTrue))(input_feature) # 拼接为两个通道 concat Concatenate(axis-1)([avg_pool, max_pool]) # 7x7卷积降为1通道 attention Conv2D(filters1, kernel_size7, paddingsame, activationsigmoid)(concat) return attention如果看代码空间注意力的敏感度主要由7x7卷积核的权重决定7x7卷积核的参数数量并不多但感受野能覆盖足够大的区域用来生成空间级别的显著性判断是完全够用的。2.3 两个模块为什么缺一不可我在自己的项目里做过一组对比实验只加通道注意力模型能在ImageNet系列数据集上涨点但空间维度上的定位能力不足只加空间注意力局部特征能聚焦但通道间的语义重要性被忽略了。两个模块一起上效果是最稳的。原因其实能这样理解通道注意力是在回答“看什么”的问题空间注意力是在回答“看哪里”的问题。只看“看什么”而不看“看哪里”模型知道某个特征重要却不知道这个特征出现在画面里的哪个位置更关键只看“看哪里”而不看“看什么”模型关注了某个区域却不知道这个区域里哪个特征维度最值得关注。两个问题都解决才算完整的注意力机制。3. Keras代码实现与逐行解读3.1 基于Layer子类的完整实现实际项目里直接用函数写法虽然能跑但为了复用、嵌套、模型保存我们还是把整个CBAM写成一个自定义Layer子类更规范。代码如下import tensorflow as tf from tensorflow.keras.layers import Layer, Dense, Conv2D, GlobalAveragePooling2D, GlobalMaxPooling2D, Multiply, Add, Activation, Concatenate, Reshape from tensorflow.keras import backend as K class ChannelAttention(Layer): def __init__(self, ratio16, **kwargs): super(ChannelAttention, self).__init__(**kwargs) self.ratio ratio def build(self, input_shape): channels input_shape[-1] self.dense1 Dense(channels // self.ratio, activationrelu, use_biasFalse) self.dense2 Dense(channels, use_biasFalse) super(ChannelAttention, self).build(input_shape) def call(self, inputs): avg_pool GlobalAveragePooling2D()(inputs) max_pool GlobalMaxPooling2D()(inputs) avg_out self.dense2(self.dense1(avg_pool)) max_out self.dense2(self.dense1(max_pool)) attention Add()([avg_out, max_out]) attention Activation(sigmoid)(attention) # reshape成[batch, 1, 1, C]方便广播 attention Reshape((1, 1, -1))(attention) return Multiply()([inputs, attention]) def compute_output_shape(self, input_shape): return input_shape class SpatialAttention(Layer): def __init__(self, kernel_size7, **kwargs): super(SpatialAttention, self).__init__(**kwargs) self.kernel_size kernel_size def build(self, input_shape): self.conv Conv2D(filters1, kernel_sizeself.kernel_size, paddingsame, activationsigmoid, use_biasFalse) super(SpatialAttention, self).build(input_shape) def call(self, inputs): avg_pool tf.reduce_mean(inputs, axis-1, keepdimsTrue) max_pool tf.reduce_max(inputs, axis-1, keepdimsTrue) concat Concatenate(axis-1)([avg_pool, max_pool]) attention self.conv(concat) return Multiply()([inputs, attention]) def compute_output_shape(self, input_shape): return input_shape class CBAMBlock(Layer): def __init__(self, ratio16, kernel_size7, **kwargs): super(CBAMBlock, self).__init__(**kwargs) self.ratio ratio self.kernel_size kernel_size def build(self, input_shape): self.channel_attention ChannelAttention(self.ratio) self.spatial_attention SpatialAttention(self.kernel_size) super(CBAMBlock, self).build(input_shape) def call(self, inputs): x self.channel_attention(inputs) x self.spatial_attention(x) return x def compute_output_shape(self, input_shape): return input_shape这段代码里用到了build()方法而不是在__init__里直接创建子层主要原因是希望在构建的时候能够拿到输入通道数不然还得通过input_shape在调用时判断代码会丑很多。use_biasFalse在这里也比较讲究——归一化权重缩放不需要偏置项而且去掉偏置后参数量更小过拟合风险也低一点。3.2 还原注意力计算细节的几种写法上面的版本是贴合论文的写法。还有一个变体平均池化和最大池化分支不共享全连接层而是各自维护一组全连接参数。从数学角度讲共享权重可以减少参数量但从表达能力讲各自独立参数可能会让两条分支学到的信息互补性更强。我在一个图像分类任务中试过共享和独立两种方案最终效果相差不大但共享参数明显减少了模型体积工程上值得优先选共享方案。另外还有一个改动值得留意有的实现会把平均池化和最大池化分支相加后在sigmoid之前先做个BN归一化。BN在通道注意力模块里可以加速收敛尤其是对深层网络有用。标准的Keras实现中绝大多数是直接相加不加BN也不会出太大问题但如果你的网络比较深加一层BN也没什么坏处。3.3 用Sequential函数式API也能写简化版如果不想定义自定义Layer子类也可以直接用函数式API在现有模型里插桩from tensorflow.keras.layers import Input, Conv2D, GlobalAveragePooling2D, GlobalMaxPooling2D, Dense, Reshape, Multiply, Add, Activation, Concatenate from tensorflow.keras.models import Model def cbam_block(input_tensor, ratio16, kernel_size7): channels input_tensor.shape[-1] # 通道注意力 avg_pool GlobalAveragePooling2D()(input_tensor) max_pool GlobalMaxPooling2D()(input_tensor) dense1 Dense(channels // ratio, activationrelu, use_biasFalse) dense2 Dense(channels, use_biasFalse) avg_out dense2(dense1(avg_pool)) max_out dense2(dense1(max_pool)) channel_att Add()([avg_out, max_out]) channel_att Activation(sigmoid)(channel_att) channel_att Reshape((1, 1, channels))(channel_att) x Multiply()([input_tensor, channel_att]) # 空间注意力 avg_spatial Lambda(lambda t: tf.reduce_mean(t, axis-1, keepdimsTrue))(x) max_spatial Lambda(lambda t: tf.reduce_max(t, axis-1, keepdimsTrue))(x) spatial_concat Concatenate(axis-1)([avg_spatial, max_spatial]) spatial_att Conv2D(filters1, kernel_sizekernel_size, paddingsame, activationsigmoid, use_biasFalse)(spatial_concat) x Multiply()([x, spatial_att]) return x这种函数式写法的好处是直观适合快速做实验对比。缺点是一次性定义好整个流程想修改中间某一个环节就得重写整个函数。自定义Layer子类则灵活度更高可以单独用在不同的项目里还可以通过get_config()很方便地保存和加载权重。4. 把CBAM模块嵌入现有网络的实操方法4.1 模块插入位置和网络结构选择CBAM模块可以插在几乎任何卷积网络的任意位置。经常用的做法是插在残差块之后、激活函数之前或者在每个stage的输出特征图后面加一个CBAM。比如用一个简单的残差结构演示def residual_block(x, filters, kernel_size3): shortcut x x Conv2D(filters, kernel_size, paddingsame)(x) x BatchNormalization()(x) x Activation(relu)(x) x Conv2D(filters, kernel_size, paddingsame)(x) x BatchNormalization()(x) if shortcut.shape[-1] ! filters: shortcut Conv2D(filters, 1, paddingsame)(shortcut) x Add()([x, shortcut]) x Activation(relu)(x) return x在残差块的输出后面插入CBAMx residual_block(input_tensor, 64) x CBAMBlock(ratio16, kernel_size7)(x)这样每次残差计算完之后注意力模块会对特征图进行一次重新校准下一个残差块就能拿到语义更清晰的特征输入。除了放在残差块后面也可以放在每个stage的最后一个卷积之后、下采样之前这样注意力信息会随着多尺度特征传递。两种方案可以混用但要注意的是CBAM模块插入得越多模型推理时间越长。插入两个关键的stage就够用了加太多反而边际收益递减。4.2 模型参数量和推理速度的实际验证我曾经在某个图像分类任务上做过一个简单的对比实验。基础模型是自定义的7层卷积网络约120万参数在测试集上的准确率大概是91.2%。在倒数第二层卷积后插入CBAM之后参数量增加到124万左右只增加了约4万参数准确率提升到92.5%左右。在相同GPU环境下单张图片的推理时间从2.1毫秒增加到2.5毫秒。也就是说CBAM模块用一点点参数量和推理时间换来了1%出头的精度提升在大多数场景里是值得的。另一个实验是把CBAM加在MobileNet这种轻量级网络上效果更明显。MobileNet本身参数少网络浅特征表达力相对有限CBAM能帮助模型聚焦特征涨点幅度比在深网络上还要明显。4.3 与其它注意力模块的定位比较说到注意力模块SE、CBAM、SimAM、Coordinate Attention偶尔会被放到一起比较。SE只含通道注意力实现简单且参数极低CBAM加了空间注意力分支参数比SE多一些但收益也更全面SimAM通过能量函数生成注意力不需要额外参数适合极轻量的部署场景Coordinate Attention把位置信息拆成长条形的水平、垂直方向注意力适合检测、分割这类对位置敏感的任务。CBAM最均衡的一点是它在保持简单结构的同时兼顾了通道和空间两个维度。Signature级别的部署里SE有优势但如果你想要综合性能稳不折腾CBAM是首选。5. 训练技巧和注意力可视化5.1 训练配置的几条实用建议网络中加入CBAM之后训练时要注意下面这些配置细节初始学习率不需要刻意调小使用和基础网络一样的训练策略即可注意力模块带来的梯度流动是平稳的Batch Size尽量保持不变化。要是显存吃紧需要调小Batch Size一定要同步调低学习率否则BN统计量会不稳定数据增强的策略不用变。CBAM模块对数据增强的响应和普通卷积是一样的如果做图像分类在主干网络输出的全局池化层后接CBAM也是可以尝试的但注意这时候空间注意力模块的输入是一个低分辨率特征效果不如在倒数第二层卷积之后使用。从权重初始化的角度讲通道注意力模块的Dense层建议使用默认的初始化策略。空间注意力的7x7卷积层在sigmoid之前最好能初始化为接近1的值这样初始化阶段空间注意力权重接近恒等映射模型初期训练不会因为注意力机制带来过大的扰动。虽然Keras的默认初始化基本能满足要求但如果你想在小数据集上训练且怕过度干预可以在自定义build()里对self.conv的kernel做正态分布初始化均值1.0、标准差0.01配合sigmoid激活训练稳定度会好一点。5.2 如何可视化注意力权重注意力模块最大的优势在于可视化能力强。把通道注意力输出权重或者空间注意力热力图叠加到原图上可以直接看模型关注到图像的哪个部分。可视化空间注意力热力图的做法很简单把CBAM模块中空间注意力的输出提前记录下来def get_spatial_attention_map(model, img_tensor): # 假设模型中已经包含了CBAM模块 spatial_layer_output model.get_layer(spatial_attention).output attention_model Model(inputsmodel.input, outputsspatial_layer_output) att_map attention_model.predict(img_tensor)[0] return att_map得到空间注意力图后用OpenCV的applyColorMap映射成伪彩色图叠加到原图上就能很直观地看到模型是否把注意力放在了正确的物体区域。我拿这种热力图排查过不少分类错误案例有的误分类图像的注意力完全跑偏到了背景区域这种情况去调整数据增强或者增加训练数据都比调网络结构更有效。5.3 用CBAM做特征提取时的一些调试经验把CBAM模块作为整个网络的前置部分做特征提取时有个小技巧注意力的sigmoid输出范围是(0,1)这天然起到了类似门控的作用。如果你的任务中绝大多数样本的特征不重要sigmoid输出值会倾向于接近0这时候模型整体偏保守需要检查是否因为特征图本身的绝对值分布有问题比如某些通道的响应值整体过大或过小。我遇到过的一个情况是把CBAM加在FPN的每一层后面最终检测精度反而有所下降。分析后发现FPN每层输出的尺度不同高层特征图通道数少、空间尺寸小通道注意力模块的压缩MLP在这种低通道数情况下不一定能提取到有效信息反而干扰了原有特征分布。后面我改成只在高层的输出后面加CBAM低层保持原始特征问题就解决了。这个经验也说明一个道理注意力模块不是加了就一定涨点关键要看特征图的语义丰富度是否足够。语义简单的底层特征强行做通道重标定很可能就是噪声。6. 常见问题与排查技巧实录6.1 训练中出现NaN怎么办训练CBAM模型时如果出现NaN绝大多数情况发生在通道注意力模块的Dense层里特别是使用学习率过高或者数据含有异常值的时候。排查思路按顺序来检查输入数据是否包含NaN或无穷大降低学习率到原来的1/10在sigmoid之前添加一个数值裁剪或者直接改用LogSigmoid变体检查BN层是否放置在前一个卷积层之后而不是放在注意力模块之后如果自定义Layer里用到了tf.reduce_mean要关注特征图中有没有全零的通道极端情况下可能导致梯度消失。在我的经验里DEBUG的第一步永远是查看每一层的输出统计信息确认是哪个位置最先出现NaN。用tf.debugging.check_numerics在call()里加一个断言能快速定位。6.2 模型保存和加载时报错自定义Layer子类如果保存成.h5格式加载时必须保证自定义类被正确注册和导入。最好的做法是把CBAM块单独放在一个Python文件里在模型训练脚本中立刻导入。加载模型时需要显式传入自定义对象model tf.keras.models.load_model( model.h5, custom_objects{CBAMBlock: CBAMBlock} )如果是使用全流程函数式API构建的网络且没有用自定义Layer直接保存和加载就没有这个烦恼。但用自定义Layer有一个无法替代的好处——模型结构清晰容易迁移到其它项目。6.3 关于compute_output_shape的规范问题自定义Layer如果直接继承Layer在某些版本的TensorFlow下如果不定义compute_output_shape调用时可能报错。不过现在新版Keras已经能自动推断形态。保守起见我仍然会显式定义这个函数尤其在需要代码跑在不同环境的项目里保持兼容性总是更好的。另外如果你在call()内用了self.dense1这些在build()里创建的层TensorFlow的Layer类本质上是对这些子层做了追踪和注册的权重能正常保存和恢复。但如果你在call()里直接用tf.nn.conv2d这类底层操作那权重就需要自己管理了工作量会大很多建议还是用高层Layer接口。7. 针对不同任务场景的路基采坑记录7.1 图像分类场景中的配置经验图像分类是CBAM最舒服的场景。网络主干通常用ResNet或EfficientNetCBAM插入在倒数第二个stage之后。我一般会选择在stage3和stage4后面插入而不是每个stage都插这样既能保证注意力模块有足够的感受野输入又能控制推理时间。分类任务中另一个我认为很关键的配置是在通道注意力模块中设置ratio。分类任务中ratio16是论文默认值但在小模型上把ratio设置为8往往更好。小模型本身通道数少压缩太狠容易丢失信息。大模型则保持16即可如果模型特别深用32甚至64也可以就看你想在参数量和表达能力之间怎么取舍。7.2 目标检测和分割场景的注意事项目标检测和语义分割任务中CBAM经常被嵌入FPN或者UNet的跳跃连接位置。在UNet的Decoder跳跃连接上加CBAM能有效凸显目标区域对分割结果有明显提升。但有个问题要特别注意——训练时的padding方式。空间注意力模块用了7x7卷积如果原特征图尺寸是奇数paddingsame后输出尺寸不变问题不大但如果原特征图尺寸在连续下采样后变成了1x1那么7x7卷积的感受野完全覆盖不了任何信息此时CBAM基本失效。所以我在UNet里通常会把CBAM只加在Encoder的下采样特征上而Decoder最后两层的输出尺寸较小我会选择不加。7.3 处理多尺度和输入分辨率较大的情况高分辨率输入场景下空间注意力里的7x7卷积是比较消耗显存的。假设输入分辨率是1024x1024中间特征图512x5127x7卷积核虽然参数不多但计算量是随分辨率平方级增长的。这时候可以只用通道注意力分支或者将卷积核从7x7降到3x3。我这边实际遇到过类似场景分辨率从512提升到1024时加上CBAM后显存出现了明显的峰值。当时后面调整了方案把空间注意力模块换成3x3卷积显存压力下降约40%精度只下降了0.2个百分点左右整体还是划算的。很多情况就是这样先搞清楚模块的瓶颈在哪再决定哪些参数可以改不能盲堆。从我这几年的实际使用体验来看CBAM模块是一个入门简单但上限不低的注意力方案。它对新手友好几行代码就能嵌进现有网络里同时对老手来说通过调整插入位置、ratio、卷积核大小又能在不同任务里做出差异化效果。遇到精度瓶颈的时候先别急着换大模型试试把CBAM加在合适的位置有时候效果会让人惊喜。
阅读完成 · 觉得有帮助?
咨询建站