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

yolo注意力机制手把手教学_yolo使用的注意力机制应该添加在哪个地方

作者：IT小白 | 2024-06-25 21:50:14

赞

踩

yolo使用的注意力机制应该添加在哪个地方

yolov5加入注意力机制步骤

1.common.py添加相应条件
2.yolo.py添加判断条件
3.创建属于自己的注意力yaml文件

第一步：

一、common.py

在common.py中先添加你想添加的注意力模块

  1. class h_sigmoid(nn.Module):
  2.     def __init__(self, inplace=True):
  3.         super(h_sigmoid, self).__init__()
  4.         self.relu = nn.ReLU6(inplace=inplace)
  5.  
  6.     def forward(self, x):
  7.         return self.relu(x + 3) / 6
  8.  
  9.  
  10. class h_swish(nn.Module):
  11.     def __init__(self, inplace=True):
  12.         super(h_swish, self).__init__()
  13.         self.sigmoid = h_sigmoid(inplace=inplace)
  14.  
  15.     def forward(self, x):
  16.         return x * self.sigmoid(x)
  17.  
  18.  
  19. class CoordAtt(nn.Module):
  20.     def __init__(self, inp, oup, reduction=32):
  21.         super(CoordAtt, self).__init__()
  22.         self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
  23.         self.pool_w = nn.AdaptiveAvgPool2d((1, None))
  24.  
  25.         mip = max(8, inp // reduction)
  26.  
  27.         self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)
  28.         self.bn1 = nn.BatchNorm2d(mip)
  29.         self.act = h_swish()
  30.  
  31.         self.conv_h = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
  32.         self.conv_w = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
  33.  
  34.     def forward(self, x):
  35.         identity = x
  36.  
  37.         n, c, h, w = x.size()
  38.         x_h = self.pool_h(x)
  39.         x_w = self.pool_w(x).permute(0, 1, 3, 2)
  40.  
  41.         y = torch.cat([x_h, x_w], dim=2)
  42.         y = self.conv1(y)
  43.         y = self.bn1(y)
  44.         y = self.act(y)
  45.  
  46.         x_h, x_w = torch.split(y, [h, w], dim=2)
  47.         x_w = x_w.permute(0, 1, 3, 2)
  48.  
  49.         a_h = self.conv_h(x_h).sigmoid()
  50.         a_w = self.conv_w(x_w).sigmoid()
  51.  
  52.         out = identity * a_w * a_h
  53.  
  54.         return out
  55.  
  56.  
  57. class SELayer(nn.Module):
  58.     def __init__(self, c1, r=16):
  59.         super(SELayer, self).__init__()
  60.         self.avgpool = nn.AdaptiveAvgPool2d(1)
  61.         self.l1 = nn.Linear(c1, c1 // r, bias=False)
  62.         self.relu = nn.ReLU(inplace=True)
  63.         self.l2 = nn.Linear(c1 // r, c1, bias=False)
  64.         self.sig = nn.Sigmoid()
  65.  
  66.     def forward(self, x):
  67.         b, c, _, _ = x.size()
  68.         y = self.avgpool(x).view(b, c)
  69.         y = self.l1(y)
  70.         y = self.relu(y)
  71.         y = self.l2(y)
  72.         y = self.sig(y)
  73.         y = y.view(b, c, 1, 1)
  74.         return x * y.expand_as(x)
  75.  
  76.  
  77. class eca_layer(nn.Module):
  78.     """Constructs a ECA module.
  79.     Args:
  80.         channel: Number of channels of the input feature map
  81.         k_size: Adaptive selection of kernel size
  82.     """
  83.  
  84.     def __init__(self, channel, k_size=3):
  85.         super(eca_layer, self).__init__()
  86.         self.avg_pool = nn.AdaptiveAvgPool2d(1)
  87.         self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
  88.         self.sigmoid = nn.Sigmoid()
  89.  
  90.     def forward(self, x):
  91.         # feature descriptor on the global spatial information
  92.         y = self.avg_pool(x)
  93.  
  94.         # Two different branches of ECA module
  95.         y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
  96.  
  97.         # Multi-scale information fusion
  98.         y = self.sigmoid(y)
  99.         x = x * y.expand_as(x)
  100.  
  101.         return x * y.expand_as(x)
  102.  
  103.  
  104. class ChannelAttention(nn.Module):
  105.     def __init__(self, in_planes, ratio=16):
  106.         super(ChannelAttention, self).__init__()
  107.         self.avg_pool = nn.AdaptiveAvgPool2d(1)
  108.         self.max_pool = nn.AdaptiveMaxPool2d(1)
  109.  
  110.         self.f1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
  111.         self.relu = nn.ReLU()
  112.         self.f2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
  113.         # 写法二,亦可使用顺序容器
  114.         # self.sharedMLP = nn.Sequential(
  115.         # nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False), nn.ReLU(),
  116.         # nn.Conv2d(in_planes // rotio, in_planes, 1, bias=False))
  117.  
  118.         self.sigmoid = nn.Sigmoid()
  119.  
  120.     def forward(self, x):
  121.         avg_out = self.f2(self.relu(self.f1(self.avg_pool(x))))
  122.         max_out = self.f2(self.relu(self.f1(self.max_pool(x))))
  123.         out = self.sigmoid(avg_out + max_out)
  124.         return out
  125.  
  126.  
  127. class SpatialAttention(nn.Module):
  128.     def __init__(self, kernel_size=7):
  129.         super(SpatialAttention, self).__init__()
  130.  
  131.         assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
  132.         padding = 3 if kernel_size == 7 else 1
  133.  
  134.         self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
  135.         self.sigmoid = nn.Sigmoid()
  136.  
  137.     def forward(self, x):
  138.         avg_out = torch.mean(x, dim=1, keepdim=True)
  139.         max_out, _ = torch.max(x, dim=1, keepdim=True)
  140.         x = torch.cat([avg_out, max_out], dim=1)
  141.         x = self.conv(x)
  142.         return self.sigmoid(x)
  143.  
  144.  
  145. class CBAMC3(nn.Module):
  146.     # CSP Bottleneck with 3 convolutions
  147.     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
  148.         super(CBAMC3, self).__init__()
  149.         c_ = int(c2 * e)  # hidden channels
  150.         self.cv1 = Conv(c1, c_, 1, 1)
  151.         self.cv2 = Conv(c1, c_, 1, 1)
  152.         self.cv3 = Conv(2 * c_, c2, 1)
  153.         self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
  154.         self.channel_attention = ChannelAttention(c2, 16)
  155.         self.spatial_attention = SpatialAttention(7)
  156.  
  157.         # self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
  158.  
  159.     def forward(self, x):
  160.         out = self.channel_attention(x) * x
  161.         print('outchannels:{}'.format(out.shape))
  162.         out = self.spatial_attention(out) * out
  163.         return out
  164.

二、yolo.py

在def parse_model(d, ch):函数下
在下面代码中增加你想添加的注意力名称

原始

  1.         if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
  2.                  BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]:
  3.             c1, c2 = ch[f], args[0]
  4.             if c2 != no:  # if not output
  5.                 c2 = make_divisible(c2 * gw, 8)
  6.  
  7.             args = [c1, c2, *args[1:]]
  8.             if m in [BottleneckCSP, C3, C3TR, C3Ghost]:
  9.                 args.insert(2, n)  # number of repeats
  10.                 n = 1

添加后

  1.         if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
  2.                  BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, CoordAtt]:
  3.             c1, c2 = ch[f], args[0]
  4.             if c2 != no:  # if not output
  5.                 c2 = make_divisible(c2 * gw, 8)
  6.  
  7.             args = [c1, c2, *args[1:]]
  8.             if m in [BottleneckCSP, C3, C3TR, C3Ghost]:
  9.                 args.insert(2, n)  # number of repeats
  10.                 n = 1

三.也是最后一步啦！创建自定义的yaml文件

这里我使用的是yolov5s.yaml为模板。注意力机制放置的位置并不是唯一的，需要根据你的数据集来摸索测试，可能别人放这儿涨点了，但是你放这儿没有效果，俗称“玄学”。
CA.yaml代码如下（示例）：

  1. # YOLOv5        本文内容由网友自发贡献，转载请注明出处：https://www.wpsshop.cn/w/IT小白/article/detail/757493
    推荐阅读
    相关标签

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