如安在YOLOv8网络中添加自界说留意力机制

在目标检测任务中，加入留意力机制可以提升模型的检测结果。本文将介绍如安在YOLOv8模型中集成多种留意力机制，如 SimAM、ShuffleAttention、TripletAttention、MHSA、CBAM 和 EMA，以增强模型对图像特征的提取能力。我们将展示每个留意力机制的代码示例，并讨论如何将这些模块添加到YOLOv8网络中。

---

目次

• 1. 留意力机制示例代码
  
  
  
  • 1.1 SimAM 模块代码
    
  • 1.2 ShuffleAttention 模块代码
    
  • 1.3 TripletAttention 模块代码
    
  • 1.4 MHSA 模块代码
    
  • 1.5 CBAM 模块代码
    
  • 1.6 EMA 模块代码
    
  • 1.7 ECA 模块代码
    
  
  
• 2. 添加留意力机制的步骤
  
  
  
  • 2.1 修改YOLOv8的配置文件
    
  • 2.2 编写自界说留意力机制模块
    
  • 2.3 训练和验证留意力机制
    
  
  
• 3. 总结
  

---

1. 留意力机制示例代码

下面介绍六种常用的留意力机制模块，并提供代码示例。每种留意力机制都有其独特的优点，可以根据任务需求选择最适合的机制。
1.1 SimAM 模块代码

SimAM (Simple Attention Module) 是一种轻量级的留意力机制；通过简单的操纵实现了留意力机制的结果，适用于对计算资源敏感的项目。它适合那些盼望在提升模型性能的同时，尽量淘汰计算开销的任务，比如嵌入式设备上的实时目标检测。。

1. import torch
2. import torch.nn as nn
5. class SimAM(torch.nn.Module):
6.     def __init__(self, e_lambda=1e-4):
7.         super(SimAM, self).__init__()
9.         self.activaton = nn.Sigmoid()
10.         self.e_lambda = e_lambda
12.     def __repr__(self):
13.         s = self.__class__.__name__ + '('
14.         s += ('lambda=%f)' % self.e_lambda)
15.         return s
17.     @staticmethod
18.     def get_module_name():
19.         return "simam"
21.     def forward(self, x):
22.         b, c, h, w = x.size()
24.         n = w * h - 1
26.         x_minus_mu_square = (x - x.mean(dim=[2, 3], keepdim=True)).pow(2)
27.         y = x_minus_mu_square / (4 * (x_minus_mu_square.sum(dim=[2, 3], keepdim=True) / n + self.e_lambda)) + 0.5
29.         return x * self.activaton(y)
32. # if __name__ == '__main__':
33. #     input = torch.randn(3, 64, 7, 7)
34. #     model = SimAM()
35. #     outputs = model(input)
36. #     print(outputs.shape)

复制代码

1.2 ShuffleAttention 模块代码

ShuffleAttention 适合需要全局特征交互的场景；通过通道洗牌操纵重新排列特征，确保模型能够在不同通道间传递信息，提升特征的全局表达能力。对于需要处理复杂、具有多样性特征的图像（如交通场景、复杂的自然环境），这种机制能有用提升模型的感知能力。

