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

EffectiveSEModule模块

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

将EffectiveSEModule模块添加到MMYOLO中

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

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

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

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

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

6. 在终端运行:

   python setup.py install

7. 安装对应版本的timm

   pip install timm==0.6.13
   

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

修改后EffectiveSE.py

import torch
from timm.models.layers.create_act import create_act_layer
from torch import nn as nn
from mmyolo.registry import MODELS

@MODELS.register_module()
class EffectiveSEModule(nn.Module):
    def __init__(self, in_channels, add_maxpool=False, gate_layer='hard_sigmoid'):
        super(EffectiveSEModule, self).__init__()
        self.add_maxpool = add_maxpool
        self.fc = nn.Conv2d(in_channels, in_channels, kernel_size=1, padding=0)
        self.gate = create_act_layer(gate_layer)

    def forward(self, x):
        x_se = x.mean((2, 3), keepdim=True)
        if self.add_maxpool:
            # experimental codepath, may remove or change
            x_se = 0.5 * x_se + 0.5 * x.amax((2, 3), keepdim=True)
        x_se = self.fc(x_se)
        return x * self.gate(x_se)

if __name__ == '__main__':
    input=torch.randn(50,512,7,7)
    Ese = EffectiveSEModule(512)
    output=Ese(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 import CoTAttention
from .ECA import ECAAttention
from .EffectiveSE import EffectiveSEModule
__all__ = ['CBAM', 'BiLevelRoutingAttention', 'DoubleAtten