TPH-YOLOv5 | 基于Transformer的YOLOv5小目标检测器 | 四头加注意力

TPH-YOLOv5 | 基于Transformer的YOLOv5小目标检测器 | 四头加注意力_第1张图片
论文地址:https://arxiv.org/pdf/2108.11539.pdf
项目地址:https://github.com/cv516Buaa/tph-yolov5

在无人机捕获的场景中进行对象检测是最近的一项热门任务。由于无人机总是在不同的高度航行,物体尺度变化剧烈,给网络优化带来了负担。此外,高速和低空飞行会在密集的物体上带来运动模糊,这给物体识别带来了很大的挑战。为了解决上述两个问题,我们提出了 TPH-YOLOv5。在 YOLOv5 的基础上,我们增加了一个预测头来检测不同尺度的物体。然后我们用 Transformer Prediction Heads (TPH) 替换原来的预测头,以利用自注意力机制探索预测潜力。我们还集成了卷积块注意模型 (CBAM),以在具有密集对象的场景中找到注意区域。为了进一步改进我们提出的 TPH-YOLOv5,我们提供了许多有用的策略,例如数据增强、多尺度测试、多模型集成和利用额外的分类器。对数据集 VisDrone2021 的广泛实验表明,TPH-YOLOv5 具有良好的性能,在无人机捕获的场景中具有令人印象深刻的可解释性。在 DET-test-challenge 数据集上,TPH-YOLOv5 的 AP 结果为 39.18%,比之前的 SOTA 方法(DPNetV3)好 1.81%。在 VisDrone Challenge 2021 中,TPHYOLOv5 获得第 5 名,并与第 1 名模型(AP 39.43%)取得了良好的匹配结果。与基线模型(YOLOv5)相比,TPH-YOLOv5 提高了约 7%,令人鼓舞且具有竞争力

解决的问题

TPH-YOLOv5旨在解决无人机影像中存在的两个问题:

  • 因无人机在不同的高度飞行,物体的尺度变化剧烈。
  • 高速和低空飞行对排列密集的物体带来了运动模糊。

主要改进

TPH-YOLOv5是在YOLOv5的基础上做了下列改进:

  • 新增了一个检测头来检测更小尺度的物体。
  • 用transformer prediction heads(TPH)替换原来的预测头部。
  • 将CBAM集成到YOLOv5中,帮助网络在大区域覆盖的图像中找到感兴趣的区域。
  • 其它一系列小tricks。

TPH-YOLOv5网络结构如下

TPH-YOLOv5 | 基于Transformer的YOLOv5小目标检测器 | 四头加注意力_第2张图片

TPH模块

作者使用了一个Transformer Encoder来代替一些卷积和CSP结构,将Transformer在视觉中应用,也是目前的主流趋势,Transformer具有独特的注意力机制,效果比原先更好。
TPH-YOLOv5 | 基于Transformer的YOLOv5小目标检测器 | 四头加注意力_第3张图片

CBAM模块

TPH-YOLOv5 | 基于Transformer的YOLOv5小目标检测器 | 四头加注意力_第4张图片


我发现作者公布的代码和图中的代码不一样,所以自己按照上面的图复现了一个,除了检测头以外,完全按照原文内容,这里我们可以参考这篇文章的结构改进自己的模型。因为这些模块我们文件里已经有了,所以我们直接改配置文件就可以了。

# YOLOv5  by Ultralytics, GPL-3.0 license
# 迪菲赫尔曼 https://blog.csdn.net/weixin_43694096?spm=1000.2115.3001.5343

# Parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:
  - [19,27,  44,40,  38,94]  # P3/8
  - [96,68,  86,152,  180,137]  # P4/16
  - [140,301,  303,264,  238,542]  # P5/32
  - [436,615,  739,380,  925,792]  # P6/64

