当前位置:   article > 正文

YOLOV5 添加注意力机制,以添加ECAttention为例(系列)

YOLOV5 添加注意力机制,以添加ECAttention为例(系列)

总述:在yolov5的网络结构中添加相关的注意力机制已经不是什么新鲜的事儿了,各种新奇的注意力机制近年来也是层出不穷,当然,选择什么类型的注意力机制、采取什么分布规范的数据集以及将注意力机制添加在网络的什么位置都会或多或少地对结果如P\R\mAP等产生好的、坏的影响,这里的确是一千个读者有一千个哈姆雷特,接下来从代码的角度讲详细的介绍如何在主干网络中添加ECA注意机制。

1、在common.py函数里面添加ECA代码。

2、在yolo.py中注册ECA。

3、修改对应的yaml文件。

步骤:

1、在common.py函数里面添加ECA代码。

  1. class ECA(nn.Module):
  2. """Constructs a ECA module.
  3. Args:
  4. channel: Number of channels of the input feature map
  5. k_size: Adaptive selection of kernel size
  6. """
  7. def __init__(self, c1,c2, k_size=3):
  8. super(ECA, self).__init__()
  9. self.avg_pool = nn.AdaptiveAvgPool2d(1)
  10. self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
  11. self.sigmoid = nn.Sigmoid()
  12. def forward(self, x):
  13. # feature descriptor on the global spatial information
  14. y = self.avg_pool(x)
  15. # print(y.shape,y.squeeze(-1).shape,y.squeeze(-1).transpose(-1, -2).shape)
  16. # Two different branches of ECA module
  17. # 50*C*1*1
  18. #50*C*1
  19. #50*1*C
  20. y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
  21. # Multi-scale information fusion
  22. y = self.sigmoid(y)
  23. return x * y.expand_as(x)

2、在yolo.py中注册ECA。对yolo.py详细修改见下面代码。

  1. if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP,
  2. C3, C3TR]:
  3. c1, c2 = ch[f], args[0]
  4. if c2 != no: # if not output
  5. c2 = make_divisible(c2 * gw, 8)
  6. args = [c1, c2, *args[1:]]
  7. if m in [BottleneckCSP, C3,eca_layer]:
  8. args.insert(2, n) # number of repeats
  9. n = 1
  10. elif m is nn.BatchNorm2d:
  11. args = [ch[f]]
  12. elif m is Concat:
  13. c2 = sum([ch[x] for x in f])
  14. elif m is Detect:
  15. args.append([ch[x] for x in f])
  16. if isinstance(args[1], int): # number of anchors
  17. args[1] = [list(range(args[1] * 2))] * len(f)
  18. elif m is Contract:
  19. c2 = ch[f] * args[0] ** 2
  20. elif m is Expand:
  21. c2 = ch[f] // args[0] ** 2
  22. elif m is eca_layer:
  23. channel=args[0]
  24. channel=make_divisible(channel*gw,8)if channel != no else channel
  25. args=[channel]
  26. else:
  27. c2 = ch[f]

3、修改对应的yaml文件。

  1. # parameters
  2. nc: 2 # number of classes
  3. depth_multiple: 0.33 # model depth multiple
  4. width_multiple: 0.50 # layer channel multiple
  5. # anchors
  6. anchors:
  7. - [10,13, 16,30, 33,23] # P3/8
  8. - [30,61, 62,45, 59,119] # P4/16
  9. - [116,90, 156,198, 373,326] # P5/32
  10. # YOLOv5 backbone
  11. backbone:
  12. # [from, number, module, args]
  13. [[-1, 1, Focus, [64, 3]], # 0-P1/2
  14. [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
  15. [-1, 3, C3, [128]],
  16. [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
  17. [-1, 9, C3, [256]],
  18. [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
  19. [-1, 9, C3, [512]],
  20. [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
  21. [-1, 1, SPP, [1024, [5, 9, 13]]],
  22. [-1, 3, C3, [1024, False]], # 9
  23. [-1,1,ECA,[1024]], #ECA
  24. ]
  25. # YOLOv5 head
  26. head:
  27. [[-1, 1, Conv, [512, 1, 1]],
  28. [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  29. [[-1, 6], 1, Concat, [1]], # cat backbone P4
  30. [-1, 3, C3, [512, False]], # 13
  31. [-1, 1, Conv, [256, 1, 1]],
  32. [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  33. [[-1, 4], 1, Concat, [1]], # cat backbone P3
  34. [-1, 3, C3, [256, False]], # 17 (P3/8-small)
  35. [-1, 1, Conv, [256, 3, 2]],
  36. [[-1, 15], 1, Concat, [1]], # cat head P4
  37. [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
  38. [-1, 1, Conv, [512, 3, 2]],
  39. [[-1, 11], 1, Concat, [1]], # cat head P5
  40. [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
  41. [[18, 21, 24], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
  42. ]

自学心得:正如开头所述,添加什么、添加何处、采用什么数据集等等一切取决于你的实际工作是什么。最后祝大家多发SCI!!!

欢迎讨论,有问题直接评论区!看到我会第一时间解答,同时也欢迎各位积极讨论!!!

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

闽ICP备14008679号