devbox
小蓝xlanll
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

YOLOv5-Lite 详解教程 | 嚼碎所有原理、训练自己数据集、TensorRT部署落地应有尽有...

作者：小蓝xlanll | 2024-02-25 10:35:20

赞

踩

yolov5-lite检测的图像最小尺寸

点击上方“3D视觉工坊”，选择“星标”

干货第一时间送达

04b6d0d464e11733e9e1402913473774.png

作者丨ChaucerG

来源丨集智书童

53325154b85b82b3fa720d3b093d8e5e.png

YOLOv5 Lite在YOLOv5的基础上进行一系列消融实验，使其更轻（Flops更小，内存占用更低，参数更少），更快（加入shuffle channel，yolov5 head进行通道裁剪，在320的input_size至少能在树莓派4B上的推理速度可以达到10+FPS），更易部署（摘除Focus层和4次slice操作，让模型量化精度下降在可接受范围内）。

912ee073b23c3031fb5dfbc049392a1f.png

1输入端方法

1、Mosaic数据增强

YOLOv5 Lite的输入端采用了和YOLOv5、YOLOv4一样的Mosaic数据增强的方式。其实Mosaic数据增强的作者也是来自YOLOv5团队的成员，不过，随机缩放、随机裁剪、随机排布的方式进行拼接，对于小目标的检测效果还是很不错的。

3b83438777de5bbac214f7ff26e11334.png

为什么要进行Mosaic数据增强呢？

在平时训练模型时，一般来说小目标的AP比中目标和大目标低很多。而Coco数据集中也包含大量的小目标，但比较麻烦的是小目标的分布并不均匀。

首先看下小、中、大目标的定义：

57cbc582ee8c367a24f2e7ec03cb3e70.png

可以看到小目标的定义是目标框的长宽0×0~32×32之间的物体。

d56419e0623059f3e062132844104e12.png

但在整体的数据集中，小、中、大目标的占比并不均衡。

如上表所示，Coco数据集中小目标占比达到41.4%，数量比中目标和大目标都要多。但在所有的训练集图片中，只有52.3%的图片有小目标，而中目标和大目标的分布相对来说更加均匀一些。

针对上述状况采用了Mosaic数据增强的方式，主要有几个优点：

  1. 丰富数据集

随机使用4张图片，随机缩放，再随机分布进行拼接，大大丰富了检测数据集，特别是随机缩放增加了很多小目标，让网络的鲁棒性更好。

  1. 减少GPU使用

可能会有人说，随机缩放和普通的数据增强也可以做到类似的效果，在同等size的输入下，普通的数据增强只能看到一张图像，而Mosaic增强训练时可以直接计算4张图片的数据，这样即使一个GPU也可以达到比较好的效果。

2、自适应Anchor计算

YOLOv5 Lite依旧沿用YOLOv5的Anchor计算方式，我们知道，在YOLO算法之中，针对不同的数据集，都会设置固定的Anchor。

在网络训练中，网络在初始锚框的基础上输出预测框，进而和Ground Truth进行比对，计算两者差距，再反向更新，迭代网络参数。

可以看出Anchor也是比较重要的一部分，比如Yolov5在Coco数据集上初始设定的锚框：

7b023effb9e592a4658ee78736db6aff.png

第1行是在最大的特征图上的锚框;

第2行是在中间的特征图上的锚框;

第3行是在最小的特征图上的锚框;

自适应计算Anchor的流程如下：

  1. 载入数据集，得到数据集中所有数据的wh;

  2. 将每张图片中wh的最大值等比例缩放到指定大小img_size，较小边也相应缩放;

  3. 将bboxes从相对坐标改成绝对坐标（乘以缩放后的wh）;

  4. 筛选bboxes，保留wh都大于等于两个像素的bboxes;

  5. 使用k-means聚类得到n个anchors（掉k-means包 涉及一个白化操作）;

  6. 使用遗传算法随机对anchors的wh进行变异，如果变异后效果变得更好（使用anchor_fitness方法计算得到的fitness（适应度）进行评估）就将变异后的结果赋值给anchors，如果变异后效果变差就跳过，默认变异1000次;

1.3 自适应缩放图片

在常用的目标检测算法中，不同的图片长宽都不相同，因此常用的方式是将原始图片统一缩放到一个标准尺寸，再送入检测网络中。比如Yolo算法中常用416×416，608×608等尺寸。

但Yolov5代码中对此进行了改进，也是Yolov5推理速度能够很快的一个不错的trick。作者认为，在项目实际使用时，很多图片的长宽比不同。因此缩放填充后，两端的黑边大小都不同，而如果填充的比较多，则存在信息冗余，影响推理速度。

图像高度上两端的黑边变少了，在推理时，计算量也会减少，即目标检测速度会得到提升。

ecc4ca9ad95525200b5a02f07f95843b.png

下面根据上图进行计算一下，主要是展示推理时的计算：

  • 计算缩放比例

原始图像的尺寸为640×427，与640的输入尺寸计算得到2个缩放系数分别为1.0和1.499，这里选择较小的1.0参与缩放计算；

  • 计算缩放后的尺寸

这里将原始尺寸乘以缩放系数1.0，可以分别得到长宽为640，427

  • 计算灰边填充数值

640-427=213，得到原本需要填充的高度。再采用numpy中np.mod取余数的方式，得到21个像素，再除以2，即得到图片高度两端需要填充的数值（为10【向上取整】和11【向下取整】），于是得到推理结果的尺寸为640×448。

注意

只是在测试，使用模型推理时，才采用缩减灰边的方式，提高目标检测，推理的速度。为什么np.mod函数的后面用32？因为Yolov5的网络经过5次下采样，而2的5次方，等于32。所以至少要去掉32的倍数，以免产生尺度太小走不完stride的问题，再进行取余。

2模型架构

619130994789cff7986471b91578f81e.png
目标检测全流程
3d64b3fa7caa6d7f208a5c777d64f04f.png
YOLOv5 Lite-E模型架构图

3.1 去除Focus层

为了充分理解，先来回顾一下Focus这个OP吧：

  1. class Focus(nn.Module):
  2.     # Focus wh information into c-space
  3.     def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
  4.         super().__init__()
  5.         self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
  6.        
  7.     def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
  8.         return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))

从直观理解，其实就是将图像进行切片，类似于下采样取值，这样得到4个图像：

75d592fd7571f8af63e9a8f54ea76a02.png

相当于是将空间信息绕到了通道信息中，cat后输入通道变成4倍，再通过conv得到新的featuremap,这样做的好处就是保持了下采样的信息没有丢失，图像的信息都保留了下来，但是在浅层中的应用，作者也表示了单纯是从计算量和参数量的角度上去设计，因为如上的所述的信息保存在浅层的意义并不大。所以后来查了下，作者的设计原因就是：为了减少浮点数和提高速度，而不是增加featuremap。

故，1个Focus可以替代更多的卷积层，而将空间信息聚焦到通道空间，这也会减少1像素的回归精度，而大部分的检测回归精度都不会接近1，这也是为什么Focus要放在输入的第1层上的原因。

一句话解释：Focus为了压缩网络层去提速。但是在YOLOv5 V6.0中，作者经过试验后，结论就是使用卷积替代Focus获得了更好的性能，且没有之前的一些局限性和副作用，因此此后的迭代中YOLO v5均去除了Focus操作。

YOLOv5 Lite也在此之前也不约而同的选择了摘除Focus层，避免多次采用slice操作，对于的芯片，特别是不含GPU、NPU加速的芯片，频繁的slice操作只会让缓存占用严重，加重计算处理的负担。同时，在芯片部署的时候，Focus层的转化对新手极度不友好。

3.2 ShuffleNet Backbone

YOLOv5 Lite的Backbone选择的是ShuffleNet，为什么是ShuffleNet呢？

