devbox
羊村懒王
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

YOLO V5（6.0）添加CBAM、SENet、CA注意力机制手把手教程_yolov5加cbam

作者：羊村懒王 | 2024-03-22 21:23:30

赞

踩

yolov5加cbam

1.添加CBAM

1)修改commom.py

在代码最后添加

  1. #CBAM
  2. class ChannelAttention(nn.Module):
  3.     def __init__(self, in_planes, ratio=16):
  4.         super(ChannelAttention, self).__init__()
  5.         self.avg_pool = nn.AdaptiveAvgPool2d(1)
  6.         self.max_pool = nn.AdaptiveMaxPool2d(1)
  7.         # in_planes // ratio 这里会出现如下警告：
  8.         # UserWarning: __floordiv__ is deprecated(被舍弃了), and its behavior will change in a future version of pytorch.
  9.         # It currently rounds toward 0 (like the 'trunc' function NOT 'floor').
  10.         # This results in incorrect rounding for negative values.
  11.         # To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'),
  12.         # or for actual floor division, use torch.div(a, b, rounding_mode='floor').
  13.         # kernel = torch.DoubleTensor([*(x[0].shape[2:])]) // torch.DoubleTensor(list((m.output_size,))).squeeze()
  14.         self.f1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
  15.         self.relu = nn.ReLU()
  16.         self.f2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
  17.      
  18.  
  19.         self.sigmoid = nn.Sigmoid()
  20.  
  21.     def forward(self, x):
  22.         # 全局平均池化—>MLP两层卷积
  23.         avg_out = self.f2(self.relu(self.f1(self.avg_pool(x))))
  24.         # 全局最大池化—>MLP两层卷积
  25.         max_out = self.f2(self.relu(self.f1(self.max_pool(x))))
  26.         out = self.sigmoid(avg_out + max_out)
  27.         return out
  28.  
  29.  
  30. class SpatialAttention(nn.Module):
  31.     def __init__(self, kernel_size=7):
  32.         super(SpatialAttention, self).__init__()
  33.  
  34.         assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
  35.         padding = 3 if kernel_size == 7 else 1
  36.  
  37.         self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
  38.         self.sigmoid = nn.Sigmoid()
  39.  
  40.     def forward(self, x):
  41.         # 基于channel的全局平均池化(channel=1)
  42.         avg_out = torch.mean(x, dim=1, keepdim=True)
  43.         # 基于channel的全局最大池化(channel=1)
  44.         max_out, _ = torch.max(x, dim=1, keepdim=True)
  45.         # channel拼接(channel=2)
  46.         x = torch.cat([avg_out, max_out], dim=1)
  47.         # channel=1
  48.         x = self.conv(x)
  49.         return self.sigmoid(x)
  50.  
  51.     
  52. class CBAMBottleneck(nn.Module):
  53.     # ch_in, ch_out, shortcut, groups, expansion, ratio, kernel_size
  54.     def __init__(self, c1, c2, shortcut=True, g=1, e=0.5, ratio=16, kernel_size=7):
  55.         super(CBAMBottleneck, self).__init__()
  56.         c_ = int(c2 * e)  # hidden channels
  57.         self.cv1 = Conv(c1, c_, 1, 1)
  58.         self.cv2 = Conv(c_, c2, 3, 1, g=g)
  59.         self.add = shortcut and c1 == c2
  60.         # 加入CBAM模块
  61.         self.channel_attention = ChannelAttention(c2, ratio)
  62.         self.spatial_attention = SpatialAttention(kernel_size)
  63.     
  64.     def forward(self, x):
  65.         # 考虑加入CBAM模块的位置：bottleneck模块刚开始时、bottleneck模块中shortcut之前，这里选择在shortcut之前
  66.         x2 = self.cv2(self.cv1(x))  # x和x2的channel数相同
  67.         # 在bottleneck模块中shortcut之前加入CBAM模块
  68.         out = self.channel_attention(x2) * x2
  69.         # print('outchannels:{}'.format(out.shape))
  70.         out = self.spatial_attention(out) * out
  71.         return x + out if self.add else out
  72.  
  73.  
  74. class C3CBAM(C3):
  75.     # C3 module with CBAMBottleneck()
  76.     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
  77.         super().__init__(c1, c2, n, shortcut, g, e)  # 引入C3(父类)的属性
  78.         c_ = int(c2 * e)  # hidden channels
  79.         self.m = nn.Sequential(*(CBAMBottleneck(c_, c_, shortcut) for _ in range(n)))
  80.  
  81.  
  82.

2）修改yolo.py
原代码为：

  1.  if m in {
  2.                 Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
  3.                 BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x}:
  4.             c1, c2 = ch[f], args[0]

在后面添加 CBAMBottleneck, C3CBAM

修改后为：

  1.  if m in {
  2.                 Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
  3.                 BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x,CBAMBottleneck, C3CBAM}:
  4.             c1, c2 = ch[f], args[0]

3）在models文件夹下创建yaml文件，命名为yolov5-cbam.yaml

内容为

  1. # YOLOv5 
    声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/羊村懒王/article/detail/290979
    推荐阅读
    相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。