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

改进yolov8: yolov8添加注意力机制_yolov8可以在哪里加入注意力机制

作者：我家小花儿 | 2024-04-14 23:22:43

赞

踩

yolov8可以在哪里加入注意力机制

目录

1.  添加注意力机制

2.  注意力机制的注册和调用

3.  修改yaml配置文件

4.  注意力机制代码(GAM/EMA/Triplet)

1.  添加注意力机制

在ultralytics-main/ultralytics/nn/modules/conv.py中添加EMA注意力机制模块：

2.  注意力机制的注册和调用

  • 在ultralytics-main/ultralytics/nn/modules/__init__.py中注册引用EMA注意力机制:

  • 在ultralytics-main/ultralytics/nn/tasks.py中引用:

  • 在tasks.py中写入调用方式：

  1.  
  2.  # """**************add Attention***************"""
  3.         elif m in {GAM_Attention, EMA}:
  4.             c1, c2 = ch[f], args[0]
  5.             if c2 != nc:  # if not output
  6.                 c2 = make_divisible(min(c2, max_channels) * width, 8)
  7.             args = [c1, c2, *args[1:]]
  8.         
  9.         elif m is TripletAttention:
  10.             c1, c2 = ch[f], args[0]
  11.             if c2 != nc:
  12.                 c2 = make_divisible(min(c2, max_channels) * width, 8)
  13.             args = [c1, *args[1:]]

3.  修改yaml配置文件

在ultralytics-main/ultralytics/models/v8/yolov8s-p6-attention.yaml配置文件：

  1. # Ultralytics YOLO 
  2. # YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
  3.  
  4. # Parameters
  5. nc: 80 # number of classes
  6. scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8-SPPCSPC.yaml with scale 'n'
  7.  # [depth, width, max_channels]
  8.  n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers,  3157200 parameters,  3157184 gradients,  8.9 GFLOPs
  9.  s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients,  28.8 GFLOPs
  10.  m: [0.67, 0.75, 768]  # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients,  79.3 GFLOPs
  11.  l: [1.00, 1.00, 512]  # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
  12.  x: [1.00, 1.25, 512]  # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs
  13.  
  14. # YOLOv8.0x6 backbone
  15. backbone:
  16.  # [from, repeats, module, args]
  17.  - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
  18.  - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
  19.  - [-1, 3, C2f, [128, True]]
  20.  - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
  21.  - [-1, 6, C2f, [256, True]]
  22.  - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
  23.  - [-1, 6, C2f, [512, True]]
  24.  - [-1, 1, Conv, [768, 3, 2]] # 7-P5/32
  25.  - [-1, 3, C2f, [768, True]]
  26.  - [-1, 1, Conv, [1024, 3, 2]] # 9-P6/64
  27.  - [-1, 3, C2f, [1024, True]]
  28.  - [-1, 1, SPPF, [1024, 5]] # 11
  29.  
  30. # YOLOv8.0x6 head
  31. head:
  32.   - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  33.   - [[-1, 8], 1, Concat, [1]] # cat backbone P5
  34.   - [-1, 3, C2, [768, False]] # 14
  35.  
  36.   - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  37.   - [[-1, 6], 1, Concat, [1]] # cat backbone P4
  38.   - [-1, 3, C2, [512, False]] # 17
  39.  
  40.   - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  41.   - [[-1, 4], 1, Concat, [1]] # cat backbone P3
  42.   - [-1, 3, C2, [256, False]] # 20 (P3/8-small)
  43.   - [-1, 1, GAM_Attention, [256, 256]]
  44.   # - [-1, 1, TripletAttention, [256]]
  45.   # - [-1, 1, EMA, [256]]
  46.  
  47.   - [-1, 1, Conv, [256, 3, 2]]
  48.   - [[-1, 17], 1, Concat, [1]] # cat head P4
  49.   - [-1, 3, C2, [512, False]] # 23 (P4/16-medium)
  50.   - [-1, 1, GAM_Attention, [512, 512]]
  51.   # - [-1, 1, TripletAttention, [512]]
  52.   # - [-1, 1, EMA, [512]]
  53.  
  54.   - [-1, 1, Conv, [512, 3, 2]]
  55.   - [[-1, 14], 1, Concat, [1]] # cat head P5
  56.   - [-1, 3, C2, [768, False]] # 26 (P5/32-large)
  57.   - [-1, 1, GAM_Attention, [768, 768]]
  58.   # - [-1, 1, TripletAttention, [768]]
  59.   # - [-1, 1, EMA, [768]]
  60.  
  61.   - [-1, 1, Conv, [768, 3, 2]]
  62.   - [[-1, 11], 1, Concat, [1]] # cat head P6
  63.   - [-1, 3, C2, [1024, False]] # 29 (P6/64-xlarge)
  64.   - [-1, 1, GAM_Attention, [1024, 1024]]
  65.   # - [-1, 1, TripletAttention, [1024]]
  66.   # - [-1, 1, EMA, [1024]]
  67.  
  68.   - [[21, 25, 29, 33], 1, Detect, [nc]] # Detect(P3, P4, P5, P6)