这里给出轻量化设计的4个准则：

  1. 同等通道大小可以最小化内存访问量;

  2. 过量使用组卷积会增加MAC;

  3. 网络过于碎片化（特别是多路）会降低并行度;

  4. 不能忽略元素级操作（比如shortcut和Add）。

同时，YOLOv5 Lite避免多次使用C3 Leyer以及高通道的C3 Layer

C3 Leyer是YOLOv5作者提出的CSPBottleneck改进版本，它更简单、更快、更轻，在近乎相似的损耗上能取得更好的结果。但C3 Layer采用多路分离卷积，测试证明，频繁使用C3 Layer以及通道数较高的C3 Layer，占用较多的缓存空间，减低运行速度。

为什么通道数越高的C3 Layer会对cpu不太友好?，主要还是因为Shufflenetv2的G1准则，通道数越高，hidden channels与c1、c2的阶跃差距更大，来个不是很恰当的比喻，想象下跳一个台阶和十个台阶，虽然跳十个台阶可以一次到达，但是你需要助跑，调整，蓄力才能跳上，可能花费的时间更久

8580401e323e4e206d81159bbe909302.png

Backbone

针对ShuffleNet v2，作者首先复盘了ShuffleNetV1的问题，认为目前比较关键的问题是如何在全卷积或者分组卷积中维护大多数的卷积是平衡的。针对这个目标，作者提出了Channel Split的操作，同时构建了ShuffleNetV2。

ae1f9538cbbca67f1cab9d98f8757044.png

上图中（a）（b）是ShuffleNetV1的结构，而后面的（c）（d）是ShuffleNetV2的层结构，也是YOLOv5 Lite中的主要结构，分别对应的是结构图中的SFB1_X和SFB2_X

d4e4feb402e06a0a21c71767ec37c95b.png
SFB1_X结构对应图（d）结构
e9a2a1878ff4e25e78e3bb0213c308ac.png
SFB2_X结构对应图（c）结构

下面稍微讲一下笔者结合论文的理解：

Channel Split操作将整个特征图分为c’组（假设为A组）和c-c’（假设为B组）两个部分，主要有3个好处：

整个特征图分为2个组了，但是这样的分组又不像分组卷积一样，增加了卷积时的组数，符合准则2；这样分开之后，将A组认为是通过short-cut通道的，而B组经过的bottleneck层的输入输出的通道数就可以保持一致，符合准则1；

同时由于最后使用的concat操作，没有用TensorAdd操作，符合准则4；

可以看到，这样一个简单的通道分离的操作带来了诸多好处；但是从理论上来说，这样的结构是否还符合short-cut的初衷（即bottleneck学到的是残差Residual部分）？这里笔者也不好妄加揣测，但是可以想到的是经过后面的Channel Shuffle的乱序之后，每个通道应该都会经过一次bottleneck结构。

上述的结构是不改变输入输出通道数和特征图大小的情况，而池化操作使用图（d）代替了，跟ShuffleNetV1类似，经过这样的结构之后，图像通道数扩张为原先的2倍。

YOLO v5 Lite在Backbone中还摘除shufflenetv2 backbone的1024 conv 和 5×5 pooling。

3.3 Neck

在目标检测领域，为了更好的提取融合特征，通常在Backbone和输出层，会插入一些层，这个部分称为Neck。相当于目标检测网络的颈部，也是非常关键的。

而YOLO v5 Lite也不例外的使用了FPN+PAN的结构，但是Lite对yolov5 head进行通道剪枝，剪枝细则参考了ShuffleNet v2的设计准则，同时改进了YOLOv4中的FPN+PAN的结构，具体就是：

  • 为了最优化内存的访问和使用，选择了使用相同的通道数量（e模型Neck通道为96）；

  • 为了进一步优化内存的使用，选择了使用原始的PANet结构，还原YOLOv4的cat操作为add操作；

ff098896277db8b84e9723e73c1106a9.png
FPN+PAN的结构
d9719779774208d8abd8142c5684cf18.png
YOLOv5 Lite的Neck层
8e44a0e32f900f6283475de2c03e04bb.png
YOLOv5 Lite的FPN+PAN整体结构图

这样结合操作，FPN层自顶向下传达强语义特征（High-Level特征），而特征金字塔则自底向上传达强定位特征（Low-Level特征），两两联手，从不同的主干层对不同的检测层进行特征聚合。

FPN+PAN借鉴的是18年CVPR的PANet，当时主要应用于图像分割领域，但Alexey将其拆分应用到Yolov4中，进一步提高特征提取的能力。

230386c60eaa2cb7d765ffe7692fda10.png

3.4 Head

对于Head部分，YOLO V5 Lite并没有对YOLOv5进行改进，所以可以看到三个紫色箭头处的特征图是40×40、20×20、10×10。以及最后Prediction中用于预测的3个特征图：

①==>40×40×255

②==>20×20×255

③==>10×10×255

d74bb8cb3ad6d2999bd4fda6ac406364.png

3.5 Anchor机制及坐标变换

1、Anchor机制

对于YOLOv5，Anchor对应与Yolov3则恰恰相反，对于所设置的Anchor：

  • 第一个Yolo层是最大的特征图40×40，对应最小的anchor box。

  • 第二个Yolo层是中等的特征图20×20，对应中等的anchor box。

  • 第三个Yolo层是最小的特征图10×10，对应最大的anchor box。

  1. # anchors:
  2. #   - [10,13, 16,30, 33,23]  # P3/8
  3. #   - [30,61, 62,45, 59,119]  # P4/16
  4. #   - [116,90, 156,198, 373,326]  # P5/32
2、样本匹配策略

在yolo v3&v4中，Anchor匹配策略和SSD、Faster RCNN类似：保证每个gt bbox有一个唯一的Anchor进行对应，匹配规则就是IOU最大，并且某个gt不能在三个预测层的某几层上同时进行匹配。不考虑一个gt bbox对应多个Anchor的场合，也不考虑Anchor是否设置合理。

这里先说一下YOLOv3的匹配策略：

假设一个图中有一个目标，这个被分割成三种格子的形式，分割成13×13 、26 × 26、52 × 52 。

  • 这个目标中心坐标下采样8倍，（416/8=52），会落在 52 × 52 这个分支的所有格子中的某一个格子，落在的格子会产生3个anchor，3个anchor和目标（已经下采样8倍的目标框）分别计算iou，得到3个iou，凡是iou大于阈值0.3的，就记为正样本，就会将label[0]中这个iou大于0.3的anchor的相应位置 赋上真实框的值。

  • 这个目标中心坐标下采样16倍，（416/16=26），会落在 26 × 26 这个分支的所有格子中的某一个格子，落在的格子会产生3个anchor，3个anchor和目标（已经下采样16倍的目标框）分别计算iou，得到三个iou，凡是iou大于阈值0.3的，就记为正样本，就会将label[1]中这个iou大于0.3的anchor的相应位置 赋上真实框的值。

  • 这个目标中心坐标下采样32倍，（416/32=13），会落在 13 × 13 这个分支的所有格子中的某一个格子，落在的格子会产生3个anchor，3个anchor和目标（已经下采样32倍的目标框）分别计算iou，得到三个iou，凡是iou大于阈值0.3的，就记为正样本，就会将label[2]中这个iou大于0.3的anchor的相应位置 赋上真实框的值。

  • 如果目标所有的anchor，9个anchor，iou全部小于阈值0.3，那么选择9个anchor中和下采样后的目标框iou最大的，作为正样本，将目标真实值赋值给相应的anchor的位置。

总的来说，就是将目标先进行3种下采样，分别和目标落在的网格产生的 9个anchor分别计算iou，大于阈值0.3的记为正样本。如果9个iou全部小于0.3，那么和目标iou最大的记为正样本。对于正样本，我们在label上 相对应的anchor位置上，赋上真实目标的值。

而yolov5采用了跨网格匹配规则，增加正样本Anchor数目的做法：

