【YOLO改进】主干SE模块(基于MMYOLO)

SE模块

论文链接:https://arxiv.org/abs/1709.01507

将SE模块添加到MMYOLO中

1. 将开源代码SE.py文件复制到mmyolo/models/plugins目录下

2. 导入MMYOLO用于注册模块的包: from mmyolo.registry import MODELS

3. 确保 class SEAttention中的输入维度为in_channels(因为MMYOLO会提前传入输入维度参数，所以要保持参数名的一致)

4. 利用@MODELS.register_module()将“class SEAttention(nn.Module)”注册:

5. 修改mmyolo/models/plugins/__init__.py文件

6. 在终端运行:

   python setup.py install

7. 修改对应的配置文件，并且将plugins的参数“type”设置为“SEAttention”，可参考【YOLO改进】主干插入注意力机制模块CBAM（基于MMYOLO）-优快云博客

修改后的SE.py

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

@MODELS.register_module()
class SEAttention(nn.Module):

    def __init__(self, in_channels=512,reduction=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(in_channels // reduction, in_channels, bias=False),
            nn.Sigmoid()
        )


    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)


if __name__ == '__main__':
    input=torch.randn(50,512,7,7)
    se = SEAttention(channel=512,reduction=8)
    output=se(input)
    print(output.shape)

修改后的__init__.py

# Copyright (c) OpenMMLab. All rights reserved.
from .cbam import CBAM
from .Biformer import BiLevelRoutingAttention
from .A2Attention import DoubleAttention
from .CoordAttention import CoordAtt
from .CoTAttention