【YOLO改进】主干插入CoordAttention模块(基于MMYOLO)

原创

已于 2024-04-09 09:20:46 修改 · 924 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO

于 2024-04-09 09:11:54 首次发布

CoordAttention注意力机制模块

论文地址：http://arxiv.org/abs/2103.02907

将CoordAttention模块添加到MMYOLO中

将开源代码CoordAttention.py文件复制到mmyolo/models/plugins目录下
导入MMYOLO用于注册模块的包: from mmyolo.registry import MODELS
确保class CoordAtt中的输入维度为in_channels(因为MMYOLO会提前传入输入维度参数，所以要保持参数名的一致)
利用@MODELS.register_module()将“class CoordAtt(nn.Module)”注册:
修改mmyolo/models/plugins/__init__.py文件
在终端运行:
```
python setup.py install
```
修改对应的配置文件，并且将plugins的参数“type”设置为“CoordAtt”，可参考【YOLO改进】主干插入注意力机制模块CBAM（基于MMYOLO）-优快云博客

修改后的CoordAttention.py

import torch
import torch.nn as nn
from mmyolo.registry import MODELS




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)

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

        mip = max(8, in_channels // reduction)

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

        self.conv_h = nn.Conv2d(mip, in_channels, kernel_size=1, stride=1, padding=0)
        self.conv_w = nn.Conv2d(mip, in_channels, 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