# YOLOv5 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],       # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],    # 1-P2/4
   [-1, 3, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],    # 3-P3/8
   [-1, 9, C3, [256]],
   [-1, 1, Conv, [512, 3, 2]],    # 5-P4/16
   [-1, 9, C3, [512]],
   [-1, 1, Conv, [768, 3, 2]],    # 7-P5/32
   [-1, 1, SPP, [1024, [3, 5, 7]]],
   [-1, 3, C3TR, [1024, False]],  # 9
  ]

# YOLOv5 head
head:
  [[-1, 1, Conv, [768, 1, 1]],  # 10
   [-1, 1, nn.Upsample, [None, 2, 'nearest']], #11
   [[-1, 6], 1, Concat, [1]],   # 12 cat backbone P5
   [-1, 3, C3, [768, False]],   # 13
   [-1, 1, CBAM, [768]],        # 14

   [-1, 1, Conv, [512, 1, 1]],  # 15
   [-1, 1, nn.Upsample, [None, 2, 'nearest']], #16
   [[-1, 4], 1, Concat, [1]],   # 17 cat backbone P4
   [-1, 3, C3, [512, False]],   #  18
   [-1, 1, CBAM, [512]],        # 19

   [-1, 1, Conv, [256, 1, 1]],  # 20
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],  #21
   [[-1, 2], 1, Concat, [1]],   # 22 cat backbone P3
   [-1, 3, C3TR, [256, False]], # 23 (P3/8-small)
   [-1, 1, CBAM, [256]],        # 24

   [-1, 1, Conv, [256, 3, 2]],  # 25
   [[-1, 20], 1, Concat, [1]],  # cat head P4 #26
   [-1, 3, C3TR, [512, False]], # 27 (P4/16-medium)
   [-1, 1, CBAM, [512]],        # 28

   [-1, 1, Conv, [512, 3, 2]],  # 29
   [[-1, 15], 1, Concat, [1]],  # 30 cat head P5
   [-1, 3, C3TR, [768, False]], # 31 (P5/32-large)
   [-1, 1, CBAM, [768]],        # 32

   [-1, 1, Conv, [768, 3, 2]],  # 33
   [[-1, 10], 1, Concat, [1]],  # 34  cat head P6
   [-1, 3, C3TR, [1024, False]],# 35 (P6/64-xlarge)

   [[23, 27, 31, 35], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5, P6)
  ]
模型 参数量parameters 计算量GFLOPs
TPH-YOLOv5 10009510 34.8

本人更多Yolov5实战内容导航

  1. 手把手带你调参Yolo v5 (v6.2)(推理)强烈推荐

  2. 手把手带你调参Yolo v5 (v6.2)(训练)

  3. 手把手带你调参Yolo v5 (v6.2)(验证)

  4. 如何快速使用自己的数据集训练Yolov5模型

  5. 手把手带你Yolov5 (v6.2)添加注意力机制(一)(并附上30多种顶会Attention原理图)强烈推荐

  6. 手把手带你Yolov5 (v6.2)添加注意力机制(二)(在C3模块中加入注意力机制)

  7. Yolov5如何更换激活函数?

  8. Yolov5如何更换BiFPN?

  9. Yolov5 (v6.2)数据增强方式解析

  10. Yolov5更换上采样方式( 最近邻 / 双线性 / 双立方 / 三线性 / 转置卷积)

  11. Yolov5如何更换EIOU / alpha IOU / SIoU?

  12. Yolov5更换主干网络之《旷视轻量化卷积神经网络ShuffleNetv2》

  13. YOLOv5应用轻量级通用上采样算子CARAFE

  14. 空间金字塔池化改进 SPP / SPPF / SimSPPF / ASPP / RFB / SPPCSPC / SPPFCSPC

  15. 用于低分辨率图像和小物体的模块SPD-Conv

  16. GSConv+Slim-neck 减轻模型的复杂度同时提升精度

  17. 头部解耦 | 将YOLOX解耦头添加到YOLOv5 | 涨点杀器


你可能感兴趣的:(YOLOv5改进实战,1024程序员节,论文,人工智能)