ToB企服应用市场:ToB评测及商务社交产业平台

标题: 如安在YOLOv8网络中添加自界说留意力机制 [打印本页]
作者: 商道如狼道    时间: 2024-11-4 20:00
标题: 如安在YOLOv8网络中添加自界说留意力机制
在目标检测任务中，加入留意力机制可以提升模型的检测结果。本文将介绍如安在YOLOv8模型中集成多种留意力机制，如 SimAM、ShuffleAttention、TripletAttention、MHSA、CBAM 和 EMA，以增强模型对图像特征的提取能力。我们将展示每个留意力机制的代码示例，并讨论如何将这些模块添加到YOLOv8网络中。
目次


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



欢迎光临 ToB企服应用市场:ToB评测及商务社交产业平台 (https://dis.qidao123.com/) Powered by Discuz! X3.4