mobilenetv2添加CBAM注意力模块

### 如何在 MobileNetV2 中集成 CBAM 模块 为了在 MobileNetV2 中集成 CBAM (Convolutional Block Attention Module)，可以在每一层卷积操作之后加入 CBAM 来增强模型的空间和通道注意力机制。这有助于提高模型的特征表达能力和性能。 下面是一个具体的实现方式： #### 1. 定义 CBAM 类 首先定义 `CBAM` 类来封装通道注意力和空间注意力功能[^2]。 ```python import torch from torch import nn class ChannelAttention(nn.Module): def __init__(self, in_planes, ratio=8): super(ChannelAttention, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) # Shared MLP layers self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False) self.relu1 = nn.ReLU() self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x)))) max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x)))) out = avg_out + max_out return self.sigmoid(out) class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super(SpatialAttention, self).__init__() assert kernel_size in (3, 7), 'kernel size must be 3 or 7' padding = 3 if kernel_size == 7 else 1 self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) x = torch.cat([avg_out, max_out], dim=1) x = self.conv1(x) return self.sigmoid(x) class CBAM(nn.Module): def __init__(self, channel_in, reduction_ratio=8, spatial_kernel=7): super(CBAM, self).__init__() self.channel_attention = ChannelAttention(channel_in, reduction_ratio) self.spatial_attention = SpatialAttention(spatial_kernel) def forward(self, x): out = self.channel_attention(x) * x out = self.spatial_attention(out) * out return out ``` #### 2. 修改 MobileNetV2 的 InvertedResidual 卷积模块 接着修改 MobileNetV2 的 `InvertedResidual` 层，使其能够应用 CBAM 模块。 ```python def make_divisible(v, divisor, min_value=None): """ This function is taken from the original tf repo. It ensures that all layers have a channel number that is divisible by 8 It can be seen here: https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py :param v: :param divisor: :param min_value: :return: """ if min_value is None: min_value = divisor new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) # Make sure that round down does not go down by more than 10%. if new_v < 0.9 * v: new_v += divisor return new_v class ConvBNReLU(nn.Sequential): def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1): padding = (kernel_size - 1) // 2 super(ConvBNReLU, self).__init__( nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False), nn.BatchNorm2d(out_planes), nn.ReLU6(inplace=True) ) class InvertedResidualWithCBAM(nn.Module): def __init__(self, inp, oup, stride, expand_ratio): super(InvertedResidualWithCBAM, self).__init__() self.stride = stride hidden_dim = int(round(inp * expand_ratio)) self.use_res_connect = self.stride == 1 and inp == oup layers = [] if expand_ratio != 1: # pw layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1)) layers.extend([ # dw ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim), # pw-linear nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup), ]) self.conv = nn.Sequential(*layers) self.cbam = CBAM(oup) def forward(self, x): identity = x output = self.conv(x) output = self.cbam(output) if self.use_res_connect: return identity + output else: return output ``` #### 3. 替换原始的 InvertedResidual 模块 最后一步是在构建整个 MobileNetV2 架构时替换掉原有的 `InvertedResidual` 模块为带有 CBAM 的版本。 ```python class MobileNetV2WithCBAM(nn.Module): def __init__(self, num_classes=1000, width_mult=1.0): super(MobileNetV2WithCBAM, self).__init__() block = InvertedResidualWithCBAM input_channel = 32 last_channel = 1280 inverted_residual_setting = [ # t, c, n, s [1, 16, 1, 1], [6, 24, 2, 2], [6, 32, 3, 2], [6, 64, 4, 2], [6, 96, 3, 1], [6, 160, 3, 2], [6, 320, 1, 1], ] features = [ConvBNReLU(3, input_channel, stride=2)] for t, c, n, s in inverted_residual_setting: output_channel = make_divisible(c * width_mult) for i in range(n): stride = s if i == 0 else 1 features.append(block(input_channel, output_channel, stride, expand_ratio=t)) input_channel = output_channel features.append(ConvBNReLU(input_channel, last_channel, kernel_size=1)) self.features = nn.Sequential(*features) self.classifier = nn.Sequential( nn.Dropout(0.2), nn.Linear(last_channel, num_classes), ) # Initialize weights for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out') if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.BatchNorm2d): nn.init.ones_(m.weight) nn.init.zeros_(m.bias) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, 0, 0.01) nn.init.zeros_(m.bias) def forward(self, x): x = self.features(x) x = x.mean([2, 3]) # Global average pooling x = self.classifier(x) return x ``` 通过上述代码实现了将 CBAM 集成到 MobileNetV2 中的方法，从而增强了模型对于不同位置以及重要性的感知能力。

### 如何在 MobileNetV2 中集成 CBAM 模块 为了在 MobileNetV2 中集成 CBAM (Convolutional Block Attention Module)，可以在每个卷积层之后引入 CBAM 来增强特征表示。CBAM 同时考虑了通道和空间维度上的注意力机制，这有助于提高模型的表达能力和准确性。 下面是一个具体的实现方案： #### 导入库并定义 CBAM 组件 ```python import torch import torch.nn as nn class ChannelAttention(nn.Module): def __init__(self, in_planes, ratio=16): super(ChannelAttention, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False) self.relu1 = nn.ReLU() self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x)))) max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x)))) out = avg_out + max_out return self.sigmoid(out) class SpatialAttention(nn.Module): def __init__(self, kernel_size=7): super(SpatialAttention, self).__init__() assert kernel_size in (3, 7), 'kernel size must be 3 or 7' padding = 3 if kernel_size == 7 else 1 self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) x = torch.cat([avg_out, max_out], dim=1) x = self.conv1(x) return self.sigmoid(x) class CBAM(nn.Module): def __init__(self, channel_in, reduction_ratio=16, spatial_kernel=7): super(CBAM, self).__init__() self.channel_attention = ChannelAttention(channel_in, reduction_ratio) self.spatial_attention = SpatialAttention(spatial_kernel) def forward(self, x): out = self.channel_attention(x) * x out = self.spatial_attention(out) * out return out ``` #### 修改 MobileNetV2 的部分结构以加入 CBAM 假设已经有一个预训练好的 `mobilenet_v2` 模型实例，则可以通过继承原有类的方式添加新的组件： ```python from torchvision.models import mobilenet_v2 class MobilenetV2WithCBAM(mobilenet_v2.MobileNetV2): def __init__(self, num_classes=1000, **kwargs): super(MobilenetV2WithCBAM, self).__init__(num_classes=num_classes, **kwargs) # 遍历所有子模块，找到倒数第二个瓶颈层之前的位置插入CBAM for name, module in list(self.named_modules()): if isinstance(module, mobilenet_v2.InvertedResidual): cbam_layer = CBAM(module.out_channels).cuda() setattr(getattr(self.features[int(name.split('.')[0])], '_modules'), f'cbam_{name}', cbam_layer) def forward(self, x): for idx, layer in enumerate(self.features._modules.values()): x = layer(x) if hasattr(layer, "cbam"): x = getattr(layer, "cbam")(x) x = self.classifier(x.view(-1, self.last_channel)) return x ``` 上述代码展示了如何创建一个新的基于原始 MobileNet V2 架构的新类，并在其内部适当位置插入 CBAM 层[^1]。