本文是關于論文《CBAM: Convolutional Block Attention Module》的閱讀筆記。
一、概述
深度學習中的注意力機制是模仿人類的注意力産生的，當我們在看一副圖像時，圖像中有些地方（比如人物照中的人的面部）會引起我們的注意力，而其他地方則可能會忽視。這相當于給整幅圖像加了一個可視化的權重。注意力機制可以分為通道注意力和空間注意力等。簡單來說，注意力機制就是給每個通道或整個空間特征圖乘以一個權重，以表示它們的重要性。
CBAM包含兩部分，一是通道的注意力，二是空間的注意力。CBAM可以用以下公式表示：
F ′ = M c ( F ) ⊗ F F ′ ′ = M s ( F ′ ) ⊗ F ′ \begin{aligned} \mathbf{F}^{\prime} &=\mathbf{M}_{\mathbf{c}}(\mathbf{F}) \otimes \mathbf{F} \\ \mathbf{F}^{\prime \prime} &=\mathbf{M}_{\mathbf{s}}\left(\mathbf{F}^{\prime}\right) \otimes \mathbf{F}^{\prime} \end{aligned} F′F′′=Mc(F)⊗F=Ms(F′)⊗F′
其中， F F F是特征圖， M c M_c Mc和 M s M_s Ms分别表示基于通道的和基于空間的注意力， ⊗ \otimes ⊗表示逐元素相乘， F ′ F^\prime F′和 F ′ ′ F^{\prime\prime} F′′分别表示進行了通道注意力和空間注意力後的輸出特征圖。由于CBAM子產品的輸入和輸出大小相同，是以它可以插入到已有模型的任意位置。
二、通道注意力子產品
通道注意力子產品先對輸入特征圖在空間次元分别進行最大池化和平均池化，然後分别經過一個共享權重的MLP（多層感覺機） ，然後将兩者的輸出做逐元素的相加，再經過sigmoid激活函數，得到通道注意力權重，将該權重和輸入特征圖做逐元素的乘法，就實作了通道上的注意力機制。
在空間次元分别進行最大池化和平均池化，其實就是隻保留通道的次元，其他次元為1，得到一個一維的矢量，然後再經過共享權重的MLP，MLP由兩層卷積操作組成，第一個卷積後使用ReLU激活函數，第二個卷積後使用sigmoid激活函數得到輸出權重，然後将每個通道乘以其對應的權重。上述過程可以表示為：
M c ( F ) = σ ( MLP ( A vg Pool ( F ) ) + M L P ( Max Pool ( F ) ) ) = σ ( W 1 ( W 0 ( F a v g c ) ) + W 1 ( W 0 ( F m a x c ) ) ) \begin{aligned} \mathbf{M}_{\mathbf{c}}(\mathbf{F}) &=\sigma(\operatorname{MLP}(A \operatorname{vg} \operatorname{Pool}(\mathbf{F}))+M L P(\operatorname{Max} \operatorname{Pool}(\mathbf{F}))) \\ &=\sigma\left(\mathbf{W}_{\mathbf{1}}\left(\mathbf{W}_{\mathbf{0}}\left(\mathbf{F}_{\mathbf{a v g}}^{\mathbf{c}}\right)\right)+\mathbf{W}_{\mathbf{1}}\left(\mathbf{W}_{\mathbf{0}}\left(\mathbf{F}_{\mathbf{m a x}}^{\mathbf{c}}\right)\right)\right) \end{aligned} Mc(F)=σ(MLP(AvgPool(F))+MLP(MaxPool(F)))=σ(W1(W0(Favgc))+W1(W0(Fmaxc)))
其中 σ \sigma σ表示sigmoid激活函數， W 0 W_0 W0和 W 1 W_1 W1分别表示兩個卷積操作， F a v g c F_{avg}^c Favgc和 F m a x c F_{max}^c Fmaxc分别表示平均池化和最大池化。
三、空間注意力子產品
先對輸入特征圖做基于通道的最大池化和平均池化操作，将得到的兩個特征圖做基于通道的拼接，再對其進行卷積操作和sigmoid激活函數，得到通道數為1，特征圖大小和輸入相同的注意力權重，再與輸入特征圖做逐元素的乘法，就完成了空間注意力機制。該過程可以表示為：
M s ( F ) = σ ( f 7 × 7 ( [ AvgPool ( F ) ; Max Pool ( F ) ] ) ) = σ ( f 7 × 7 ( F a v g s ; F m a x s ] ) ) \begin{aligned} \mathbf{M}_{\mathbf{s}}(\mathbf{F}) &=\sigma\left(f^{7 \times 7}([\text {AvgPool}(\mathbf{F}) ; \operatorname{Max} \operatorname{Pool}(\mathbf{F})])\right) \\ &\left.=\sigma\left(f^{7 \times 7}\left(\mathbf{F}_{\mathbf{a v g}}^{\mathbf{s}} ; \mathbf{F}_{\mathbf{m a x}}^{\mathbf{s}}\right]\right)\right) \end{aligned} Ms(F)=σ(f7×7([AvgPool(F);MaxPool(F)]))=σ(f7×7(Favgs;Fmaxs]))
其中 f 7 × 7 f^{7\times7} f7×7表示卷積核大小為 7 × 7 7\times7 7×7的卷積操作。
四、代碼
import torch
import torch.nn as nn
import torch.nn.functional as F
class CBAM_Module(nn.Module):
def __init__(self, dim, in_channels, ratio, kernel_size):
super(CBAM_Module, self).__init__()
self.avg_pool = getattr(nn, "AdaptiveAvgPool{0}d".format(dim))(1)
self.max_pool = getattr(nn, "AdaptiveMaxPool{0}d".format(dim))(1)
conv_fn = getattr(nn, "Conv{0}d".format(dim))
self.fc1 = conv_fn(in_channels, in_channels // ratio, kernel_size=1, padding=0)
self.relu = nn.ReLU()
self.fc2 = conv_fn(in_channels // ratio, in_channels, kernel_size=1, padding=0)
self.sigmoid = nn.Sigmoid()
assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
padding = 3 if kernel_size == 7 else 1
self.conv = conv_fn(2, 1, kernel_size=kernel_size, stride=1, padding=padding)
def forward(self, x):
# Channel attention module:（Mc(f) = σ(MLP(AvgPool(f)) + MLP(MaxPool(f)))）
module_input = x
avg = self.fc2(self.relu(self.fc1(self.avg_pool(x))))
mx = self.fc2(self.relu(self.fc1(self.max_pool(x))))
x = self.sigmoid(avg + mx)
x = module_input * x
# Spatial attention module:Ms (f) = σ( f7×7( AvgPool(f) ; MaxPool(F)] )))
module_input = x
avg = torch.mean(x, dim=1, keepdim=True)
mx, _ = torch.max(x, dim=1, keepdim=True)
x = torch.cat((avg, mx), dim=1)
x = self.sigmoid(self.conv(x))
x = module_input * x
return x