YOLOv8改进实战 | 引入多维协作注意模块MCA，实现暴力涨点
YOLOv8专栏导航：点击此处跳转
前言
YOLOv8 是由 YOLOv5 的发布者 Ultralytics 发布的最新版本的 YOLO。它可用于对象检测、分割、分类任务以及大型数据集的学习，并且可以在包括 CPU 和 GPU 在内的各种硬件上执行。
YOLOv8是一种尖端的、最先进的 (SOTA) 模型，它建立在以前成功的 YOLO 版本的基础上，并引入了新的功能和改进，以进一步提高性能和灵活性。YOLOv8 旨在快速、准确且易于使用，这也使其成为对象检测、图像分割和图像分类任务的绝佳选择。具体创新包括一个新的骨干网络、一个新的 Ancher-Free 检测头和一个新的损失函数，还支持YOLO以往版本，方便不同版本切换和性能对比。
目录
- 一、MCANet介绍
- 二、代码实现
- 代码目录
- 注册模块
- 配置yaml文件
- 三、模型测试
- 四、模型训练
- 五、总结
一、MCANet介绍
论文链接：MCA: Multidimensional collaborative attention in deep convolutional neural networks for image recognition
Pytorch code：MCANet
作者提出了一种基于高效轴向注意力的多尺度交叉轴注意（MCA）方法。MCA通过计算两个并行轴向注意力之间的双向交叉注意力，以更好地捕获全局信息。此外，为了处理病变区域或器官在个体大小和形状上的显著变化，还在每个轴向注意力路径中使用不同大小的条形卷积核进行多次卷积，以提高编码空间信息的效率。
顶部分支用于捕捉空间维度 W 中特征之间的交互作用。同样，中间分支用于捕捉空间维度 H 中特征之间的交互作用。底部分支负责捕捉通道之间的交互作用。在前两个分支中，我们使用置换操作来捕捉通道维度和两个空间维度之间的长程依赖关系。最后，在集成阶段，来自所有三个分支的输出通过简单平均进行聚合。
二、代码实现
代码目录
- 按下面文件夹结构创建文件（相比于在原有
ultralytics/nn/modules文件夹下的相关文件中直接添加更便于管理）- ultralytics- nn- extra_modules- __init__.py- attention.py- modules
在ultralytics/nn/extra_modules/__init__.py中添加：
from .attention import *
在ultralytics/nn/extra_modules/attention.py中添加：
import torch
from torch import nn
import math__all__ = ['MCALayer', 'MCAGate']class StdPool(nn.Module):def __init__(self):super(StdPool, self).__init__()def forward(self, x):b, c, _, _ = x.size()std = x.view(b, c, -1).std(dim=2, keepdim=True)std = std.reshape(b, c, 1, 1)return stdclass MCAGate(nn.Module):def __init__(self, k_size, pool_types=['avg', 'std']):super(MCAGate, self).__init__()self.pools = nn.ModuleList([])for pool_type in pool_types:if pool_type == 'avg':self.pools.append(nn.AdaptiveAvgPool2d(1))elif pool_type == 'max':self.pools.append(nn.AdaptiveMaxPool2d(1))elif pool_type == 'std':self.pools.append(StdPool())else:raise NotImplementedErrorself.conv = nn.Conv2d(1, 1, kernel_size=(1, k_size), stride=1, padding=(0, (k_size - 1) // 2), bias=False)self.sigmoid = nn.Sigmoid()self.weight = nn.Parameter(torch.rand(2))def forward(self, x):feats = [pool(x) for pool in self.pools]if len(feats) == 1:out = feats[0]elif len(feats) == 2:weight = torch.sigmoid(self.weight)out = 1 / 2 * (feats[0] + feats[1]) + weight[0] * feats[0] + weight[1] * feats[1]else:assert False, "Feature Extraction Exception!"out = out.permute(0, 3, 2, 1).contiguous()out = self.conv(out)out = out.permute(0, 3, 2, 1).contiguous()out = self.sigmoid(out)out = out.expand_as(x)return x * outclass MCALayer(nn.Module):def __init__(self, inp, no_spatial=False):super(MCALayer, self).__init__()lambd = 1.5gamma = 1temp = round(abs((math.log2(inp) - gamma) / lambd))kernel = temp if temp % 2 else temp - 1self.h_cw = MCAGate(3)self.w_hc = MCAGate(3)self.no_spatial = no_spatialif not no_spatial:self.c_hw = MCAGate(kernel)def forward(self, x):x_h = x.permute(0, 2, 1, 3).contiguous()x_h = self.h_cw(x_h)x_h = x_h.permute(0, 2, 1, 3).contiguous()x_w = x.permute(0, 3, 2, 1).contiguous()x_w = self.w_hc(x_w)x_w = x_w.permute(0, 3, 2, 1).contiguous()if not self.no_spatial:x_c = self.c_hw(x)x_out = 1 / 3 * (x_c + x_h + x_w)else:x_out = 1 / 2 * (x_h + x_w)return x_out
注册模块
在ultralytics/nn/tasks.py文件开头添加：
from ultralytics.nn.extra_modules import *
在ultralytics/nn/tasks.py文件中parse_model函数添加：
elif m in {MCALayer}:args = [ch[f], *args]
配置yaml文件
yolov8-mca.yaml
# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect# Parameters
nc: 80 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'# [depth, width, max_channels]n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPss: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPsm: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPsl: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPsx: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs# YOLOv8.0n backbone
backbone:# [from, repeats, module, args]- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4- [-1, 3, C2f, [128, True]]- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8- [-1, 6, C2f, [256, True]]- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16- [-1, 6, C2f, [512, True]]- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32- [-1, 3, C2f, [1024, True]]- [-1, 1, SPPF, [1024, 5]] # 9# YOLOv8.0n head
head:- [-1, 1, nn.Upsample, [None, 2, 'nearest']]- [[-1, 6], 1, Concat, [1]] # cat backbone P4- [-1, 3, C2f, [512]] # 12- [-1, 1, nn.Upsample, [None, 2, 'nearest']]- [[-1, 4], 1, Concat, [1]] # cat backbone P3- [-1, 3, C2f, [256]] # 15 (P3/8-small)- [-1, 1, MCALayer, []]- [-1, 1, Conv, [256, 3, 2]]- [[-1, 12], 1, Concat, [1]] # cat head P4- [-1, 3, C2f, [512]] # 19 (P4/16-medium)- [-1, 1, MCALayer, []]- [-1, 1, Conv, [512, 3, 2]]- [[-1, 9], 1, Concat, [1]] # cat head P5- [-1, 3, C2f, [1024]] # 23 (P5/32-large)- [-1, 1, MCALayer, []]- [[16, 20, 24], 1, Detect, [nc]] # Detect(P3, P4, P5)三、模型测试
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLOmodel = YOLO("yolov8n-mca.yaml") # build a new model from scratch
from n params module arguments0 -1 1 464 ultralytics.nn.modules.conv.Conv [3, 16, 3, 2]1 -1 1 4672 ultralytics.nn.modules.conv.Conv [16, 32, 3, 2]2 -1 1 7360 ultralytics.nn.modules.block.C2f [32, 32, 1, True]3 -1 1 18560 ultralytics.nn.modules.conv.Conv [32, 64, 3, 2]4 -1 2 49664 ultralytics.nn.modules.block.C2f [64, 64, 2, True]5 -1 1 73984 ultralytics.nn.modules.conv.Conv [64, 128, 3, 2]6 -1 2 197632 ultralytics.nn.modules.block.C2f [128, 128, 2, True]7 -1 1 295424 ultralytics.nn.modules.conv.Conv [128, 256, 3, 2]8 -1 1 460288 ultralytics.nn.modules.block.C2f [256, 256, 1, True]9 -1 1 164608 ultralytics.nn.modules.block.SPPF [256, 256, 5]10 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']11 [-1, 6] 1 0 ultralytics.nn.modules.conv.Concat [1]12 -1 1 148224 ultralytics.nn.modules.block.C2f [384, 128, 1]13 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']14 [-1, 4] 1 0 ultralytics.nn.modules.conv.Concat [1]15 -1 1 37248 ultralytics.nn.modules.block.C2f [192, 64, 1]16 -1 1 15 ultralytics.nn.extra_modules.attention.MCALayer[64]17 -1 1 36992 ultralytics.nn.modules.conv.Conv [64, 64, 3, 2]18 [-1, 12] 1 0 ultralytics.nn.modules.conv.Concat [1]19 -1 1 123648 ultralytics.nn.modules.block.C2f [192, 128, 1]20 -1 1 15 ultralytics.nn.extra_modules.attention.MCALayer[128]21 -1 1 147712 ultralytics.nn.modules.conv.Conv [128, 128, 3, 2]22 [-1, 9] 1 0 ultralytics.nn.modules.conv.Concat [1]23 -1 1 493056 ultralytics.nn.modules.block.C2f [384, 256, 1]24 -1 1 17 ultralytics.nn.extra_modules.attention.MCALayer[256]25 [16, 20, 24] 1 751507 ultralytics.nn.modules.head.Detect [1, [64, 128, 256]]
YOLOv8n-mca summary: 282 layers, 3011090 parameters, 3011074 gradients, 8.2 GFLOPs
四、模型训练
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO# Load a model
model = YOLO("yolov8n-mca.yaml") # build a new model from scratch# Use the model
model.train(data="./mydata/data.yaml",epochs=300,batch=32,imgsz=640,workers=8,device=0,project="runs/train",name='exp') # train the model
五、总结
- 模型的训练具有很大的随机性，您可能需要点运气和更多的训练次数才能达到最高的 mAP。
