当前位置:   article > 正文

YOLOv8改进算法之添加CA注意力机制_yolov8添加注意力机制

yolov8添加注意力机制

1. CA注意力机制

CA(Coordinate Attention)注意力机制是一种用于加强深度学习模型对输入数据的空间结构理解的注意力机制。CA 注意力机制的核心思想是引入坐标信息,以便模型可以更好地理解不同位置之间的关系。如下图:

4cf2f4e3554f46929a6bbb5db3e79108.png

1. 输入特征: CA 注意力机制的输入通常是一个特征图,它通常是卷积神经网络(CNN)中的某一层的输出,具有以下形状:[C, H, W],其中:

  • C 是通道数,表示特征图中的不同特征通道。
  • H 是高度,表示特征图的垂直维度。
  • W 是宽度,表示特征图的水平维度。

2. 全局平均池化: CA 注意力机制首先对输入特征图进行两次全局平均池化,一次在宽度方向上,一次在高度方向上。这两次操作分别得到两个特征映射:

  • 在宽度方向上的平均池化得到特征映射 [C, H, 1]
  • 在高度方向上的平均池化得到特征映射 [C, 1, W]

这两个特征映射分别捕捉了在宽度和高度方向上的全局特征。

3. 合并宽高特征: 将上述两个特征映射合并,通常通过简单的堆叠操作,得到一个新的特征层,形状为 [C, 1, H + W],其中 H + W 表示在宽度和高度两个方向上的维度合并在一起。

4. 卷积+标准化+激活函数: 对合并后的特征层进行卷积操作,通常是 1x1 卷积,以捕捉宽度和高度维度之间的关系。然后,通常会应用标准化(如批量标准化)和激活函数(如ReLU)来进一步处理特征,得到一个更加丰富的表示。

5. 再次分开: 分别从上述特征层中分离出宽度和高度方向的特征:

  • 一个分支得到特征层 [C, 1, H]
  • 另一个分支得到特征层 [C, 1, W]

6. 转置: 对分开的两个特征层进行转置操作,以恢复宽度和高度的维度,得到两个特征层分别为 [C, H, 1][C, 1, W]

7. 通道调整和 Sigmoid: 对两个分开的特征层分别应用 1x1 卷积,以调整通道数,使其适应注意力计算。然后,应用 Sigmoid 激活函数,得到在宽度和高度维度上的注意力分数。这些分数用于指示不同位置的重要性。

8. 应用注意力: 将原始输入特征图与宽度和高度方向上的注意力分数相乘,得到 CA 注意力机制的输出。

 

2. YOLOv8添加CA注意力机制

加入注意力机制,在ultralytics包中的nn包的modules里添加CA注意力模块,我这里选择在conv.py文件中添加CA注意力机制。

4334c4e3f5ae4109b59163920b5157a9.png

CA注意力机制代码如下:

  1. import torch
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  4. class h_sigmoid(nn.Module):
  5. def __init__(self, inplace=True):
  6. super(h_sigmoid, self).__init__()
  7. self.relu = nn.ReLU6(inplace=inplace)
  8. def forward(self, x):
  9. return self.relu(x + 3) / 6
  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. def forward(self, x):
  15. return x * self.sigmoid(x)
  16. class CoordAtt(nn.Module):
  17. def __init__(self, inp, reduction=32):
  18. super(CoordAtt, self).__init__()
  19. self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
  20. self.pool_w = nn.AdaptiveAvgPool2d((1, None))
  21. mip = max(8, inp // reduction)
  22. self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)
  23. self.bn1 = nn.BatchNorm2d(mip)
  24. self.act = h_swish()
  25. self.conv_h = nn.Conv2d(mip, inp, kernel_size=1, stride=1, padding=0)
  26. self.conv_w = nn.Conv2d(mip, inp, kernel_size=1, stride=1, padding=0)
  27. def forward(self, x):
  28. identity = x
  29. n, c, h, w = x.size()
  30. x_h = self.pool_h(x)
  31. x_w = self.pool_w(x).permute(0, 1, 3, 2)
  32. y = torch.cat([x_h, x_w], dim=2)
  33. y = self.conv1(y)
  34. y = self.bn1(y)
  35. y = self.act(y)
  36. x_h, x_w = torch.split(y, [h, w], dim=2)
  37. x_w = x_w.permute(0, 1, 3, 2)
  38. a_h = self.conv_h(x_h).sigmoid()
  39. a_w = self.conv_w(x_w).sigmoid()
  40. out = identity * a_w * a_h
  41. return out

CA注意力机制的注册和引用如下:

 ultralytics/nn/modules/_init_.py文件中:

b5965c8994e84095842122c96556e16b.jpeg

  ultralytics/nn/tasks.py文件夹中:

823386aa86634340aff6aa1701b3507e.png

 在tasks.py中的parse_model中添加如下代码:

  1. elif m in {CoordAtt}:
  2. args=[ch[f],*args]

新建相应的yolov8s-CA.yaml文件,代码如下:

  1. # Ultralytics YOLO 本文内容由网友自发贡献,转载请注明出处:【wpsshop博客】
    推荐阅读
    相关标签