对于任何一个输出层，yolov5抛弃了Max-IOU匹配规则而采用shape匹配规则，计算标签box和当前层的anchors的宽高比，即:wb/wa,hb/ha。如果宽高比大于设定的阈值说明该box没有合适的anchor，在该预测层之间将这些box当背景过滤掉。

  1. # r为目标wh和锚框wh的比值，比值在0.25到4之间的则采用该种锚框预测目标
  2. r = t[:, :, 4:6] / anchors[:, None]  # wh ratio：计算标签box和当前层的anchors的宽高比，即:wb/wa,hb/ha
  3. # 将比值和预先设置的比例anchor_t对比，符合条件为True，反之False
  4. j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t']  # compare

对于剩下的bbox，计算其落在哪个网格内，同时利用四舍五入规则，找出最近的2个网格，将这3个网格都认为是负责预测该bbox的，可以发现粗略估计正样本数相比前yolo系列，增加了3倍。code如下：

  1. # Offsets
  2. # 得到相对于以左上角为坐标原点的坐标
  3. gxy = t[:, 2:4]  # grid xy
  4. # 得到相对于右下角为坐标原点的坐标
  5. gxi = gain[[2, 3]] - gxy  # inverse
  6. # 这两个条件可以用来选择靠近的两个邻居网格
  7. # jk和lm是判断gxy的中心点更偏向哪里
  8. j, k = ((gxy % 1 < g) & (gxy > 1)).T
  9. l, m = ((gxi % 1 < g) & (gxi > 1)).T
  10. j = torch.stack((torch.ones_like(j), j, k, l, m))
  11. # yolov5不仅用目标中心点所在的网格预测该目标，还采用了距目标中心点的最近两个网格
  12. # 所以有五种情况，网格本身，上下左右，这就是repeat函数第一个参数为5的原因
  13. t = t.repeat((5, 1, 1))[j]
  14. # 这里将t复制5个，然后使用j来过滤
  15. # 第一个t是保留所有的gtbox，因为上一步里面增加了一个全为true的维度，
  16. # 第二个t保留了靠近方格左边的gtbox，
  17. # 第三个t保留了靠近方格上方的gtbox，
  18. # 第四个t保留了靠近方格右边的gtbox，
  19. # 第五个t保留了靠近方格下边的gtbox，
  20. offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]

对于YOLOv5，不同于yolov3,yolov4的是：其gt box可以跨层预测，即有些gt box在多个预测层都算正样本；同时其gt box可匹配的anchor数可为3~9个，显著增加了正样本的数量。不再是gt box落在那个网格就只由该网格内的anchor来预测，而是根据中心点的位置增加2个邻近的网格的anchor来共同预测。

如下图所示，绿点表示该gt bbox中心，现在需要额外考虑其2个最近的邻域网格的anchor也作为该gt bbox的正样本，明显增加了正样本的数量。

ece9e1faa2ddf0d861b07fe56c0ae907.png 
  1. # 输入参数pred为网络的预测输出，它是一个list包含三个检测头的输出tensor。
  2. def build_targets(self, p, targets):
  3.     '''
  4.     build_targets()函数的作用：找出与该gtbox最匹配的先验框（anchor）
  5.     '''
  6.    # 这里na为Anchor框种类数 nt为目标数 这里的na为3，nt为2
  7.     na, nt = self.na, targets.shape[0]  # number of anchors, targets
  8.    # 类别 边界框 索引 锚框
  9.     tcls, tbox, indices, anch = [], [], [], []
  10.     # 利用gain来计算目标在某一个特征图上的位置信息，初始化为1
  11.     gain = torch.ones(7, device=targets.device)  # normalized to gridspace gain
  12.     # 第2个维度复制nt遍
  13.     ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt)  # same as .repeat_interleave(nt)
  14.     # targets.shape = (na, nt, 7)（3，2，7）给每个目标加上Anchor框索引
  15.     targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2)  # append anchor indices
  16.  
  17.     g = 0.5  # bias
  18.     # 上下左右4个网格
  19.     off = torch.tensor([[0, 0],
  20.                         [1, 0], [0, 1], [-1, 0], [0, -1],  # j,k,l,m
  21.                        ], device=targets.device).float() * g  # offsets
  22.    # 处理每个检测层（3个）
  23.     for i in range(self.nl):
  24.     '''
  25.     tensor([[[ 1.25000,  1.62500],  #10,13, 16,30, 33,23 每一个数除以8
  26.              [ 2.00000,  3.75000],
  27.              [ 4.12500,  2.87500]],
  29.             [[ 1.87500,  3.81250], #30,61, 62,45, 59,119 每一个数除以16
  30.              [ 3.87500,  2.81250],
  31.              [ 3.68750,  7.43750]],
  33.             [[ 3.62500,  2.81250], #116,90, 156,198, 373,326 每一个数除以32
  34.              [ 4.87500,  6.18750],
  35.              [11.65625, 10.18750]]])
  36.     '''
  37.         # 3个anchors，已经除以当前特征图对应的stride
  38.         anchors = self.anchors[i]
  39.         # 将target中归一化后的xywh映射到3个尺度（80,80， 40,40， 20,20）的输出需要的放大系数
  40.         gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]]  # xyxy gain
  41.  
  42.         # 将xywh映射到当前特征图，即乘以对应的特征图尺寸,targets*gain将归一化的box乘以特征图尺度，将box坐标分别投影到对应的特征图上
  43.         t = targets * gain
  44.         if nt:
  45.             # r为目标wh和锚框wh的比值，比值在0.25到4之间的则采用该种锚框预测目标
  46.             # 计算当前tartget的wh和anchor的wh比值
  47.             # 如果最大比值大于预设值model.hyp['anchor_t']=4，则当前target和anchor匹配度不高，不强制回归，而把target丢弃
  48.             r = t[:, :, 4:6] / anchors[:, None]  # wh ratio：计算标签box和当前层的anchors的宽高比，即:wb/wa,hb/ha
  49.             # 筛选满足条件1/hyp['anchor_t] < target_wh / anchor_wh < hyp['anchor_t]的框
  50.             #.max(2)对第3维度的值进行max，将比值和预先设置的比例anchor_t对比，符合条件为True，反之False
  51.             j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t']  # compare
  52.             # 根据j筛选符合条件的坐标
  53.             t = t[j]  # filter
  54.             # Offsets
  55.             # 得到相对于以左上角为坐标原点的坐标
  56.             gxy = t[:, 2:4]  # grid xy
  57.             # 得到相对于右下角为坐标原点的坐标
  58.             gxi = gain[[2, 3]] - gxy  # inverse
  59.             # 这2个条件可以用来选择靠近的2个临近网格
  60.             # jk和lm是判断gxy的中心点更偏向哪里（如果中心点的相对左上角的距离大于1，小于1.5，则满足临近选择的条件）
  61.             j, k = ((gxy % 1 < g) & (gxy > 1)).T
  62.             # jk和lm是判断gxi的中心点更偏向哪里（如果中心点的相对右下角的距离大于1，小于1.5，则满足临近选择的条件）
  63.             l, m = ((gxi % 1 < g) & (gxi > 1)).T
  64.             j = torch.stack((torch.ones_like(j), j, k, l, m))
  65.             # yolov5不仅用目标中心点所在的网格预测该目标，还采用了距目标中心点的最近两个网格
  66.       # 所以有五种情况，网格本身，上下左右，这就是repeat函数第一个参数为5的原因
  67.             t = t.repeat((5, 1, 1))[j]
  68.             # 这里将t复制5个，然后使用j来过滤
  69.             # 第1个t是保留所有的gtbox，因为上一步里面增加了一个全为true的维度，
  70.             # 第2个t保留了靠近方格左边的gtbox，
  71.             # 第3个t保留了靠近方格上方的gtbox，
  72.             # 第4个t保留了靠近方格右边的gtbox，
  73.             # 第5个t保留了靠近方格下边的gtbox，
  74.             offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
  75.         else:
  76.             t = targets[0]
  77.             offsets = 0
  78.  
  79.         """
  80.         对每个bbox找出对应的正样本anchor。
  81.         a 表示当前bbox和当前层的第几个anchor匹配
  82.         b 表示当前bbox属于batch内部的第几张图片，
  83.         c 是该bbox的类别
  84.         gi,gj 是对应的负责预测该bbox的网格坐标
  85.         gxy 负责预测网格中心点坐标xy
  86.         gwh 是对应的bbox的wh
  87.         """
  88.         b, c = t[:, :2].long().T  # image, class b表示当前bbox属于该batch内第几张图片
  89.         gxy = t[:, 2:4]  # grid xy真实目标框的xy坐标
  90.         gwh = t[:, 4:6]  # grid wh真实目标框的宽高
  91.         gij = (gxy - offsets).long()   #取整
  92.         gi, gj = gij.T  # grid xy indices (gi,gj)是计算出来的负责预测该gt box的网格的坐标
  93.  
  94.         # Append
  95.         a = t[:, 6].long()  # anchor indices a表示当前gt box和当前层的第几个anchor匹配上了
  96.         indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))  # image, anchor, grid indices
  97.         tbox.append(torch.cat((gxy - gij, gwh), 1))  # gtbox与3个负责预测的网格的坐标偏移量
  98.         anch.append(anchors[a])  # anchors
  99.         tcls.append(c)  # class
  100.  
  101.     return tcls, tbox, indices, anch

