YOLOv8添加多种attention注意力机制

yolov8发布了快半年的时间，由于它还没有一个固定的版本，所以它的源码可能会和我下载的有所差异，但是大体上应该是差不多的。这里就不赘述yolov8相对于yolov5改了哪些东西，主要的就是一些像C3模块优化为C2f，SPP优化为SPPF等等。下面就讲一下如何在YOLOv8中加入一些自己的attention注意力机制。

这里我用CBAM这个attention举例，它是2018年提出了一种轻量的注意力模块( CBAM，Convolutional Block Attention Module )，可以在通道和空间维度上进行 Attention ，它对小目标检测有较好的鲁棒性。

论文题目：《CBAM: Convolutional Block Attention Module》
论文地址：https://arxiv.org/pdf/1807.06521.pdf

一、复制注意力机制代码

我这里有十几种不同注意力机制的代码，由于复制太麻烦了，所以就用这个CBAM的作为例子，想要其他的可以私信我哦

import numpy as np
import torch
from torch import nn
from torch.nn import init


class ChannelAttention(nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()
        self.maxpool = nn.AdaptiveMaxPool2d(1)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.se = nn.Sequential(
            nn.Conv2d(channel, channel // reduction, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(channel // reduction, channel, 1, bias=False)
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_result = self.maxpool(x)
        avg_result = self.avgpool(x)
        max_out = self.se(max_result)
        avg_out = self.se(avg_result)
        output = self.sigmoid(max_out + avg_out)
        return output


class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_result, _ = torch.max(x, dim=1, keepdim=True)
        avg_result = torch.mean(x, dim=1, keepdim=True)
        result = torch.cat([max_result, avg_result], 1)
        output = self.conv(result)
        output = self.sigmoid(output)
        return output


class CBAMBlock(nn.Module):

    def __init__(self, channel=512, reduction=16, kernel_size=7):
        super().__init__()
        self.ca = ChannelAttention(channel=channel, reduction=reduction)
        self.sa = SpatialAttention(kernel_size=kernel_size)

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        b, c, _, _ = x.size()
        out = x * self.ca(x)
        out = out * self.sa(out)
        return out

在ultralytics/nn这个目录下新建一个py文件，将复制的attention代码粘贴进去

然后在ultralytics/nn这个目录下的task.py修改一些信息

首先在开头导入一下这个模块

 注意from ultralytics.nn.后面的是你粘贴代码那个py文件的名字

然后在解析函数中加入以下代码

        elif m in {CBAMBlock}:
            args = [ch[f], *args]

大约在650行的位置，{}这个里面的是你要导入类的名称，这里还有一点就是有的注意力机制有通道数，有的不接收通道数，这个后面在详细说，目前这个按照这些步骤是没有问题的

 到这里就差不多了，这不像添加那些DCNv2、v3需要去__init__中声明一下，然后就可以修改yaml配置文件了

 我这个是在Backbone最后一层添加的，CBAMBlock就是attention中最后那个类的名字，[16,7]就是默认接受的参数，这里不填好像也可以，用的就是代码中默认的参数

至此，网络结构就修改完成了，然后就可以训练了，可以在终端输出的网络结构中看看自己的注意力机制有没有成功添加进去。