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

1、在yolov5/models下面新建一个EfficientChannelAttention.py文件,在里面放入下面的代码
(超详细)6-YOLOV5改进-添加ECA注意力机制_第1张图片
代码如下:

import torch, math
from torch import nn

class EfficientChannelAttention(nn.Module):           # Efficient Channel Attention module
    def __init__(self, c, b=1, gamma=2):
        super(EfficientChannelAttention, self).__init__()
        t = int(abs((math.log(c, 2) + b) / gamma))
        k = t if t % 2 else t + 1

        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv1 = nn.Conv1d(1, 1, kernel_size=k, padding=int(k/2), bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        out = self.avg_pool(x)
        out = self.conv1(out.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
        out = self.sigmoid(out)
        return out * x

2、找到yolo.py文件,进行更改内容
在29行加一个from models.EfficientChannelAttention import EfficientChannelAttention, 保存即可
(超详细)6-YOLOV5改进-添加ECA注意力机制_第2张图片

3、找到自己想要更改的yaml文件,我选择的yolov5s.yaml文件(你可以根据自己需求进行选择),将刚刚写好的模块EfficientChannelAttention加入到yolov5s.yaml里面,并更改一些内容。更改如下
(超详细)6-YOLOV5改进-添加ECA注意力机制_第3张图片

4、在yolo.py里面加入两行代码(335-337)
保存即可!
(超详细)6-YOLOV5改进-添加ECA注意力机制_第4张图片
运行一下
(超详细)6-YOLOV5改进-添加ECA注意力机制_第5张图片

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