一、我的版本

ultralytics=8.3.112

python=3.9

二、新建一个attention.py

在这个路径下ultralytics-8.3.112/ultralytics/nn/Attention.py新建一个Attention.py

然后把eca注意力机制的代码复制在里面

import torch
from torch import nn
from torch.nn.parameter import Parameter


class ECA(nn.Module):
    """Constructs a ECA module.
    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """

    def __init__(self, channel, k_size=3):
        super(ECA, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # feature descriptor on the global spatial information
        y = self.avg_pool(x)

        # Two different branches of ECA module
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)

        # Multi-scale information fusion
        y = self.sigmoid(y)

        return x * y.expand_as(x)

三、在task.py里面有两个操作

1、ultralytics-8.3.112/ultralytics/nn/tasks.py点开这个文件,在顶部加入

from ultralytics.nn.Attention import (ECA)

2、往下拉,在

def parse_model(d, ch, verbose=True):  # model_dict, input_channels(3)

这个里面加入即可运行

        elif m in {ECA}:
            c2 = ch[f]
            args = [c2, *args]

下图就是大概的位置

Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