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

1、在yolov5/models下面新建一个CoTAttention.py文件,在里面放入下面的代码

代码如下:

bash 复制代码
import numpy as np
import torch
from torch import flatten, nn
from torch.nn import init
from torch.nn.modules.activation import ReLU
from torch.nn.modules.batchnorm import BatchNorm2d
from torch.nn import functional as F


class CoTAttention(nn.Module):

    def __init__(self, dim=512, kernel_size=3):
        super().__init__()
        self.dim = dim
        self.kernel_size = kernel_size

        self.key_embed = nn.Sequential(
            nn.Conv2d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=4, bias=False),
            nn.BatchNorm2d(dim),
            nn.ReLU()
        )
        self.value_embed = nn.Sequential(
            nn.Conv2d(dim, dim, 1, bias=False),
            nn.BatchNorm2d(dim)
        )

        factor = 4
        self.attention_embed = nn.Sequential(
            nn.Conv2d(2 * dim, 2 * dim // factor, 1, bias=False),
            nn.BatchNorm2d(2 * dim // factor),
            nn.ReLU(),
            nn.Conv2d(2 * dim // factor, kernel_size * kernel_size * dim, 1)
        )

    def forward(self, x):
        bs, c, h, w = x.shape
        k1 = self.key_embed(x)  # bs,c,h,w
        v = self.value_embed(x).view(bs, c, -1)  # bs,c,h,w

        y = torch.cat([k1, x], dim=1)  # bs,2c,h,w
        att = self.attention_embed(y)  # bs,c*k*k,h,w
        att = att.reshape(bs, c, self.kernel_size * self.kernel_size, h, w)
        att = att.mean(2, keepdim=False).view(bs, c, -1)  # bs,c,h*w
        k2 = F.softmax(att, dim=-1) * v
        k2 = k2.view(bs, c, h, w)

        return k1 + k2

2、找到yolo.py文件,进行更改内容

在29行加一个from models.CoTAttention import CoTAttention, 保存即可

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

4、在yolo.py里面加入两行代码(335-337)

保存即可!

运行

相关推荐
Francek Chen1 小时前
【深度学习计算机视觉】07:单发多框检测(SSD)
人工智能·pytorch·深度学习·计算机视觉·单发多框检测
dami_king1 小时前
RTX4090算力应用-3D
人工智能·深度学习·3d·ai
编程武士8 小时前
从50ms到30ms:YOLOv10部署中图像预处理的性能优化实践
人工智能·python·yolo·性能优化
mCell9 小时前
长期以来我对 LLM 的误解
深度学习·llm·ollama
Ada's10 小时前
深度学习在自动驾驶上应用(二)
人工智能·深度学习·自动驾驶
机器学习之心11 小时前
198种组合算法+优化BiLSTM神经网络+SHAP分析+新数据预测+多输出!深度学习可解释分析,强烈安利,粉丝必备!
深度学习·神经网络·shap分析·新数据预测·198种组合算法·优化bilstm神经网络·多输出
charieli-fh12 小时前
指令微调数据评估与影响:构建高质量大语言模型的关键
人工智能·深度学习·语言模型
Coovally AI模型快速验证12 小时前
从避障到实时建图:机器学习如何让无人机更智能、更安全、更实用(附微型机载演示示例)
人工智能·深度学习·神经网络·学习·安全·机器学习·无人机
没有梦想的咸鱼185-1037-166312 小时前
【遥感技术】从CNN到Transformer:基于PyTorch的遥感影像、无人机影像的地物分类、目标检测、语义分割和点云分类
pytorch·python·深度学习·机器学习·数据分析·cnn·transformer
IT古董13 小时前
【第五章:计算机视觉-项目实战之图像分割实战】1.图像分割理论-(1)图像分割基础知识:定义、任务描述、应用场景、标注格式
yolo·目标检测·计算机视觉