ultralytics-yolov8-加入attention(以eca注意力机制为例)
·
一、我的版本
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]
下图就是大概的位置
更多推荐
所有评论(0)