1.BiLevelRoutingAttention简介
论文发表时间:2023.03.15
摘要:作为视觉变换器的核心构建模块,
注意力是捕获远程依赖性的强大工具
。然而,这种强大功能是有代价的:由于需要计算所有空间位置之间的成对令牌交互,它会带来巨大的计算负担和沉重的内存占用。一系列工作试图通过引入手工制作的和内容不可知的稀疏性到注意力中来减轻这个问题
,例如限制注意力操作在局部窗口、轴向条纹或扩展窗口内部。与这些方法相反,我们提出了一种新颖的动态稀疏注意力通过双层路由来实现更灵活的计算分配与内容感知
。具体来说,对于一个查询,不相关的键值对首先在粗略的区域级别被过滤掉,然后在剩余候选区域(即路由区域)的并集中应用细粒度令牌到令牌的注意力。我们提供了一个简单而有效的提出的双层路由注意力实现
,它利用稀疏性来节省计算和内存
,同时只涉及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)