YOLO11改进|注意力机制篇|引入CoT注意力机制
目录
- 一、【CoT】注意力机制
- 1.1【CoT】注意力介绍
- 1.2【CoT】核心代码
- 二、添加【CoT】注意力机制
- 2.1STEP1
- 2.2STEP2
- 2.3STEP3
- 2.4STEP4
- 三、yaml文件与运行
- 3.1yaml文件
- 3.2运行成功截图
一、【CoT】注意力机制
1.1【CoT】注意力介绍
下图是【CoT】的结构图,让我们简单分析一下运行过程和优势
处理过程:
- 提取查询和键:
- 输入图像首先通过两个并行的卷积操作生成 查询(Query) 和 键(Key) 特征。其中,查询是通过一个 1×1的卷积得到的,而键是通过一个 3×3的卷积得到的静态特征。这两个特征代表了不同维度上的信息,查询用于定位,键用于描述局部上下文。
- 特征匹配与注意力权重:
- 查询和键特征经过两个 1×1卷积操作,然后通过注意力机制进行特征匹配,生成动态的 注意力权重(weights)。权重的生成基于查询和键的相似性,确定了图像中哪些区域应该被关注。
- 值特征的生成与加权:
- 值(Values) 特征同样通过 1×1的卷积操作生成,这些特征包含了输入图像中的原始信息。生成的注意力权重被应用于值特征,通过动态加权将权重信息融入到值特征中,生成最终的动态特征表示。
- 输出:
- 最后,将动态加权的特征进行汇总,得到最终的输出特征。这些输出特征经过注意力机制的加持,能够更有效地表示输入图像中的重要区域。
优势: - 动态加权的自适应性:
- 通过查询和键的特征匹配,注意力机制能够为输入特征分配不同的权重。这种自适应的加权方式确保了模型能够更加灵活地聚焦于图像中最重要的部分,增强了特征表达的准确性。
- 局部和全局信息的结合:
- 通过 1×1卷积和 3×3卷积的结合,模型能够同时捕捉到图像的局部细节和全局上下文信息。查询特征定位局部重要信息,而键特征提供全局的上下文背景,二者的结合使得模型在处理复杂图像时更加高效。
- 高效的计算:
- 使用卷积操作提取查询、键和值特征,确保了模型能够高效处理大规模图像数据。相较于全局注意力机制,该模块通过局部操作降低了计算成本,同时保留了强大的特征表达能力。
- 提升特征表达的多样性:
- 通过引入注意力权重的动态加权,该模块能够灵活调整特征的表达,尤其在处理复杂场景时,能够有效提高模型对不同对象的区分能力。
1.2【CoT】核心代码
import torch
from torch import flatten, nn
from torch.nn import functional as Fclass CoTAttention(nn.Module):def __init__(self, dim=512, kernel_size=3):super().__init__()self.dim = dimself.kernel_size = kernel_sizeself.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 = 4self.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.shapek1 = self.key_embed(x) # bs,c,h,wv = self.value_embed(x).view(bs, c, -1) # bs,c,h,wy = torch.cat([k1, x], dim=1) # bs,2c,h,watt = self.attention_embed(y) # bs,c*k*k,h,watt = 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*wk2 = F.softmax(att, dim=-1) * vk2 = k2.view(bs, c, h, w)return k1 + k2
二、添加【CoT】注意力机制
2.1STEP1
首先找到ultralytics/nn文件路径下新建一个Add-module的python文件包【这里注意一定是python文件包,新建后会自动生成_init_.py】,如果已经跟着我的教程建立过一次了可以省略此步骤,随后新建一个CoTpy文件并将上文中提到的注意力机制的代码全部粘贴到此文件中,如下图所示
2.2STEP2
在STEP1中新建的_init_.py文件中导入增加改进模块的代码包如下图所示
2.3STEP3
找到ultralytics/nn文件夹中的task.py文件,在其中按照下图添加
2.4STEP4
定位到ultralytics/nn文件夹中的task.py文件中的def parse_model(d, ch, verbose=True): # model_dict, input_channels(3)函数添加如图代码,【如果不好定位可以直接ctrl+f搜索定位】
三、yaml文件与运行
3.1yaml文件
以下是添加【CoT】注意力机制在Backbone中的yaml文件,大家可以注释自行调节,效果以自己的数据集结果为准
# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLO11 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect# Parameters
nc: 80 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n'# [depth, width, max_channels]n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPss: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPsm: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPsl: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPsx: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs# YOLO11n backbone
backbone:# [from, repeats, module, args]- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2- [-1, 1, Conv, [128,3,2]] # 1-P2/4- [-1, 2, C3k2, [256, False, 0.25]]- [-1, 1, Conv, [256,3,2]] # 3-P3/8- [-1, 2, C3k2, [512, False, 0.25]]- [-1, 1, Conv, [512,3,2]] # 5-P4/16- [-1, 2, C3k2, [512, True]]- [-1, 1, Conv, [1024,3,2]] # 7-P5/32- [-1, 2, C3k2, [1024, True]]- [-1, 1, SPPF, [1024, 5]] # 9- [-1,1,CoTAttention,[]]- [-1, 2, C2PSA, [1024]] # 10# YOLO11n head
head:- [-1, 1, nn.Upsample, [None, 2, "nearest"]]- [[-1, 6], 1, Concat, [1]] # cat backbone P4- [-1, 2, C3k2, [512, False]] # 13- [-1, 1, nn.Upsample, [None, 2, "nearest"]]- [[-1, 4], 1, Concat, [1]] # cat backbone P3- [-1, 2, C3k2, [256, False]] # 16 (P3/8-small)- [-1, 1, Conv, [256, 3, 2]]- [[-1, 14], 1, Concat, [1]] # cat head P4- [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium)- [-1, 1, Conv, [512, 3, 2]]- [[-1, 11], 1, Concat, [1]] # cat head P5- [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large)- [[17, 20, 23], 1, Detect, [nc]] # Detect(P3, P4, P5)
以上添加位置仅供参考,具体添加位置以及模块效果以自己的数据集结果为准
3.2运行成功截图
OK 以上就是添加【CoT】注意力机制的全部过程了,后续将持续更新尽情期待