最后返回四个列表：

  • class：类别

  • tbox：gtbox与3个负责预测的网格的xy坐标偏移量，gtbox的宽高

  • indices：b表示当前gtbox属于该batch内第几张图片，a表示gtbox与anchors的对应关系，负责预测的网格纵坐标，负责预测的网格横坐标

  • anch：最匹配的anchors

yolov5增加正样本的方法，最多可增大到原来的3倍，大大增加了正样本的数量，加速了模型的收敛。

目标检测重中之重可以理解为Anchor的匹配策略，当下流行的Anchor-Free不过换了一种匹配策略罢了。当下真正可创新之处在于更优的匹配策略。

3、正样本个数的增加策略

yolov5共有3个预测分支（FPN、PAN结构），共有9种不同大小的anchor，每个预测分支上有3种不同大小的anchor。Yolov5算法通过以下3种方法大幅增加正样本个数：

  1. 跨预测分支预测：假设一个ground truth框可以和2个甚至3个预测分支上的anchor匹配，则这2个或3个预测分支都可以预测该ground truth框，即一个ground truth框可以由多个预测分支来预测。

  2. 跨网格预测：假设一个ground truth框落在了某个预测分支的某个网格内，则该网格有左、上、右、下4个邻域网格，根据ground truth框的中心位置，将最近的2个邻域网格也作为预测网格，也即一个ground truth框可以由3个网格来预测；

  3. 跨anchor预测：假设一个ground truth框落在了某个预测分支的某个网格内，该网格具有3种不同大小anchor，若ground truth可以和这3种anchor中的多种anchor匹配，则这些匹配的anchor都可以来预测该ground truth框，即一个ground truth框可以使用多种anchor来预测。

4、坐标变换

对于之前的YOLOv3和YOLOv4，使用的是如下图所示的坐标表示形式：

f632bd2839dc3c37fb6f5c8fefee6aaf.png
YOLOv3和YOLOv4的坐标表示
a5d4ffcb047186f6e05fae40f4e6b11c.png

、、、分别是即边界框bbox相对于feature map的位置和宽高；

和分别代表feature map中grid cell的左上角坐标，在yolo中每个grid cell在feature map中的宽和高均为1；

和分别代表Anchor映射到feature map中的的宽和高，anchor box原本设定是相对于坐标系下的坐标，需要除以stride如32映射到feature map坐标系中；

、、、这4个参数化坐标是网络学习的目标，其中,是预测的坐标偏移值，和是尺度缩放，sigma代表sigmoid函数。

与faster rcnn和ssd中的参数化坐标不同的是x和y的回归方式，YOLO v3&v4使用了sigmoid函数进行偏移量的规则化，而faster和ssd中对x，y除以anchor的宽和高进行规则化。

YOLOv5参数化坐标的方式和yolo v3&v4是不一样的，如下：

5f3abcd6a603a701e1fe5c4a0d89a8a3.png

YOLOv5的坐标表示方法

用公式表示如下：

9e26f7753be66c4efc19bb9e800fb502.png

  1. xy = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i]  # xy
  2. wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh

这样，pxy的取值范围是[-0.5,1.5]，pwh的取值范围是(0,4×anchors[i])，这是因为采用了跨网格匹配规则，要跨网格预测了。

为什么这么改造呢？

91e3da23bc315fa3322185119f0daf46.png

可以看出，对于不同的和，当他们大于零比较多时，YOLOv5的反馈更加平滑，相对于v3、v4也就更容易收敛。

5、为什么进行坐标参数化？
为什么要学习偏移而不是实际值？

Anchor已经粗略地“框住了”输入图像中的目标，明显的一个问题是，框的不够准确。因为受限于Anchor的生成方式，Anchor的坐标永远都是固定的那几个。所以，如果我们预测相对于Anchor的offset，那么，就可以通过预测的offset调整锚框位置，从而得到更精准的bounding box。

为什么要学习偏移系数而不是偏移量？

首先，对于预测的bounding box的w和h可以通过anchor进行缩放，但有一个基本的要求，就是h和w都必须为正值，而网络最后一层的预测输出是没法保证正负的，所以最简单的方法就是对预测输出求exp，这样就保证了预测值恒为正。那么反过来，对预测目标就是求log。

其次，对cx和cy除以anchor的宽和高的处理是为了做尺度归一化。例如，大的box的绝对偏移量一般较大，而小的box的绝对偏移量一般较小，除以宽和高消除这种影响。即两个框大小不一，但相对值却一致。

为什么都要进行Sigmoid计算？

yolov5需要的训练数据的label是根据原图尺寸归一化了的，这样做是因为怕大的边框的影响比小的边框影响大，因此做了归一化的操作，这样大的和小的边框都会被同等看待了，而且训练也容易收敛。所以在网络输出的部分也需要对输出进行归一化操作，因此选择了Sigmoid计算。

4输出端

4.1 优化方法

YOLO V5的作者提供了2个优化函数Adam和SGD，并都预设了与之匹配的训练超参数。默认为SGD。YOLO V4使用SGD。

YOLO V5的作者建议是，如果需要训练较小的自定义数据集，Adam是更合适的选择，尽管Adam的学习率通常比SGD低。但是如果训练大型数据集，对于YOLOV5来说SGD效果比Adam好。

实际上学术界上对于SGD和Adam哪个更好，一直没有统一的定论，取决于实际项目情况。

4.2 损失函数

8e86ebddfea73dba84630509be5e16b4.png

通过上图我们可以看到，对于图中的目标，都会输出class_num+4+1长度的向量，比如针对coco数据集有80个类别，就会输出长度为85的特征向量，其中所包含的内容如下图所示：

74b176db0c480da33312d5d300659ffa.png

图中向量包含4个坐标信息，一个包含目标概率和80个类别得分，换句话解释就是“这个图像中是否有目标（物体出现的概率）？有的话是什么（80类的类别得分）？然后就是这个目标物体在哪里（box坐标位置）？”

其实面对上述的3个输出，也对应YOLOv5的3个分支的，其分别是obj分支、cls分支和box分支。

1、obj分支

obj分支输出的是该anchor中是否含有物体的概率，默认使用BCEWithLogits Loss。

BCEWithLogitsLoss是将BCELoss（BCE：Binary cross entropy）和sigmoid融合了，也就是说省略了sigmoid这个步骤；BCELoss的数学公式如下：

