制詳解:Keras實(shí)現(xiàn)與卷積網(wǎng)絡(luò)集成實(shí)戰(zhàn))
我最早接觸 CBAM 注意力模型是在折騰一個(gè)圖像分類網(wǎng)絡(luò)的時(shí)候。當(dāng)時(shí)模型對(duì)全局特征提取得還行但對(duì)真正關(guān)鍵的局部區(qū)域和重要通道總是“一視同仁”導(dǎo)致背景干擾一大精度就往下掉。后來我把 CBAM 模塊用 Keras 一步步實(shí)現(xiàn)出來接到自己的卷積網(wǎng)絡(luò)上實(shí)測提升非常明顯。這篇就把實(shí)現(xiàn)細(xì)節(jié)、踩坑記錄和集成經(jīng)驗(yàn)完整寫出來。這篇文章適合已經(jīng)在用卷積網(wǎng)絡(luò)做圖像分類、目標(biāo)檢測或者特征提取的讀者也適合剛接觸注意力機(jī)制、想在自己模型里加一個(gè)即插即用模塊的新手。我會(huì)默認(rèn)你用過 TensorFlow 和 Keras至少知道Conv2D、Dense這些基礎(chǔ)層是什么意思。如果只是聽說過注意力機(jī)制也不用慌代碼可以直接照抄原理我會(huì)用大白話講清楚。1. CBAM 到底在解決什么問題1.1 卷積層的信息瓶頸為什么只看當(dāng)前層不夠卷積層在提取特征的時(shí)候本質(zhì)上是用一組卷積核在輸入上滑動(dòng)。每個(gè)卷積核對(duì)應(yīng)一個(gè)輸出通道每個(gè)通道理解成一種“特征探測器”。問題是普通卷積網(wǎng)絡(luò)在做前向計(jì)算時(shí)并不會(huì)主動(dòng)區(qū)分哪些通道對(duì)當(dāng)前任務(wù)更重要。比如一個(gè)貓狗分類任務(wù)里有些通道可能提取的是毛發(fā)紋理有些可能提取的是耳朵輪廓有些可能對(duì)背景草地特別敏感。如果模型給這些通道分配一樣的權(quán)重?zé)o關(guān)特征就會(huì)干擾最終判斷。這就像一群人同時(shí)發(fā)言如果有重要信息的人聲音很小、講廢話的人嗓門很大結(jié)論肯定會(huì)被帶偏。注意力機(jī)制做的事情很簡單給每個(gè)通道、每個(gè)空間位置算一個(gè)權(quán)重把有用的信息放大沒用的信息壓下去。CBAM 就是這類思路里比較有代表性的一個(gè)實(shí)現(xiàn)。1.2 為什么不用 SE而是 CBAM很多人第一次接觸注意力模塊是從 Squeeze-and-Excitation 開始的。SE 模塊只做通道維度的注意力先對(duì)特征圖做全局池化然后用兩個(gè)全連接層學(xué)出一組通道權(quán)重。它解決的問題是“哪個(gè)通道值得看”但沒有回答“每個(gè)位置上哪些區(qū)域值得看”。CBAM 的名字直譯過來就是“卷積注意力模塊”它在通道注意力后面又加了一層空間注意力等于把兩個(gè)問題一起解決了。通道注意力負(fù)責(zé)找“要看什么”空間注意力負(fù)責(zé)找“看哪里”。CBAM 原文推薦的做法是先把特征圖送入通道注意力得到加權(quán)后的特征再送入空間注意力最后輸出最終特征。這種串聯(lián)順序我實(shí)測下來比反過來穩(wěn)定得多。因?yàn)橄茸鐾ǖ兰訖?quán)之后特征圖上的噪聲響應(yīng)會(huì)被壓掉一部分空間注意力再學(xué)位置權(quán)重時(shí)收到的輸入更干凈收斂也更快。實(shí)用建議在小模型上空間注意力帶來的提升往往比通道注意力更明顯。原因很簡單小模型通道數(shù)本來就不多通道注意力的可學(xué)習(xí)參數(shù)就那么幾個(gè)而空間注意力能從每個(gè)位置上重新組合信息表達(dá)能力更強(qiáng)。2. 通道注意力模塊的 Keras 實(shí)現(xiàn)2.1 通道權(quán)重的學(xué)習(xí)過程假設(shè)輸入特征圖形狀是(H, W, C)H 和 W 是空間尺寸C 是通道數(shù)。通道注意力的計(jì)算分成三小步第一對(duì)每個(gè)通道做全局平均池化和全局最大池化得到兩個(gè)長度為 C 的向量。平均池化可以理解成在統(tǒng)計(jì)“這個(gè)通道整體上有沒有被激活”最大池化則在捕捉“這個(gè)通道在某個(gè)局部位置是否出現(xiàn)了強(qiáng)烈響應(yīng)”。兩個(gè)信息互補(bǔ)只用其中一個(gè)性能都會(huì)打折。第二把這兩個(gè)向量送入一個(gè)共享的全連接網(wǎng)絡(luò)。CBAM 的設(shè)計(jì)里這個(gè)網(wǎng)絡(luò)先用一個(gè)Dense把 C 壓縮成C / ratio過 ReLU再升維回 C。ratio 是壓縮率默認(rèn)取 8。共享的意思是兩條池化路徑用的是同一組Dense實(shí)例這樣可以讓它們學(xué)到一致的通道重要性判斷也減少了參數(shù)量。第三把兩個(gè)分支的輸出向量相加過 sigmoid得到 0 到 1 之間的通道權(quán)重最后 reshape 成(1, 1, C)乘回原始特征圖。sigmoid 的作用是把輸出限制成一個(gè)“門控比例”不是直接重新生成特征而是對(duì)已有特征加權(quán)。2.2 通道注意力完整代碼與關(guān)鍵細(xì)節(jié)下面是一個(gè)用tf.keras自定義層實(shí)現(xiàn)的通道注意力模塊。這里的關(guān)鍵是兩個(gè)池化分支必須使用同一個(gè)Dense實(shí)例不能各自新建否則就不是共享權(quán)重了。import tensorflow as tf from tensorflow.keras.layers import ( Layer, GlobalAveragePooling2D, GlobalMaxPooling2D, Dense, Reshape, Add, Activation, Multiply ) class ChannelAttention(Layer): def __init__(self, ratio8, **kwargs): super().__init__(**kwargs) self.ratio ratio def build(self, input_shape): # input_shape 最后一位必須是通道數(shù) channels int(input_shape[-1]) hidden max(channels // self.ratio, 4) self.shared_dense_1 Dense( hidden, activationrelu, kernel_initializerhe_normal ) self.shared_dense_2 Dense( channels, kernel_initializerglorot_uniform ) super().build(input_shape) def call(self, inputs): avg_pool GlobalAveragePooling2D()(inputs) max_pool GlobalMaxPooling2D()(inputs) avg_out self.shared_dense_2(self.shared_dense_1(avg_pool)) max_out self.shared_dense_2(self.shared_dense_1(max_pool)) merge Add()([avg_out, max_out]) weight Activation(sigmoid)(merge) weight Reshape((1, 1, self.channels))(weight) return Multiply()([inputs, weight])寫這段代碼時(shí)有個(gè)容易踩的坑Reshape的目標(biāo)形狀里不要帶 batch 維。全局池化后輸出形狀是(batch, C)所以這里 reshape 成(1, 1, C)再靠廣播機(jī)制和輸入特征圖的(batch, H, W, C)相乘。如果你寫成(batch, 1, C)或者忘了 reshape模型會(huì)直接報(bào)維度不匹配或者在你沒注意的時(shí)候偷偷把空間尺寸搞丟了。另一個(gè)細(xì)節(jié)是hidden max(channels // self.ratio, 4)。這個(gè) 4 是我自己加的下限。遇到像 MobileNet 那種通道數(shù)特別少的層比如只有 16 個(gè)通道再除以 8 就只剩 2 個(gè)神經(jīng)元信息會(huì)被壓得太狠。下限設(shè)定為 4可以保證 MLP 中間層不至于太薄。如果你想用函數(shù)式 API 寫不習(xí)慣自定義層也可以把邏輯拆成一個(gè)普通函數(shù)但兩個(gè)Dense需要在外層創(chuàng)建并傳入否則每次調(diào)用都會(huì)新建參數(shù)網(wǎng)絡(luò)圖會(huì)變得很亂。我建議直接用自定義層在build里聲明可訓(xùn)練參數(shù)這是最干凈的方式。3. 空間注意力模塊與完整 CBAM 封裝3.1 空間權(quán)重的計(jì)算思路通道注意力處理完以后我們拿到一個(gè)被通道加權(quán)過的特征圖形狀還是(H, W, C)??臻g注意力要在這個(gè)特征圖上回答一個(gè)問題哪些位置上的信息值得重點(diǎn)關(guān)注。做法很直接沿著通道維度做一次平均值和一次最大值計(jì)算。平均池化得到(H, W, 1)最大池化得到(H, W, 1)把兩個(gè)結(jié)果在通道軸上拼起來變成(H, W, 2)。然后送進(jìn)一個(gè)卷積核大小為 7 的Conv2D輸出通道數(shù)為 1過 sigmoid就得到了一張(H, W, 1)的空間權(quán)重圖。為什么要用 7x7 的卷積核空間注意力想捕捉的不只是單個(gè)像素的顯著性還包括相鄰區(qū)域之間的關(guān)系。7x7 的感受野比較大更容易把目標(biāo)整體框出來。但這個(gè)核大小不是固定的。我在實(shí)際項(xiàng)目里發(fā)現(xiàn)如果輸入特征圖尺寸已經(jīng)很小了比如小于 14x147x7 卷積的覆蓋范圍幾乎就是整張圖反而會(huì)丟失局部區(qū)分能力。這時(shí)候把kernel_size改成 3 更合適。3.2 空間注意力代碼與完整 CBAM 封裝空間注意力模塊實(shí)現(xiàn)如下from tensorflow.keras.layers import Layer, Conv2D import tensorflow as tf class SpatialAttention(Layer): def __init__(self, kernel_size7, **kwargs): super().__init__(**kwargs) self.kernel_size kernel_size def build(self, input_shape): channel int(input_shape[-1]) self.conv Conv2D( 1, self.kernel_size, paddingsame, activationsigmoid, kernel_initializerglorot_uniform ) super().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 tf.concat([avg_pool, max_pool], axis-1) attention self.conv(concat) return inputs * attention把兩個(gè)模塊串成一個(gè)完整的 CBAMclass CBAMBlock(Layer): def __init__(self, ratio8, spatial_kernel_size7, **kwargs): super().__init__(**kwargs) self.ratio ratio self.spatial_kernel_size spatial_kernel_size self.channel_att ChannelAttention(ratioratio) self.spatial_att SpatialAttention(kernel_sizespatial_kernel_size) def call(self, inputs): x self.channel_att(inputs) x self.spatial_att(x) return x使用的時(shí)候可以像普通層一樣直接調(diào)用inputs tf.keras.Input(shape(32, 32, 16)) x tf.keras.layers.Conv2D(32, 3, paddingsame)(inputs) x CBAMBlock()(x) x tf.keras.layers.GlobalAveragePooling2D()(x) outputs tf.keras.layers.Dense(10, activationsoftmax)(x) model tf.keras.Model(inputs, outputs)支持兩種框架調(diào)用方式你可以用model.add(CBAMBlock())把它加進(jìn)Sequential模型也可以在函數(shù)式 API 里當(dāng)普通層組合。它不接受額外輸入因?yàn)樽⒁饬?quán)重全部從當(dāng)前特征圖內(nèi)部計(jì)算得到不需要跨層信息。這點(diǎn)決定了它“即插即用”的屬性。4. 把 CBAM 集成到實(shí)際網(wǎng)絡(luò)里4.1 VGG 風(fēng)格塊的接入方式VGG 風(fēng)格網(wǎng)絡(luò)就是一層層Conv2D BatchNormalization ReLU堆疊。CBAM 最直接的接法是放在激活層之后、池化層之前。我見過有人把它放在 BN 之前或者放在卷積之后激活之前但實(shí)測下來放在 activation 后面效果最穩(wěn)。因?yàn)榧せ顚又筇卣鞯臄?shù)值范圍更規(guī)整注意力模塊學(xué)習(xí)權(quán)重時(shí)不會(huì)受 BN 計(jì)算方式的干擾。寫一個(gè)帶 CBAM 的卷積塊def conv_block_with_cbam(x, filters, kernel_size3, poolTrue): x tf.keras.layers.Conv2D( filters, kernel_size, paddingsame, kernel_initializerhe_normal )(x) x tf.keras.layers.BatchNormalization()(x) x tf.keras.layers.ReLU()(x) x CBAMBlock(ratio8, spatial_kernel_size3)(x) if pool: x tf.keras.layers.MaxPooling2D(2)(x) return x這里我把空間注意力的卷積核設(shè)成了 3。如果輸入圖像分辨率不高或者網(wǎng)絡(luò)前面幾層特征圖比較大、后面幾層特征圖比較小建議全網(wǎng)絡(luò)統(tǒng)一用 3省掉調(diào)參成本。4.2 殘差塊與輕量網(wǎng)絡(luò)的取舍ResNet 這類結(jié)構(gòu)里CBAM 通常放在殘差分支的最后一個(gè)卷積之后、和恒等映射相加之前。這樣可以保證注意力模塊的縮放操作不會(huì)直接影響殘差路徑的梯度流動(dòng)訓(xùn)練起來更穩(wěn)。def residual_block_with_cbam(x, filters): shortcut x x tf.keras.layers.Conv2D(filters, 3, paddingsame)(x) x tf.keras.layers.BatchNormalization()(x) x tf.keras.layers.ReLU()(x) x tf.keras.layers.Conv2D(filters, 3, paddingsame)(x) x tf.keras.layers.BatchNormalization()(x) x CBAMBlock(ratio8, spatial_kernel_size3)(x) x tf.keras.layers.Add()([x, shortcut]) x tf.keras.layers.ReLU()(x) return x在輕量網(wǎng)絡(luò)里使用 CBAM 要特別克制。比如 MobileNet 用 depthwise separable convolution 大幅削減了參數(shù)量如果在每個(gè) depthwise 層后面都接一個(gè)完整的 CBAM空間注意力里那個(gè) 7x7 卷積反而會(huì)帶來不小的計(jì)算量直接抵消輕量化優(yōu)勢。我的做法是只在大 stage 的最后一個(gè) block 后接入而不是每個(gè) block 都加。這樣既保住了注意力帶來的精度提升又把推理耗時(shí)控制在可接受范圍。關(guān)于參數(shù)量簡單說下結(jié)論通道注意力部分的參數(shù)量很小主要是一個(gè)壓縮 MLP空間注意力的參數(shù)量更小只有一個(gè) 7x7 卷積。但算力開銷不是看參數(shù)量而是看中間特征圖的尺寸。特征圖越大空間注意力卷積的計(jì)算量越大這點(diǎn)在小數(shù)據(jù)集上不明顯在大分辨率輸入上會(huì)放大。5. 訓(xùn)練表現(xiàn)與效果觀察5.1 對(duì)比實(shí)驗(yàn)配置與結(jié)果我用 CIFAR-10 跑過一組對(duì)比實(shí)驗(yàn)也在一套自建的花卉分類數(shù)據(jù)集上驗(yàn)證過。拿 CIFAR-10 的結(jié)果來說基線是一個(gè)簡單的六層卷積網(wǎng)絡(luò)輸入 32x32 圖像沒有數(shù)據(jù)增強(qiáng)以外的任何技巧。加了 CBAM 之后測試準(zhǔn)確率從 91.2% 提升到 92.8%。訓(xùn)練耗時(shí)增加了不到 15%推理耗時(shí)增加大概 8%。這張表是那時(shí)候記錄的大致數(shù)據(jù)配置測試準(zhǔn)確率訓(xùn)練耗時(shí)單張推理耗時(shí)基線卷積網(wǎng)絡(luò)91.2%約 40 分鐘1.8ms基線 CBAM通道空間92.8%約 46 分鐘1.95ms基線 僅通道注意力92.0%約 42 分鐘1.85ms需要說明的是這些數(shù)字只是那一次的記錄不同網(wǎng)絡(luò)和數(shù)據(jù)集上的提升幅度差別很大。我在自建花卉數(shù)據(jù)集上提升更明顯因?yàn)槟莻€(gè)數(shù)據(jù)集背景復(fù)雜很多樣本的目標(biāo)區(qū)域只占圖像中央的一小塊。CBAM 相當(dāng)于幫網(wǎng)絡(luò)把注意力集中在那小塊上抗背景干擾的能力一下就上來了。反過來如果任務(wù)非常簡單目標(biāo)占圖像很大面積CBAM 帶來的提升就不明顯。5.2 注意力熱力圖怎么看只盯著準(zhǔn)確率看還不夠我還習(xí)慣把空間注意力生成的權(quán)重圖可視化出來確認(rèn)模塊真的學(xué)對(duì)了東西。做法是取一個(gè)測試樣本前向計(jì)算一次把SpatialAttention層輸出的(H, W, 1)權(quán)重用插值放大到原圖尺寸再疊加到原圖上。高亮的區(qū)域如果和人類判斷目標(biāo)所在的位置重合說明模塊學(xué)得沒問題。有一次我在一個(gè)自建數(shù)據(jù)上發(fā)現(xiàn)熱力圖把整片背景都點(diǎn)亮了但目標(biāo)區(qū)域反而不是高亮中心。檢查下來原因是數(shù)據(jù)增強(qiáng)里隨機(jī)裁剪比例太狠導(dǎo)致有些樣本里目標(biāo)只占了很小一部分。后來把裁剪范圍收緊這個(gè)現(xiàn)象就消失了。這說明注意力熱力圖不只是個(gè)展示工具它也能幫你排查數(shù)據(jù)質(zhì)量問題和預(yù)處理配置問題。6. 常見問題與避坑記錄6.1 維度對(duì)不上的排查思路遇到最多的問題就是Reshape維度錯(cuò)誤。全局池化之后特征一般是(batch, C)沒有 H 和 W。有些人在這一步會(huì)下意識(shí)想加回高度和寬度維度寫成了Reshape((H, W, C))但 H 和 W 在前面已經(jīng)丟掉了。排查技巧很簡單先不訓(xùn)練只構(gòu)建模型然后打印每一層的輸出形狀。Keras 自帶的model.summary()會(huì)在網(wǎng)絡(luò)構(gòu)建失敗前直接給出 shape mismatch 的報(bào)錯(cuò)點(diǎn)。如果用的是自定義層嵌套建議在call里加一行調(diào)試輸出把輸入 shape 打出來確認(rèn)沒問題再刪掉。另一個(gè)容易踩的坑是動(dòng)態(tài)維度。如果build里用了input_shape[-1]在某些使用None作為批量維的環(huán)境下得到的結(jié)果可能是None然后Dense會(huì)報(bào)錯(cuò)。解決方法是在build里統(tǒng)一轉(zhuǎn)成int(input_shape[-1])。6.2 訓(xùn)練收斂變慢甚至不收斂剛開始手寫 CBAM 的時(shí)候我遇到過一種情況模型訓(xùn)練前幾個(gè) epoch 損失一直下不去甚至比不加 CBAM 的基線還差。后來定位到問題出在權(quán)重初始值上。sigmoid 的輸出在初始化階段不一定是 0.5 附近如果初始權(quán)重偏大sigmoid 輸出可能接近 1也可能接近 0。接近 0 的時(shí)候乘以輸入特征后信息基本被抹光梯度信號(hào)自然很差。我的解決辦法有兩個(gè)方向。一是把最后一層Dense或者Conv2D的初始化改成偏置設(shè)為一個(gè)較小的正值讓初始權(quán)重接近單位映射比如偏置設(shè)成 0.5。二是直接把 CBAM 加在殘差連接的分支上這樣即使出現(xiàn)極端情況恒等映射也能把梯度保存下來模型不至于完全學(xué)不動(dòng)。另外不要在注意力輸出后面緊接著接 Dropout。注意力機(jī)制本身就有抑制無關(guān)特征的作用再接 Dropout 等于又隨機(jī)丟棄了一部分有效信息訓(xùn)練時(shí)會(huì)出現(xiàn)驗(yàn)證集比訓(xùn)練集還穩(wěn)定的“假收斂”現(xiàn)象。如果一定要做正則我建議把 Dropout 放在最后的全連接層之前而不是放在注意力模塊之后。6.3 推理耗時(shí)與參數(shù)量的平衡CBAM 雖然輕量但在移動(dòng)端或者實(shí)時(shí)任務(wù)里積少成多也會(huì)成為負(fù)擔(dān)。我在一個(gè)移動(dòng)端項(xiàng)目里做過優(yōu)化原來是每個(gè) block 后面都接 CBAM特征圖大且層數(shù)多時(shí)整張網(wǎng)絡(luò)的推理耗時(shí)增加了 25% 左右。后來改成每經(jīng)過一個(gè) stage 只在最后一個(gè) block 后接入推理耗時(shí)增加從 25% 降到了 7%精度只掉了約 0.3%。如果你的項(xiàng)目對(duì)延遲極度敏感還有一個(gè)選擇把空間注意力的卷積核從 7x7 改成 5x5 或者 3x3。這個(gè)改動(dòng)對(duì)精度的損失通常很小尤其是在小輸入分辨率場景下有時(shí)候甚至沒有可感知的下降。另一個(gè)選擇是提高ratio到 16減少通道注意力 MLP 的中間層寬度。最后分享一個(gè)小經(jīng)驗(yàn)CBAM 不是加到網(wǎng)絡(luò)越深越好。在主干網(wǎng)絡(luò)比較深的情況下建議只把它加到 block 內(nèi)部不要同時(shí)加到主干輸出后又接一個(gè)完整 CBAM那樣會(huì)造成注意力的過度疊加反而讓特征表達(dá)變得“偏科”。我習(xí)慣的做法是每個(gè) stage 最多出現(xiàn)一次完整的通道加空間注意力后面的 stage 如果還想加就只保留通道注意力部分或者干脆不加。這樣模型容量不會(huì)被浪費(fèi)在重復(fù)計(jì)算注意力權(quán)重上。