(超详细)9-YOLOV5改进-添加EffectiveSEModule注意力机制

发布时间:2024年01月21日

1、在yolov5/models下面新建一个EffectiveSEModule.py文件,在里面放入下面的代码
在这里插入图片描述

代码如下:

import torch
from torch import nn as nn
from timm.models.layers.create_act import create_act_layer


class EffectiveSEModule(nn.Module):
    def __init__(self, channels, add_maxpool=False, gate_layer='hard_sigmoid'):
        super(EffectiveSEModule, self).__init__()
        self.add_maxpool = add_maxpool
        self.fc = nn.Conv2d(channels, 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)

2、找到yolo.py文件,进行更改内容
在29行加一个from models.EffectiveSEModule import EffectiveSEModule, 保存即可
在这里插入图片描述

3、找到自己想要更改的yaml文件,我选择的yolov5s.yaml文件(你可以根据自己需求进行选择),将刚刚写好的模块EffectiveSEModule加入到yolov5s.yaml里面,并更改一些内容。更改如下
在这里插入图片描述

4、在yolo.py里面加入两行代码(335-337)
保存即可!
在这里插入图片描述

文章来源:https://blog.csdn.net/qq_44421796/article/details/135726604
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。