258e9a3c92821412e538a00726a7f57d.png

  1. class BCEBlurWithLogitsLoss(nn.Module):
  2.     # BCEwithLogitLoss() with reduced missing label effects.
  3.     def __init__(self, alpha=0.05):
  4.         super(BCEBlurWithLogitsLoss, self).__init__()
  5.         self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none')  # must be nn.BCEWithLogitsLoss()
  6.         self.alpha = alpha
  7.  
  8.     def forward(self, pred, true):
  9.         loss = self.loss_fcn(pred, true)
  10.         pred = torch.sigmoid(pred)  # prob from logits
  11.         dx = pred - true  # reduce only missing label effects
  12.         # dx = (pred - true).abs()  # reduce missing label and false label effects
  13.         alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4))
  14.         loss *= alpha_factor
  15.         return loss.mean()
2、cls分支

cls分支输出的是该anchor属于哪一类的概率，也默认使用BCEWithLogits Loss。

  1. class BCEBlurWithLogitsLoss(nn.Module):
  2.     # BCEwithLogitLoss() with reduced missing label effects.
  3.     def __init__(self, alpha=0.05):
  4.         super(BCEBlurWithLogitsLoss, self).__init__()
  5.         self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none')  # must be nn.BCEWithLogitsLoss()
  6.         self.alpha = alpha
  7.  
  8.     def forward(self, pred, true):
  9.         loss = self.loss_fcn(pred, true)
  10.         pred = torch.sigmoid(pred)  # prob from logits
  11.         dx = pred - true  # reduce only missing label effects
  12.         # dx = (pred - true).abs()  # reduce missing label and false label effects
  13.         alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4))
  14.         loss *= alpha_factor
  15.         return loss.mean()

例如，对于coco数据集上训练的YOLO的每个anchor的维度都是85，前5个属性是(Cx,Cy,w,h,confidence)，confidence对应obj，后80个维度对应cls。

3、box分支

这里的box分支输出的便是物体的具体位置信息了，通过前面对于坐标参数化的分析可以知道，具体的输出4个值为、、以及，然后通过前面的参数化反转方式与GT进行计算loss，对于回归损失，yolov3使用的loss是smooth l1损失。Yolov5的边框(Bounding box)回归的损失函数默认使用的是CIoU，不是GIoU，不是DIoU，是CIoU。

b81bd3cb7cdbd9b34ba7d3a173dddf78.png

回归损失的发展

下面用一张图粗略看一下IoU,GIoU,DIoU,CIoU：

ca36248c6a319e1bd4a086a91f724aae.png4eb7edc2785d1e835562e8c82aaef9ff.png

式中，、，、分别代表候选框的中心点坐标。

下面大概说一下每个IOU损失的局限性：

IoU Loss 有2个缺点：

  • 当预测框和目标框不相交时，IoU(A,B)=0时，不能反映A,B距离的远近，此时损失函数不可导，IoU Loss 无法优化两个框不相交的情况。

  • 假设预测框和目标框的大小都确定，只要两个框的相交值是确定的，其IoU值是相同时，IoU值不能反映两个框是如何相交的。

GIoU Loss 有1个缺点：

  • 当目标框完全包裹预测框的时候，IoU和GIoU的值都一样，此时GIoU退化为IoU, 无法区分其相对位置关系；

DIoU Loss 有1个缺点：

  • 当预测框的中心点的位置都一样时, DIoU无法区分候选框位置的质量；

综合IoU、GIoU、DIoU的种种局限性，总结一个好的bounding box regressor包含3个要素：

  1. 、Overlapping area

  2. 、Central point distance

  3. 、Aspect ratio

因此，YOLOv5使用的是CIoU Loss：

24784c4d72b9eca4709c0d8c9d163686.png

  1. iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True)  # iou(prediction, target)
  2. lbox += (1.0 - iou).mean()  # iou loss
  3.  
  4. def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
  5.     # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
  6.     box2 = box2.T
  7.  
  8.     # Get the coordinates of bounding boxes
  9.     if x1y1x2y2:  # x1, y1, x2, y2 = box1
  10.         b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
  11.         b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
  12.     else:  # transform from xywh to xyxy
  13.         b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
  14.         b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
  15.         b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
  16.         b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2
  17.  
  18.     # Intersection area
  19.     inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
  20.             (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
  21.  
  22.     # Union Area
  23.     w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
  24.     w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
  25.     union = w1 * h1 + w2 * h2 - inter + eps
  26.  
  27.     iou = inter / union
  28.     if CIoU or DIoU or GIoU:
  29.         cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1)  # convex (smallest enclosing box) width
  30.         ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1)  # convex height
  31.         if CIoU or DIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
  32.             c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squared
  33.             rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
  34.                     (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center distance squared
  35.             if CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
  36.                 v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
  37.                 with torch.no_grad():
  38.                     alpha = v / (v - iou + (1 + eps))
  39.                 return iou - (rho2 / c2 + v * alpha)  # CIoU
  40.             return iou - rho2 / c2  # DIoU
  41.         c_area = cw * ch + eps  # convex area
  42.         return iou - (c_area - union) / c_area  # GIoU https://arxiv.org/pdf/1902.09630.pdf
  43.     return iou  # IoU
4、Loss计算
  1. def compute_loss(p, targets, model):  # predictions, targets, model
  2.     device = targets.device
  3.     lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
  4.     tcls, tbox, indices, anchors = build_targets(p, targets, model)  # targets
  5.     h = model.hyp  # hyperparameters
  6.  
  7.     # Define criteria
  8.     BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([h['cls_pw']])).to(device)
  9.     BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([h['obj_pw']])).to(device)
  10.  
  11.     # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
  12.     cp, cn = smooth_BCE(eps=0.0)
  13.  
  14.     # Focal loss
  15.     g = h['fl_gamma']  # focal loss gamma
  16.     if g > 0:
  17.         BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
  18.    。。。。。。

4.3、后处理之DIoU NMS

350d4a5cab406fa9ce50afb1ee800a52.pngcf7b3ea43d636cb70b90bc2323650038.pngb58f4620cdb39a0755270a4081159a5e.png

在上图重叠的摩托车检测中，中间的摩托车因为考虑边界框中心点的位置信息，也可以回归出来。因此在重叠目标的检测中，DIOU_nms的效果优于传统的nms。

为什么不用CIoU NMS呢？

因为前面讲到的CIOU loss，是在DIOU loss的基础上，添加的影响因子，包含ground truth标注框的信息，在训练时用于回归。但在测试过程中，并没有ground truth的信息，不用考虑影响因子，因此直接用DIOU NMS即可。

5YOLOv5 Lite训练自己的数据集

5.1 git clone仓库代码

clone YOLOv5 Lite代码并下载coco的预训练权重。

  1. $ git clone https://github.com/ppogg/YOLOv5-Lite
  2. $ cd YOLOv5-Lite
  3. $ pip install -r requirements.txt

5.2 处理数据集格式

这里可以直接参考coco128的数据集形式进行整理：

文件夹目录如下图所示：

d53d913b201c194263d2871390892719.png8712f998765ca36bf3b56e6a52682a83.png

5.3 配置超参数

主要是配置data文件夹下的coco128.yaml中的数据集位置和种类：

e57f80af087c00b4e917f58931456a98.png7180d87c82c76e7fb07af4c817b81b41.png

5.4 配置模型

这里主要是配置models目录下的模型yaml文件，主要是进去后修改nc这个参数来进行类别的修改。

e7cb39bab76db0a44906f1f4782eec17.png

修改 nc 参数

目前支持的模型种类如下所示：

e1e5d8d083dc0f34f2813c873877fa80.png

5.3 训练

  1. $ python train.py --data coco.yaml --cfg v5lite-e.yaml --weights v5lite-e.pt --batch-size 128
  2.                                          v5lite-s.yaml --weights v5lite-s.pt --batch-size 128
  3.                                          v5lite-c.yaml           v5lite-c.pt               96
  4.                                          v5lite-g.yaml           v5lite-g.pt               64

如果您是多卡进行训练，则：

$ python -m torch.distributed.launch --nproc_per_node 2 train.py

5.4 检测结果