1. import torch
2. from torch import nn
3. from torch.nn import init
4. from torch.nn.parameter import Parameter
6. class ShuffleAttention(nn.Module):
8.     def __init__(self, channel=512, reduction=16, G=8):
9.         super().__init__()
10.         self.G = G
11.         self.channel = channel
12.         self.avg_pool = nn.AdaptiveAvgPool2d(1)
13.         self.gn = nn.GroupNorm(channel // (2 * G), channel // (2 * G))
14.         self.cweight = Parameter(torch.zeros(1, channel // (2 * G), 1, 1))
15.         self.cbias = Parameter(torch.ones(1, channel // (2 * G), 1, 1))
16.         self.sweight = Parameter(torch.zeros(1, channel // (2 * G), 1, 1))
17.         self.sbias = Parameter(torch.ones(1, channel // (2 * G), 1, 1))
18.         self.sigmoid = nn.Sigmoid()
20.     def init_weights(self):
21.         for m in self.modules():
22.             if isinstance(m, nn.Conv2d):
23.                 init.kaiming_normal_(m.weight, mode='fan_out')
24.                 if m.bias is not None:
25.                     init.constant_(m.bias, 0)
26.             elif isinstance(m, nn.BatchNorm2d):
27.                 init.constant_(m.weight, 1)
28.                 init.constant_(m.bias, 0)
29.             elif isinstance(m, nn.Linear):
30.                 init.normal_(m.weight, std=0.001)
31.                 if m.bias is not None:
32.                     init.constant_(m.bias, 0)
34.     @staticmethod
35.     def channel_shuffle(x, groups):
36.         b, c, h, w = x.shape
37.         x = x.reshape(b, groups, -1, h, w)
38.         x = x.permute(0, 2, 1, 3, 4)
40.         # flatten
41.         x = x.reshape(b, -1, h, w)
43.         return x
45.     def forward(self, x):
46.         b, c, h, w = x.size()
47.         # group into subfeatures
48.         x = x.view(b * self.G, -1, h, w)  # bs*G,c//G,h,w
50.         # channel_split
51.         x_0, x_1 = x.chunk(2, dim=1)  # bs*G,c//(2*G),h,w
53.         # channel attention
54.         x_channel = self.avg_pool(x_0)  # bs*G,c//(2*G),1,1
55.         x_channel = self.cweight * x_channel + self.cbias  # bs*G,c//(2*G),1,1
56.         x_channel = x_0 * self.sigmoid(x_channel)
58.         # spatial attention
59.         x_spatial = self.gn(x_1)  # bs*G,c//(2*G),h,w
60.         x_spatial = self.sweight * x_spatial + self.sbias  # bs*G,c//(2*G),h,w
61.         x_spatial = x_1 * self.sigmoid(x_spatial)  # bs*G,c//(2*G),h,w
63.         # concatenate along channel axis
64.         out = torch.cat([x_channel, x_spatial], dim=1)  # bs*G,c//G,h,w
65.         out = out.contiguous().view(b, -1, h, w)
67.         # channel shuffle
68.         out = self.channel_shuffle(out, 2)
69.         return out
72. if __name__ == '__main__':
73.     input = torch.randn(50, 512, 7, 7)
74.     se = ShuffleAttention(channel=512, G=8)
75.     output = se(input)

复制代码

1.3 TripletAttention 模块代码

TripletAttention 适合需要捕捉多方向特征的场景；在通道上引入了三个方向的留意力（水平、垂直、深度），能够帮助模型更好地感知多方向上的特征。这种机制特别适用于那些需要捕捉方向性信息的任务，比如门路标记检测和自然场景理解。

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

复制代码

1.4 MHSA 模块代码

MHSA (Multi-Head Self-Attention) 是常用于Transformer模型的留意力机制，适合大规模上下文建模的场景；通过多头自留意力的机制，能够帮助模型捕捉图像中的长距离依赖关系。它在需要处理上下文信息的任务中体现出色，如自然场景中的多物体检测。对于需要全局信息并且图像内物体之间具有复杂相互关系的任务，MHSA 是理想的选择。

1. import torch
2. import torch.nn as nn
4. class MHSA(nn.Module):
5.     def __init__(self, n_dims, width=14, height=14, heads=4, pos_emb=False):
6.         super(MHSA, self).__init__()
8.         self.heads = heads
9.         self.query = nn.Conv2d(n_dims, n_dims, kernel_size=1)
10.         self.key = nn.Conv2d(n_dims, n_dims, kernel_size=1)
11.         self.value = nn.Conv2d(n_dims, n_dims, kernel_size=1)
12.         self.pos = pos_emb
13.         if self.pos:
14.             self.rel_h_weight = nn.Parameter(torch.randn([1, heads, (n_dims) // heads, 1, int(height)]),
15.                                              requires_grad=True)
16.             self.rel_w_weight = nn.Parameter(torch.randn([1, heads, (n_dims) // heads, int(width), 1]),
17.                                              requires_grad=True)
18.         self.softmax = nn.Softmax(dim=-1)
20.     def forward(self, x):
21.         n_batch, C, width, height = x.size()
22.         q = self.query(x).view(n_batch, self.heads, C // self.heads, -1)
23.         k = self.key(x).view(n_batch, self.heads, C // self.heads, -1)
24.         v = self.value(x).view(n_batch, self.heads, C // self.heads, -1)
25.         content_content = torch.matmul(q.permute(0, 1, 3, 2), k)  # 1,C,h*w,h*w
26.         c1, c2, c3, c4 = content_content.size()
27.         if self.pos:
28.             content_position = (self.rel_h_weight + self.rel_w_weight).view(1, self.heads, C // self.heads, -1).permute(
29.                 0, 1, 3, 2)  # 1,4,1024,64
31.             content_position = torch.matmul(content_position, q)  # ([1, 4, 1024, 256])
32.             content_position = content_position if (
33.                     content_content.shape == content_position.shape) else content_position[:, :, :c3, ]
34.             assert (content_content.shape == content_position.shape)
35.             energy = content_content + content_position
36.         else:
37.             energy = content_content
38.         attention = self.softmax(energy)
39.         out = torch.matmul(v, attention.permute(0, 1, 3, 2))  # 1,4,256,64
40.         out = out.view(n_batch, C, width, height)
41.         return out
43. # if __name__ == '__main__':
44. #     input = torch.randn(50, 512, 7, 7)
45. #     mhsa = MHSA(n_dims=512)
46. #     output = mhsa(input)
47. #     print(output.shape)

复制代码

1.5 CBAM 模块代码

CBAM（Convolutional Block Attention Module）适合需要联合通道和空间特征的场景；通过联合通道留意力和空间留意力，帮助网络更加精准地捕捉图像中的关键区域。它适用于大多数量标检测任务，特别是当需要细化某些特定物体的检测时，比如在自动驾驶中的行人检测或交通标记检测。

1. import torch
2. from torch import nn
4. class ChannelAttention(nn.Module):
5.     # Channel-attention module https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet
6.     def __init__(self, channels: int) -> None:
7.         super().__init__()
8.         self.pool = nn.AdaptiveAvgPool2d(1)
9.         self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
10.         self.act = nn.Sigmoid()
12.     def forward(self, x: torch.Tensor) -> torch.Tensor:
13.         return x * self.act(self.fc(self.pool(x)))
16. class SpatialAttention(nn.Module):
17.     # Spatial-attention module
18.     def __init__(self, kernel_size=7):
19.         super().__init__()
20.         assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
21.         padding = 3 if kernel_size == 7 else 1
22.         self.cv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
23.         self.act = nn.Sigmoid()
25.     def forward(self, x):
26.         return x * self.act(self.cv1(torch.cat([torch.mean(x, 1, keepdim=True), torch.max(x, 1, keepdim=True)[0]], 1)))
29. class CBAM(nn.Module):
30.     # Convolutional Block Attention Module
31.     def __init__(self, c1, kernel_size=7):  # ch_in, kernels
32.         super().__init__()
33.         self.channel_attention = ChannelAttention(c1)
34.         self.spatial_attention = SpatialAttention(kernel_size)
36.     def forward(self, x):
37.         return self.spatial_attention(self.channel_attention(x))

复制代码

1.6 EMA 模块代码

EMA（Efficient Multi-Head Attention）适合盼望在多头自留意力中提升服从的场景；它通过淘汰计算复杂度而提升性能，适用于大规模数据集的训练。它在保持留意力机制强盛的特征捕捉能力的同时，还能显著低落计算成本，适合高性能要求的任务场景。

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

复制代码

1.7 ECA 模块代码

ECA（Efficient Channel Attention）适合需要高效通道留意力的场景；通过消除全毗连层，利用1D卷积进行局部交互，大大淘汰了参数量，同时仍旧保留了通道留意力的能力。它适合那些对计算资源有限定的场景，比如移动设备上进行目标检测的任务。

1. import torch
2. from torch import nn
4. class ECA(nn.Module):
5.     def __init__(self, channels: int, k_size: int = 3):
6.         super(ECA, self).__init__()
7.         self.avg_pool = nn.AdaptiveAvgPool2d(1)
8.         self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
9.         self.sigmoid = nn.Sigmoid()
11.     def forward(self, x):
12.         # Apply global average pooling
13.         y = self.avg_pool(x)
15.         # Reshape and apply 1D convolution
16.         y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
18.         # Apply sigmoid activation and element-wise multiplication
19.         return x * self.sigmoid(y)

复制代码

---

2. 添加留意力机制的步骤

2.1 修改YOLOv8的配置文件

我们可以通过在YOLOv8配置文件中指定利用留意力机制，以下是如安在第10层加入留意力机制的配置示例，以ShuffleAttention留意力机制为例，用到哪个放开哪个：
[code]# Ultralytics YOLO