yolov8添加ca注意力机制

登高去梯

已于 2024-04-07 09:55:00 修改

阅读量1.9k

点赞数 9

分类专栏： yolov8 文章标签： YOLO 深度学习 pytorch

于 2023-12-05 18:35:08 首次发布

本文链接：https://blog.csdn.net/weixin_44240141/article/details/134813886

版权

yolov8 专栏收录该内容

3 篇文章 0 订阅

订阅专栏

创建文件 coordAtt.py

位置：ultralytics/nn/modules/coordAtt.py

######################  CoordAtt  ####     start   by  AI&CV  ###############################
# https://zhuanlan.zhihu.com/p/655475515
import torch
import torch.nn as nn
import torch.nn.functional as F


class h_sigmoid(nn.Module):
    def __init__(self, inplace=True):
        super(h_sigmoid, self).__init__()
        self.relu = nn.ReLU6(inplace=inplace)

    def forward(self, x):
        return self.relu(x + 3) / 6


class h_swish(nn.Module):
    def __init__(self, inplace=True):
        super(h_swish, self).__init__()
        self.sigmoid = h_sigmoid(inplace=inplace)

    def forward(self, x):
        return x * self.sigmoid(x)


class CoordAtt(nn.Module):
    def __init__(self, inp, reduction=32):
        super(CoordAtt, self).__init__()
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))

        mip = max(8, inp // reduction)

        self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)
        self.bn1 = nn.BatchNorm2d(mip)
        self.act = h_swish()

        self.conv_h = nn.Conv2d(mip, inp, kernel_size=1, stride=1, padding=0)
        self.conv_w = nn.Conv2d(mip, inp, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        identity = x

        n, c, h, w = x.size()
        x_h = self.pool_h(x)
        x_w = self.pool_w(x).permute(0, 1, 3, 2)

        y = torch.cat([x_h, x_w], dim=2)
        y = self.conv1(y)
        y = self.bn1(y)
        y = self.act(y)

        x_h, x_w = torch.split(y, [h, w], dim=2)
        x_w = x_w.permute(0, 1, 3, 2)

        a_h = self.conv_h(x_h).sigmoid()
        a_w = self.conv_w(x_w).sigmoid()

        out = identity * a_w * a_h

        return out
######################  CoordAtt  ####     end   by  AI&CV  ###############################

conv.py中添加头文件

位置：ultralytics/nn/modules/conv.py
在这里插入图片描述

from ultralytics.nn.modules.coordAtt import CoordAtt # todo 源码修改 ~1+
'CoordAtt') # todo 源码修改（还原则删除“,'CoordAtt'”） ~1

init.py中添加头文件

位置：ultralytics/nn/modules/init.py
在这里插入图片描述

CoordAtt) # todo 源码修改（还原则删除“,CoordAtt”） ~2+
'CoordAtt') # todo 源码修改（还原则删除“,'CoordAtt'”） ~2

tasks.py 文件

位置：ultralytics/nn/tasks.py

task.py文件添加头文件

在这里插入图片描述

CoordAtt)  # todo 源码修改1 (还原则删除", CoordAtt") ~3

CBAM,CoordAtt)  # todo 源码修改1 (还原则删除", CoordAtt") ~3

task.py文件的方法中添加代码

        elif m is CoordAtt: # todo 源码修改 ~4
            """
            ch[f]:上一层的
            args[0]:第0个参数
            c1:输入通道数
            c2:输出通道数
            """
            c1, c2 = ch[f], args[0]
            # print("ch[f]:",ch[f])
            # print("args[0]:",args[0])
            # print("args:",args)
            # print("c1:",c1)
            # print("c2:",c2)
            if c2 != nc:  # if c2 not equal to number of classes (i.e. for Classify() output)
                c2 = make_divisible(c2 * width, 8)
            args = [c1, *args[1:]]