$ python path/to/detect.py --weights v5lite-e.pt --source 0 img.jpg        # image

3e22000ccecd01c77cb125d3db67f582.png

检测结果

5TensorRT部署

5.1 目标检测常见的落地形式

7571eefd949f7347447ccb6bc961a7e1.png

1、TensorRT是什么

TensorRT是推理优化器，能对训练好的模型进行优化。可以理解为只有前向传播的深度学习框架，这个框架可以将Caffe，TensorFlow的网络模型解析，然后与TensorRT中对应的层进行一一映射，把其他框架的模型统一全部转换到TensorRT中，然后在TensorRT中可以针对NVIDIA自家GPU实施优化策略，并进行部署加速。当你的网络训练完之后，可以将训练模型文件直接丢进TensorRT中，而不再需要依赖深度学习框架（Caffe，TensorFlow等)。

2、本文AI部署流程

先把onnx转化为TensorRT的Engine文件，然后让c++环境下的TensorRT直接加载Engine文件，从而构建engine，本文主要讲解onnx转换至Engine，然后进行基于TensorRT的C++推理检测。

61287a70d1ac4fc6cc053086b3602d8b.png

转换和部署模型5个基本步骤:

  • step1：获取模型

  • step2：选择batchsize

  • step3：选择精度

  • step4：模型转换

  • step5：模型部署

5.2 ONNX-TensorRT的部署流程

1、ONNX转化为TRT Engine
  1. # 导出onnx文件
  2. python export.py ---weights weights/v5lite-g.pt --batch-size 1 --imgsz 640 --include onnx --simplify
  3.  
  4. # 使用TensorRT官方的trtexec工具将onnx文件转换为engine
  5. trtexec --explicitBatch --onnx=./v5lite-g.onnx --saveEngine=v5lite-g.trt --fp16

闲话不多说，这里已经拿到了trt的engine，那么如何进行推理呢？总的来说，分为3步：

  1. 首先load你的engine，拿到一个ICudaEngine, 这个是TensorRT推理的核心；

  2. 定位模型的输入和输出，有几个输入有几个输出；

  3. forward模型，然后拿到输出，对输出进行后处理。

当然这里最核心的东西其实就两个，一个是如何导入拿到CudaEngine，第二个是比较麻烦的后处理。

2、加载TRT Engine
  1. bool Model::readTrtFile() {
  2.     std::string cached_engine;
  3.     std::fstream file;
  4.     std::cout << "loading filename from:" << engine_file << std::endl;
  5.     nvinfer1::IRuntime *trtRuntime;
  6.     file.open(engine_file, std::ios::binary | std::ios::in);
  7.  
  8.     if (!file.is_open()) {
  9.         std::cout << "read file error: " << engine_file << std::endl;
  10.         cached_engine = "";
  11.     }
  12.  
  13.     while (file.peek() != EOF) {
  14.         std::stringstream buffer;
  15.         buffer << file.rdbuf();
  16.         cached_engine.append(buffer.str());
  17.     }
  18.     file.close();
  19.  
  20.     trtRuntime = nvinfer1::createInferRuntime(gLogger.getTRTLogger());
  21.     engine = trtRuntime->deserializeCudaEngine(cached_engine.data(), cached_engine.size(), nullptr);
  22.     std::cout << "deserialize done" << std::endl;
  23.  
  24. }
  25.  
  26. // 加载 TensorRT Engine
  27. void v5Lite::LoadEngine() {
  28.     // create and load engine
  29.     std::fstream existEngine;
  30.     existEngine.open(engine_file, std::ios::in);
  31.     // 如果存在已经转换完成的TensorRT Engine文件，则直接加载
  32.     if (existEngine) {
  33.         readTrtFile(engine_file, engine);
  34.         assert(engine != nullptr);
  35.     } 
  36.     // 如果不存在已经转换完成的TensorRT Engine文件，则直接加载ONNX权重进行在线生成
  37.     else {
  38.         onnxToTRTModel(onnx_file, engine_file, engine, BATCH_SIZE);
  39.         assert(engine != nullptr);
  40.     }
  41. }
3、后处理之坐标转换

20e328f835ad96910d556a340857d621.png

YOLOv5的坐标转换方法 
  1. std::vector<std::vector<V5lite::DetectRes>> V5lite::postProcess(const std::vector<cv::Mat> &vec_Mat, float *output,
  2.                                                                 const int &outSize) {
  3.     std::vector<std::vector<DetectRes>> vec_result;
  4.     int index = 0;
  5.     for (const cv::Mat &src_img : vec_Mat)
  6.     {
  7.         std::vector<DetectRes> result;
  8.         float ratio = float(src_img.cols) / float(IMAGE_WIDTH) > float(src_img.rows) / float(IMAGE_HEIGHT)  ? float(src_img.cols) / float(IMAGE_WIDTH) : float(src_img.rows) / float(IMAGE_HEIGHT);
  9.         float *out = output + index * outSize;
  10.         int position = 0;
  11.         for (int n = 0; n < (int)grids.size(); n++)
  12.         {
  13.             for (int c = 0; c < grids[n][0]; c++)
  14.             {
  15.                 std::vector<int> anchor = anchors[n * grids[n][0] + c];
  16.                 for (int h = 0; h < grids[n][1]; h++)
  17.                     for (int w = 0; w < grids[n][2]; w++)
  18.                     {
  19.                         float *row = out + position * (CATEGORY + 5);
  20.                         position++;
  21.                         DetectRes box;
  22.                         auto max_pos = std::max_element(row + 5, row + CATEGORY + 5);
  23.                         box.prob = row[4] * row[max_pos - row];
  24.                         if (box.prob < obj_threshold)
  25.                             continue;
  26.                         box.classes = max_pos - row - 5;
  27.                         // 坐标的反参数化，和前文的坐标转换对接
  28.                         box.x = (row[0] * 2 - 0.5 + w) / grids[n][2] * IMAGE_WIDTH * ratio;
  29.                         box.y = (row[1] * 2 - 0.5 + h) / grids[n][1] * IMAGE_HEIGHT * ratio;
  30.                         box.w = pow(row[2] * 2, 2) * anchor[0] * ratio;
  31.                         box.h = pow(row[3] * 2, 2) * anchor[1] * ratio;
  32.                         result.push_back(box);
  33.                     }
  34.             }
  35.         }
  36.         NmsDetect(result);
  37.         vec_result.push_back(result);
  38.         index++;
  39.     }
  40.     return vec_result;
  41. }
