《博主简介》
小伙伴们好,我是阿旭。专注于人工智能、AIGC、python、计算机视觉相关分享研究。
✌更多学习资源,可关注公-仲-hao:【阿旭算法与机器学习】,共同学习交流~
👍感谢小伙伴们点赞、关注!
《------往期经典推荐------》
二、机器学习实战专栏【链接】,已更新31期,欢迎关注,持续更新中~~
三、深度学习【Pytorch】专栏【链接】
四、【Stable Diffusion绘画系列】专栏【链接】
《------正文------》
前言
文章提供了一个简单而有效的提出的双层路由注意力实现
,它利用稀疏性来节省计算和内存
,因此它在良好的性能和高计算效率方面都享有优势,特别是在密集预测任务中。本文详细介绍了如何在yolov8中添加BiLevelRoutingAttention
注意力机制模块,并且使用修改后的yolov8进行目标检测训练与推理
。本文提供了所有源码免费供小伙伴们学习参考,需要的可以通过文末方式自行下载。
本文改进使用的ultralytics版本为:ultralytics == 8.0.227
目录
1.BiLevelRoutingAttention简介
论文发表时间:2023.03.15
github地址:https://github.com/rayleizhu/BiFormer
paper地址:https://arxiv.org/pdf/2303.08810.pdf
摘要:作为视觉变换器的核心构建模块,
注意力是捕获远程依赖性的强大工具
。然而,这种强大功能是有代价的:由于需要计算所有空间位置之间的成对令牌交互,它会带来巨大的计算负担和沉重的内存占用。一系列工作试图通过引入手工制作的和内容不可知的稀疏性到注意力中来减轻这个问题
,例如限制注意力操作在局部窗口、轴向条纹或扩展窗口内部。与这些方法相反,我们提出了一种新颖的动态稀疏注意力通过双层路由来实现更灵活的计算分配与内容感知
。具体来说,对于一个查询,不相关的键值对首先在粗略的区域级别被过滤掉,然后在剩余候选区域(即路由区域)的并集中应用细粒度令牌到令牌的注意力。我们提供了一个简单而有效的提出的双层路由注意力实现
,它利用稀疏性来节省计算和内存
,同时只涉及GPU友好的密集矩阵乘法。基于提出的双层路由注意力构建的一个新的通用视觉变换器,名为BiFormer,随后被提出。由于BiFormer以一种查询自适应的方式只关注小部分相关令牌,不受其他不相关令牌的分心,因此它在良好的性能和高计算效率方面都享有优势,特别是在密集预测任务中。在图像分类、目标检测和语义分割等多个计算机视觉任务中的实证结果验证了我们设计的有效性。
论文主要亮点如下:
-
为原始注意力机制引入了一种新颖的双层次路由机制,这使得在查询自适应的方式下可以实现内容感知的稀疏模式。
-
将双层次路由注意力作为基本构件,我们提出了一个通用的视觉变换器命名为BiFormer。
-
在各种计算机视觉任务上的实验结果,包括图像分类、目标检测和语义分割,显示所提出的BiFormer在类似模型大小下,实现了比基线模型显著更好的性能。
1.1 网络结构
1.2 稀疏性应用方式
1.3 性能对比
2.YOLOv8添加注意力机制
替换位置与替换后网络结构示意
添加位置
替换后的YOLOv8网络结构如下:
定义注意力机制类
在ultralytics/nn/modules/block.py
中添加如下代码块,并定义BiLevelRoutingAttention
类:
并在ultralytics/nn/modules/block.py
中最上方添加如下代码:
修改指定文件
在ultralytics/nn/modules/__init__.py
文件中的添加如下代码:
在 ultralytics/nn/tasks.py
上方导入BiLevelRoutingAttention
类名,并在parse_model
解析函数中添加如下代码:
elif m in [BiLevelRoutingAttention]:
c2 = ch[f]
args = [c2, *args[0:]]
在ultralytics/cfg/models/v8
文件夹下新建yolov8-BiLevelRoutingAttention.yaml
文件,内容如下:
# Parameters
nc: 80 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
# [depth, width, max_channels]
n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPs
s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPs
m: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPs
l: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
x: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPs
# YOLOv8.0n 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, 3, C2f, [128, True]]
- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
- [-1, 6, C2f, [256, True]]
- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
- [-1, 6, C2f, [512, True]]
- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
- [-1, 3, C2f, [1024, True]]
- [-1, 1, SPPF, [1024, 5]] # 9
# YOLOv8.0n head
head:
- [-1, 1, nn.Upsample, [None, 2, 'nearest']]
- [[-1, 6], 1, Concat, [1]] # cat backbone P4
- [-1, 3, C2f, [512]] # 12
- [-1, 1, nn.Upsample, [None, 2, 'nearest']]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # 15 (P3/8-small)
- [-1, 1, BiLevelRoutingAttention, []] # 16
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 12], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 19 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 9], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # 22 (P5/32-large)
- [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)
3.加载配置文件并训练
加载yolov8-BiLevelRoutingAttention.yaml
配置文件,并运行train.py
训练代码:
#coding:utf-8
from ultralytics import YOLO
if __name__ == '__main__':
model = YOLO('ultralytics/cfg/models/v8/yolov8-BiLevelRoutingAttention.yaml')
model.load('yolov8n.pt') # loading pretrain weights
model.train(data='datasets/TomatoData/data.yaml', epochs=50, batch=4)
注意观察,打印出的网络结构是否正常修改,如下图所示:
4.模型推理
模型训练完成后,我们使用训练好的模型对图片进行检测:
#coding:utf-8
from ultralytics import YOLO
import cv2
# 所需加载的模型目录
# path = 'models/best2.pt'
path = 'runs/detect/train11/weights/best.pt'
# 需要检测的图片地址
img_path = "TestFiles/Riped tomato_8.jpeg"
# 加载预训练模型
# conf 0.25 object confidence threshold for detection
# iou 0.7 intersection over union (IoU) threshold for NMS
model = YOLO(path, task='detect')
# 检测图片
results = model(img_path)
res = results[0].plot()
# res = cv2.resize(res,dsize=None,fx=2,fy=2,interpolation=cv2.INTER_LINEAR)
cv2.imshow("YOLOv8 Detection", res)
cv2.waitKey(0)
【源码免费获取】
为了小伙伴们能够,更好的学习实践,本文已将所有代码、示例数据集、论文等相关内容打包上传,供小伙伴们学习。获取方式如下:
结束语
关于本篇文章大家有任何建议或意见,欢迎在评论区留言交流!
觉得不错的小伙伴,感谢点赞、关注加收藏哦!