yolov8添加ca注意力机制

创建文件 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
yolov8添加ca注意力机制_第1张图片

init.py中添加头文件

位置 :ultralytics/nn/modules/init.py
yolov8添加ca注意力机制_第2张图片

tasks.py 文件

位置:ultralytics/nn/tasks.py

task.py文件 添加头文件

yolov8添加ca注意力机制_第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:]]

yolov8添加ca注意力机制_第4张图片

运行效果

yolov8添加ca注意力机制_第5张图片

对比图(左:未添加cbam,右上:添加cbam,右下:添加ca)

yolov8添加ca注意力机制-出现bug

ImportError: cannot import name ‘CoordAtt’ from ‘ultralytics.nn.modules’ (D:\anaconda3\envs\torch\lib\site-packages\ultralytics\nn\modules_init_.py)

在这里插入图片描述

解决方法:拷贝项目中左图文件,到环境配置的右图目录中

yolov8添加ca注意力机制_第6张图片

ImportError: cannot import name ‘CoordAtt’ from ‘ultralytics.nn.modules.conv’ (D:\anaconda3\envs\torch\lib\site-packages\ultralytics\nn\modules\conv.py)

在这里插入图片描述

解决方法:拷贝项目中左图文件,到环境配置的右图目录中

yolov8添加ca注意力机制_第7张图片

ModuleNotFoundError: No module named ‘ultralytics.nn.modules.coordAtt’

在这里插入图片描述
解决方法:拷贝项目中左图文件,到环境配置的右图目录中
yolov8添加ca注意力机制_第8张图片

你可能感兴趣的:(YOLO,深度学习,pytorch)