4、进行模型推理
  1. // 推理整个文件夹的文件
  2. bool YOLOv5::InferenceFolder(const std::string &folder_name) {
  3.     // 读取文件夹下面的文件，并返回为一个 string vector 迭代器
  4.     std::vector<std::string> sample_images = readFolder(folder_name);
  5.     //get context
  6.     assert(engine != nullptr);
  7.     // 创建上下文，创建一些空间来存储中间值。一个engine可以创建多个context，分别执行多个推理任务。
  8.     context = engine->createExecutionContext();
  9.     assert(context != nullptr);
  10.  
  11.     // 传递给Engine的输入输出buffers指针，这里对应一个输入和一个输出
  12.     assert(engine->getNbBindings() == 2);
  13.     void *buffers[2];
  14.     std::vector<int64_t> bufferSize;
  15.     int nbBindings = engine->getNbBindings();
  16.     bufferSize.resize(nbBindings);
  17.  
  18.     for (int i = 0; i < nbBindings; ++i) {
  19.         // 获取输入或输出的维度信息
  20.         nvinfer1::Dims dims = engine->getBindingDimensions(i);
  21.         // 获取输入或输出的数据类型信息
  22.         nvinfer1::DataType dtype = engine->getBindingDataType(i);
  23.         int64_t totalSize = volume(dims) * 1 * getElementSize(dtype);
  24.         bufferSize[i] = totalSize;
  25.         std::cout << "binding" << i << ": " << totalSize << std::endl;
  26.         // &buffers是双重指针 相当于改变指针本身，这里就是把输入或输出进行向量化操作
  27.         cudaMalloc(&buffers[i], totalSize);
  28.     }
  29.  
  30.     //get stream
  31.     cudaStream_t stream;
  32.     // 创建 Stream
  33.     cudaStreamCreate(&stream);
  34.  
  35.     int outSize = bufferSize[1] / sizeof(float) / BATCH_SIZE;
  36.     // 执行推理
  37.     EngineInference(sample_images, outSize, buffers, bufferSize, stream);
  38.  
  39.     // 释放 stream 和 buffers
  40.     cudaStreamDestroy(stream);
  41.     cudaFree(buffers[0]);
  42.     cudaFree(buffers[1]);
  43.  
  44.     // destroy the engine
  45.     context->destroy();
  46.     engine->destroy();
  47. }
  48.  
  49. void YOLOv5::EngineInference(const std::vector<std::string> &image_list, const int &outSize, void **buffers,
  50.                              const std::vector<int64_t> &bufferSize, cudaStream_t stream) {
  51.     int index = 0;
  52.     int batch_id = 0;
  53.     std::vector<cv::Mat> vec_Mat(BATCH_SIZE);
  54.     std::vector<std::string> vec_name(BATCH_SIZE);
  55.     float total_time = 0;
  56.     // 遍历图像路径list
  57.     for (const std::string &image_name : image_list)
  58.     {
  59.         index++;
  60.         std::cout << "Processing: " << image_name << std::endl;
  61.         // 读取图像内容到cv_mat
  62.         cv::Mat src_img = cv::imread(image_name);
  63.         // 把图像和图像名分别保存在vec_Mat和vec_name之中
  64.         if (src_img.data)
  65.         {
  66.             vec_Mat[batch_id] = src_img.clone();
  67.             vec_name[batch_id] = image_name;
  68.             batch_id++;
  69.         }
  70.         if (batch_id == BATCH_SIZE or index == image_list.size())
  71.         {
  72.             // 声明时间戳 t_start_pre
  73.             auto t_start_pre = std::chrono::high_resolution_clock::now();
  74.             std::cout << "prepareImage" << std::endl;
  75.             std::vector<float>curInput = prepareImage(vec_Mat);
  76.             auto t_end_pre = std::chrono::high_resolution_clock::now();
  77.             // 至此，prepare Image的时间已经计算完成
  78.             float total_pre = std::chrono::duration<float, std::milli>(t_end_pre - t_start_pre).count();
  79.             std::cout << "prepare image take: " << total_pre << " ms." << std::endl;
  80.             total_time += total_pre;
  81.             batch_id = 0;
  82.             if (!curInput.data()) {
  83.                 std::cout << "prepare images ERROR!" << std::endl;
  84.                 continue;
  85.             }
  86.             // 将数据从CPU端传送到GPU端
  87.             std::cout << "host2device" << std::endl;
  88.             cudaMemcpyAsync(buffers[0], curInput.data(), bufferSize[0], cudaMemcpyHostToDevice, stream);
  89.  
  90.             // 执行推理
  91.             std::cout << "execute" << std::endl;
  92.             auto t_start = std::chrono::high_resolution_clock::now();
  93.             context->execute(BATCH_SIZE, buffers);
  94.             auto t_end = std::chrono::high_resolution_clock::now();
  95.             float total_inf = std::chrono::duration<float, std::milli>(t_end - t_start).count();
  96.             std::cout << "Inference take: " << total_inf << " ms." << std::endl;
  97.             total_time += total_inf;
  98.             std::cout << "execute success" << std::endl;
  99.             std::cout << "device2host" << std::endl;
  100.             std::cout << "post process" << std::endl;
  101.             auto r_start = std::chrono::high_resolution_clock::now();
  102.             auto *out = new float[outSize * BATCH_SIZE];
  103.             // Copy GPU端的推理结果到CPU端
  104.             cudaMemcpyAsync(out, buffers[1], bufferSize[1], cudaMemcpyDeviceToHost, stream);
  105.             // 阻塞当前程序的执行，直到所有任务都处理完毕，这样可以将计算和主机与设备之前的传输并行化，提高效率。
  106.             cudaStreamSynchronize(stream);
  107.             // 进行后处理操作
  108.             auto boxes = postProcess(vec_Mat, out, outSize);
  109.             auto r_end = std::chrono::high_resolution_clock::now();
  110.             float total_res = std::chrono::duration<float, std::milli>(r_end - r_start).count();
  111.             std::cout << "Post process take: " << total_res << " ms." << std::endl;
  112.             total_time += total_res;
  113.             for (int i = 0; i < (int)vec_Mat.size(); i++)
  114.             {
  115.                 auto org_img = vec_Mat[i];
  116.                 if (!org_img.data)
  117.                     continue;
  118.                 auto rects = boxes[i];
  119.                 for(const auto &rect : rects)
  120.                 {
  121.                     char t[256];
  122.                     sprintf(t, "%.2f", rect.prob);
  123.                     std::string name = coco_labels[rect.classes] + "-" + t;
  124.                     // 图书添加文字
  125.                     cv::putText(org_img, name, cv::Point(rect.x - rect.w / 2, rect.y - rect.h / 2 - 5), cv::FONT_HERSHEY_COMPLEX, 0.7, class_colors[rect.classes], 2);
  126.                     // 绘制矩形框
  127.                     cv::Rect rst(rect.x - rect.w / 2, rect.y - rect.h / 2, rect.w, rect.h);
  128.                     cv::rectangle(org_img, rst, class_colors[rect.classes], 2, cv::LINE_8, 0);
  129.                 }
  130.                 int pos = vec_name[i].find_last_of(".");
  131.                 std::string rst_name = vec_name[i].insert(pos, "_");
  132.                 std::cout << rst_name << std::endl;
  133.                 // 保存检测结果
  134.                 cv::imwrite(rst_name, org_img);
  135.             }
  136.             vec_Mat = std::vector<cv::Mat>(BATCH_SIZE);
  137.             delete[] out;
  138.         }
  139.     }
  140.     std::cout << "Average processing time is " << total_time / image_list.size() << "ms" << std::endl;
  141. }
5、后处理之NMS C++实现
  1. void V5lite::NmsDetect(std::vector<DetectRes> &detections) {
  2.     sort(detections.begin(), detections.end(), [=](const DetectRes &left, const DetectRes &right) {
  3.         return left.prob > right.prob;
  4.     });
  5.  
  6.     for (int i = 0; i < (int)detections.size(); i++)
  7.         for (int j = i + 1; j < (int)detections.size(); j++)
  8.         {
  9.             if (detections[i].classes == detections[j].classes)
  10.             {   // 计算DIoU的值
  11.                 float iou = IOUCalculate(detections[i], detections[j]);
  12.                 if (iou > nms_threshold)
  13.                     detections[j].prob = 0;
  14.             }
  15.         }
  16.  
  17.     detections.erase(std::remove_if(detections.begin(), detections.end(), [](const DetectRes &det)
  18.     { return det.prob == 0; }), detections.end());
  19. }
  20.  
  21. // 计算 DIOU
  22. float v5Lite::IOUCalculate(const YOLOv5::DetectRes &det_a, const YOLOv5::DetectRes &det_b) {
  23.     cv::Point2f center_a(det_a.x, det_a.y);
  24.     cv::Point2f center_b(det_b.x, det_b.y);
  25.     // 计算左上角角点坐标
  26.     cv::Point2f left_up(std::min(det_a.x - det_a.w / 2, det_b.x - det_b.w / 2),
  27.                         std::min(det_a.y - det_a.h / 2, det_b.y - det_b.h / 2));
  28.     // 计算右下角角点坐标
  29.     cv::Point2f right_down(std::max(det_a.x + det_a.w / 2, det_b.x + det_b.w / 2),
  30.                            std::max(det_a.y + det_a.h / 2, det_b.y + det_b.h / 2));
  31.  
  32.     // 计算框的中心点距离
  33.     float distance_d = (center_a - center_b).x * (center_a - center_b).x + (center_a - center_b).y * (center_a - center_b).y;
  34.     // 计算框的角点距离
  35.     float distance_c = (left_up - right_down).x * (left_up - right_down).x + (left_up - right_down).y * (left_up - right_down).y;
  36.  
  37.     float inter_l = det_a.x - det_a.w / 2 > det_b.x - det_b.w / 2 ? det_a.x - det_a.w / 2 : det_b.x - det_b.w / 2;
  38.     float inter_t = det_a.y - det_a.h / 2 > det_b.y - det_b.h / 2 ? det_a.y - det_a.h / 2 : det_b.y - det_b.h / 2;
  39.     float inter_r = det_a.x + det_a.w / 2 < det_b.x + det_b.w / 2 ? det_a.x + det_a.w / 2 : det_b.x + det_b.w / 2;
  40.     float inter_b = det_a.y + det_a.h / 2 < det_b.y + det_b.h / 2 ? det_a.y + det_a.h / 2 : det_b.y + det_b.h / 2;
  41.     if (inter_b < inter_t || inter_r < inter_l)
  42.         return 0;
  43.  
  44.     // 计算交集
  45.     float inter_area = (inter_b - inter_t) * (inter_r - inter_l);
  46.     // 计算并集
  47.     float union_area = det_a.w * det_a.h + det_b.w * det_b.h - inter_area;
  48.     if (union_area == 0)
  49.         return 0;
  50.     else
  51.         return inter_area / union_area - distance_d / distance_c;
  52. }