说明：上述注意力机制可以在head中使用，同样可以添加在backbone中，可以多次尝试选择最适合自己的任务的方式，但需要注意head中的concat层是否准确。

4.  注意力机制代码(GAM/EMA/Triplet)

  • GAM

  1. def channel_shuffle(x, groups=2):  ##shuffle channel
  2.     # RESHAPE----->transpose------->Flatten
  3.     B, C, H, W = x.size()
  4.     out = x.view(B, groups, C // groups, H, W).permute(0, 2, 1, 3, 4).contiguous()
  5.     out = out.view(B, C, H, W)
  6.     return out
  7.  
  8. class GAM_Attention(nn.Module):
  9.     # https://paperswithcode.com/paper/global-attention-mechanism-retain-information
  10.     def __init__(self, c1, c2, group=True, rate=4):
  11.         super(GAM_Attention, self).__init__()
  12.  
  13.         self.channel_attention = nn.Sequential(
  14.             nn.Linear(c1, int(c1 / rate)),
  15.             nn.ReLU(inplace=True),
  16.             nn.Linear(int(c1 / rate), c1)
  17.         )
  18.  
  19.         self.spatial_attention = nn.Sequential(
  20.  
  21.             nn.Conv2d(c1, c1 // rate, kernel_size=7, padding=3, groups=rate) if group else nn.Conv2d(c1, int(c1 / rate),
  22.                                                                                                      kernel_size=7,
  23.                                                                                                      padding=3),
  24.             nn.BatchNorm2d(int(c1 / rate)),
  25.             nn.ReLU(inplace=True),
  26.             nn.Conv2d(c1 // rate, c2, kernel_size=7, padding=3, groups=rate) if group else nn.Conv2d(int(c1 / rate), c2,
  27.                                                                                                      kernel_size=7,
  28.                                                                                                      padding=3),
  29.             nn.BatchNorm2d(c2)
  30.         )
  31.  
  32.     def forward(self, x):
  33.         b, c, h, w = x.shape
  34.         x_permute = x.permute(0, 2, 3, 1).view(b, -1, c)
  35.         x_att_permute = self.channel_attention(x_permute).view(b, h, w, c)
  36.         x_channel_att = x_att_permute.permute(0, 3, 1, 2)
  37.         # x_channel_att=channel_shuffle(x_channel_att,4) #last shuffle
  38.         x = x * x_channel_att
  39.  
  40.         x_spatial_att = self.spatial_attention(x).sigmoid()
  41.         x_spatial_att = channel_shuffle(x_spatial_att, 4)  # last shuffle
  42.         out = x * x_spatial_att
  43.         # out=channel_shuffle(out,4) #last shuffle
  44.         return out

  • EMA

  1. class EMA(nn.Module):
  2.     def __init__(self, channels, c2=None, factor=32):
  3.         super(EMA, self).__init__()
  4.         self.groups = factor
  5.         assert channels // self.groups > 0
  6.         self.softmax = nn.Softmax(-1)
  7.         self.agp = nn.AdaptiveAvgPool2d((1, 1))
  8.         self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
  9.         self.pool_w = nn.AdaptiveAvgPool2d((1, None))
  10.         self.gn = nn.GroupNorm(channels // self.groups, channels // self.groups)
  11.         self.conv1x1 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=1, stride=1, padding=0)
  12.         self.conv3x3 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=3, stride=1, padding=1)
  13.  
  14.     def forward(self, x):
  15.         b, c, h, w = x.size()
  16.         group_x = x.reshape(b * self.groups, -1, h, w)  # b*g,c//g,h,w
  17.         x_h = self.pool_h(group_x)
  18.         x_w = self.pool_w(group_x).permute(0, 1, 3, 2)
  19.         hw = self.conv1x1(torch.cat([x_h, x_w], dim=2))
  20.         x_h, x_w = torch.split(hw, [h, w], dim=2)
  21.         x1 = self.gn(group_x * x_h.sigmoid() * x_w.permute(0, 1, 3, 2).sigmoid())
  22.         x2 = self.conv3x3(group_x)
  23.         x11 = self.softmax(self.agp(x1).reshape(b * self.groups, -1, 1).permute(0, 2, 1))
  24.         x12 = x2.reshape(b * self.groups, c // self.groups, -1)  # b*g, c//g, hw
  25.         x21 = self.softmax(self.agp(x2).reshape(b * self.groups, -1, 1).permute(0, 2, 1))
  26.         x22 = x1.reshape(b * self.groups, c // self.groups, -1)  # b*g, c//g, hw
  27.         weights = (torch.matmul(x11, x12) + torch.matmul(x21, x22)).reshape(b * self.groups, 1, h, w)
  28.         return (group_x * weights.sigmoid()).reshape(b, c, h, w)

  • Triplet

    1. class BasicConv(nn.Module):   #https://arxiv.org/pdf/2010.03045.pdf
    2.     def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True,
    3.                  bn=True, bias=False):
    4.         super(BasicConv, self).__init__()
    5.         self.out_channels = out_planes
    6.         self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding,
    7.                               dilation=dilation, groups=groups, bias=bias)
    8.         self.bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.01, affine=True) if bn else None
    9.         self.relu = nn.ReLU() if relu else None
    10.  
    11.     def forward(self, x):
    12.         x = self.conv(x)
    13.         if self.bn is not None:
    14.             x = self.bn(x)
    15.         if self.relu is not None:
    16.             x = self.relu(x)
    17.         return x
    18.  
    19.  
    20. class ZPool(nn.Module):
    21.     def forward(self, x):
    22.         return torch.cat((torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1)
    23.  
    24.  
    25. class AttentionGate(nn.Module):
    26.     def __init__(self):
    27.         super(AttentionGate, self).__init__()
    28.         kernel_size = 7
    29.         self.compress = ZPool()
    30.         self.conv = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size - 1) // 2, relu=False)
    31.  
    32.     def forward(self, x):
    33.         x_compress = self.compress(x)
    34.         x_out = self.conv(x_compress)
    35.         scale = torch.sigmoid_(x_out)
    36.         return x * scale
    37.  
    38.  
    39. class TripletAttention(nn.Module):
    40.     def __init__(self, no_spatial=False):
    41.         super(TripletAttention, self).__init__()
    42.         self.cw = AttentionGate()
    43.         self.hc = AttentionGate()
    44.         self.no_spatial = no_spatial
    45.         if not no_spatial:
    46.             self.hw = AttentionGate()
    47.  
    48.     def forward(self, x):
    49.         x_perm1 = x.permute(0, 2, 1, 3).contiguous()
    50.         x_out1 = self.cw(x_perm1)
    51.         x_out11 = x_out1.permute(0, 2, 1, 3).contiguous()
    52.         x_perm2 = x.permute(0, 3, 2, 1).contiguous()
    53.         x_out2 = self.hc(x_perm2)
    54.         x_out21 = x_out2.permute(0, 3, 2, 1).contiguous()
    55.         if not self.no_spatial:
    56.             x_out = self.hw(x)
    57.             x_out = 1 / 3 * (x_out + x_out11 + x_out21)
    58.         else:
    59.             x_out = 1 / 2 * (x_out11 + x_out21)
    60.         return x_out

本文内容由网友自发贡献，转载请注明出处：https://www.wpsshop.cn/w/我家小花儿/article/detail/424628
推荐阅读
相关标签

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

  

闽ICP备14008679号