yolov8改进之EMA注意力机制

最新推荐文章于 2024-07-24 09:26:27 发布

lvnacp

最新推荐文章于 2024-07-24 09:26:27 发布

阅读量2.3k

点赞数 7

分类专栏： yolov8改进文章标签： YOLO 深度学习人工智能

本文链接：https://blog.csdn.net/lvnacp/article/details/139000082

版权

yolov8改进专栏收录该内容

4 篇文章 2 订阅

订阅专栏

论文地址:https://arxiv.org/abs/2305.13563v1

模型讲解和代码实现，即插即用提升性能

概述

通过通道降维来建模跨通道关系可能会给提取深度视觉带来副作用。因此提出了一种新的高效的多尺度注意力(EMA)模块。以保留每个通道上的信息和降低计算开销为目标，将部分通道重塑为批量维度，并将通道维度分组为多个子特征，使空间语义特征在每个特征组中均匀分布。
在这里插入图片描述
左图为CA注意力机制模型图，右图为EMA注意力机制模型图

模型优点

(1) 我们考虑一种通用方法，将部分通道维度重塑为批量维度，以避免通过通用卷积进行某种形式的降维。
(2) 除了在不进行通道降维的情况下在每个并行子网络中构建局部的跨通道交互外，我们还通过跨空间学习方法融合两个并行子网络的输出特征图。
(3) 与CBAM、NAM[16]、SA、ECA和CA相比，EMA不仅取得了更好的结果，而且在所需参数方面效率更高。

代码实现

在ultralytics-8.2.0\ultralytics\nn\文件夹下新建attention.py

import torch
from torch import nn


class EMA(nn.Module):
    def __init__(self, channels, c2=None, factor=32):
        super(EMA, self).__init__()
        self.groups = factor
        assert channels // self.groups > 0
        self.softmax = nn.Softmax(-1)
        self.agp = nn.AdaptiveAvgPool2d((1, 1))
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))
        self.gn = nn.GroupNorm(channels // self.groups, channels // self.groups)
        self.conv1x1 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=1, stride=1, padding=0)
        self.conv3x3 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        b, c, h, w = x.size()
        group_x = x.reshape(b * self.groups, -1, h, w)  # b*g,c//g,h,w
        x_h = self.pool_h(group_x)
        x_w = self.pool_w(group_x).permute(0, 1, 3, 2)
        hw = self.conv1x1(torch.cat([x_h, x_w], dim=2))
        x_h, x_w = torch.split(hw, [h, w], dim=2)
        x1 = self.gn(group_x * x_h.sigmoid() * x_w.permute(0, 1, 3, 2).sigmoid())
        x2 = self.conv3x3(group_x)
        x11 = self.softmax(self.agp(x1).reshape(b * self.groups, -1, 1).permute(0, 2, 1))
        x12 = x2.reshape(b * self.groups, c // self.groups, -1)  # b*g, c//g, hw
        x21 = self.softmax(self.agp(x2).reshape(b * self.groups, -1, 1).permute(0, 2, 1))
        x22 = x1.reshape(b * self.groups, c // self.groups, -1)  # b*g, c//g, hw
        weights = (torch.matmul(x11, x12) + torch.matmul(x21, x22)).reshape(b * self.groups, 1, h, w)
        return (group_x * weights.sigmoid()).reshape(b, c, h, w)

在ultralytics-8.2.0\ultralytics\nn\tasks.py文件的parse_model()函数中添加如下代码
在这里插入图片描述
在ultralytics-8.2.0\ultralytics\cfg\models\v8目录下新建yolov8_EMA.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 GFLOPs
#  s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients,  28.8 GFLOPs
  m: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients,  79.3 GFLOPs
#  l: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
#  x: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs

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, EMA, [8]] #13

  - [-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, EMA, [8]] #16

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 13], 1, Concat, [1]] # cat head P4
  - [-1, 3, C2f, [512]] # 19 (P4/16-medium)
  - [-1, 1, EMA, [8]] # 20

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 9], 1, Concat, [1]] # cat head P5
  - [-1, 3, C2f, [1024]] # 23 (P5/32-large)
  - [-1, 1, EMA, [8]] #24

  - [[17, 21, 25], 1, Detect, [nc]] # Detect(P3, P4, P5)

在主目录下新建main.py文件
在这里插入图片描述
我使用yolov8的yolov8m加载权重，main.py代码如下

from ultralytics import YOLO
def train_model():
    model = YOLO("yolov8_EMA.yaml").load('yolov8m.pt')
    model.train(data="rail_defects.yaml", epochs=300)

if __name__ == '__main__':
    train_model()

运行结果

                 from  n    params  module                                       arguments                     
  0                  -1  1      1392  ultralytics.nn.modules.conv.Conv             [3, 48, 3, 2]                 
  1                  -1  1     41664  ultralytics.nn.modules.conv.Conv             [48, 96, 3, 2]                
  2                  -1  2    111360  ultralytics.nn.modules.block.C2f             [96, 96, 2, True]             
  3                  -1  1    166272  ultralytics.nn.modules.conv.Conv             [96, 192, 3, 2]               
  4                  -1  4    813312  ultralytics.nn.modules.block.C2f             [192, 192, 4, True]           
  5                  -1  1    664320  ultralytics.nn.modules.conv.Conv             [192, 384, 3, 2]              
  6                  -1  4   3248640  ultralytics.nn.modules.block.C2f             [384, 384, 4, True]           
  7                  -1  1   1991808  ultralytics.nn.modules.conv.Conv             [384, 576, 3, 2]              
  8                  -1  2   3985920  ultralytics.nn.modules.block.C2f             [576, 576, 2, True]           
  9                  -1  1    831168  ultralytics.nn.modules.block.SPPF            [576, 576, 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  2   1993728  ultralytics.nn.modules.block.C2f             [960, 384, 2]                 
 13                  -1  1      1488  ultralytics.nn.attention.EMA                 [384, 8]                      
 14                  -1  1         0  torch.nn.modules.upsampling.Upsample         [None, 2, 'nearest']          
 15             [-1, 4]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 16                  -1  2    517632  ultralytics.nn.modules.block.C2f             [576, 192, 2]                 
 17                  -1  1       384  ultralytics.nn.attention.EMA                 [192, 8]                      
 18                  -1  1    332160  ultralytics.nn.modules.conv.Conv             [192, 192, 3, 2]              
 19            [-1, 13]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 20                  -1  2   1846272  ultralytics.nn.modules.block.C2f             [576, 384, 2]                 
 21                  -1  1      1488  ultralytics.nn.attention.EMA                 [384, 8]                      
 22                  -1  1   1327872  ultralytics.nn.modules.conv.Conv             [384, 384, 3, 2]              
 23             [-1, 9]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 24                  -1  2   4207104  ultralytics.nn.modules.block.C2f             [960, 576, 2]                 
 25                  -1  1      3312  ultralytics.nn.attention.EMA                 [576, 8]                      
 26        [17, 21, 25]  1   3776854  ultralytics.nn.modules.head.Detect           [2, [192, 384, 576]]          
YOLOv8_EMA summary: 327 layers, 25864150 parameters, 25864134 gradients, 79.6 GFLOPs