CMakeLists.txt如下：

  1. cmake_minimum_required(VERSION 3.5)
  2.  
  3. project(v5lite_trt)
  4.  
  5. set(CMAKE_CXX_STANDARD 14)
  6.  
  7. # CUDA
  8. find_package(CUDA REQUIRED)
  9. message(STATUS "Find CUDA include at ${CUDA_INCLUDE_DIRS}")
  10. message(STATUS "Find CUDA libraries: ${CUDA_LIBRARIES}")
  11.  
  12. # TensorRT
  13. set(TENSORRT_ROOT "/home/chaucer/TensorRT-8.0.1.6")
  14.  
  15. find_path(TENSORRT_INCLUDE_DIR NvInfer.h
  16.         HINTS ${TENSORRT_ROOT} PATH_SUFFIXES include/)
  17. message(STATUS "Found TensorRT headers at ${TENSORRT_INCLUDE_DIR}")
  18.  
  19. find_library(TENSORRT_LIBRARY_INFER nvinfer
  20.         HINTS ${TENSORRT_ROOT} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR}
  21.         PATH_SUFFIXES lib lib64 lib/x64)
  22.  
  23. find_library(TENSORRT_LIBRARY_ONNXPARSER nvonnxparser
  24.         HINTS  ${TENSORRT_ROOT} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR}
  25.         PATH_SUFFIXES lib lib64 lib/x64)
  26.  
  27. set(TENSORRT_LIBRARY ${TENSORRT_LIBRARY_INFER} ${TENSORRT_LIBRARY_ONNXPARSER})
  28. message(STATUS "Find TensorRT libs: ${TENSORRT_LIBRARY}")
  29.  
  30. # OpenCV
  31. find_package(OpenCV REQUIRED)
  32. message(STATUS "Find OpenCV include at ${OpenCV_INCLUDE_DIRS}")
  33. message(STATUS "Find OpenCV libraries: ${OpenCV_LIBRARIES}")
  34.  
  35. set(COMMON_INCLUDE ./includes/common)
  36. set(YAML_INCLUDE ./includes/yaml-cpp/include)
  37. set(YAML_LIB_DIR ./includes/yaml-cpp/libs)
  38.  
  39. include_directories(${CUDA_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR} ${OpenCV_INCLUDE_DIRS} ${COMMON_INCLUDE} ${YAML_INCLUDE})
  40. link_directories(${YAML_LIB_DIR})
  41.  
  42. add_executable(v5lite_trt main.cpp v5lite.cpp)
  43. target_link_libraries(v5lite_trt ${OpenCV_LIBRARIES} ${CUDA_LIBRARIES} ${TENSORRT_LIBRARY} yaml-cpp)
  1. mkdir build
  2. cd build
  3. cmake ..
  4. make -j8
  5. v5lite_trt ../config.yaml ../samples/
5、检测结果和时间

86635a8d57b37b6184c18839eec33fb5.png3ec9fb8eb31853277380cf41dbf5e9a2.png

6参考

[1].https://github.com/ppogg/YOLOv5-Lite
[2].https://zhuanlan.zhihu.com/p/400545131
[3].https://github.com/ultralytics/yolov5
[4].https://zhuanlan.zhihu.com/p/172121380
[5].https://zhuanlan.zhihu.com/p/143747206

本文仅做学术分享，如有侵权，请联系删文。

3D视觉精品课程推荐：

1.面向自动驾驶领域的多传感器数据融合技术

2.面向自动驾驶领域的3D点云目标检测全栈学习路线！(单模态+多模态/数据+代码)
3.彻底搞透视觉三维重建：原理剖析、代码讲解、及优化改进
4.国内首个面向工业级实战的点云处理课程
5.激光-视觉-IMU-GPS融合SLAM算法梳理和代码讲解
6.彻底搞懂视觉-惯性SLAM：基于VINS-Fusion正式开课啦
7.彻底搞懂基于LOAM框架的3D激光SLAM: 源码剖析到算法优化
8.彻底剖析室内、室外激光SLAM关键算法原理、代码和实战(cartographer+LOAM +LIO-SAM)

9.从零搭建一套结构光3D重建系统[理论+源码+实践]

10.单目深度估计方法：算法梳理与代码实现

11.自动驾驶中的深度学习模型部署实战

12.相机模型与标定(单目+双目+鱼眼）

13.重磅！四旋翼飞行器：算法与实战

14.ROS2从入门到精通：理论与实战

重磅！3DCVer-学术论文写作投稿 交流群已成立

扫码添加小助手微信，可申请加入3D视觉工坊-学术论文写作与投稿 微信交流群，旨在交流顶会、顶刊、SCI、EI等写作与投稿事宜。

同时也可申请加入我们的细分方向交流群，目前主要有3D视觉、CV&深度学习、SLAM、三维重建、点云后处理、自动驾驶、多传感器融合、CV入门、三维测量、VR/AR、3D人脸识别、医疗影像、缺陷检测、行人重识别、目标跟踪、视觉产品落地、视觉竞赛、车牌识别、硬件选型、学术交流、求职交流、ORB-SLAM系列源码交流、深度估计等微信群。

一定要备注：研究方向+学校/公司+昵称，例如：”3D视觉 + 上海交大 + 静静“。请按照格式备注，可快速被通过且邀请进群。原创投稿也请联系。

adb20072850a70c798f41cd5f1ba170e.png

▲长按加微信群或投稿

1326916accb62ae9d7e6947f7d3939c9.png

▲长按关注公众号

3D视觉从入门到精通知识星球：针对3D视觉领域的视频课程（三维重建系列、三维点云系列、结构光系列、手眼标定、相机标定、激光/视觉SLAM、自动驾驶等）、知识点汇总、入门进阶学习路线、最新paper分享、疑问解答五个方面进行深耕，更有各类大厂的算法工程人员进行技术指导。与此同时，星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息，打造成集技术与就业为一体的铁杆粉丝聚集区，近4000星球成员为创造更好的AI世界共同进步，知识星球入口：

学习3D视觉核心技术，扫描查看介绍，3天内无条件退款

776a7bb4fa0d1519da76eed7c174e675.png

 圈里有高质量教程资料、答疑解惑、助你高效解决问题

觉得有用，麻烦给个赞和在看~  

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/小蓝xlanll/article/detail/140384
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号