YOLOv10(6):YOLOv10基于TensorRT的部署(基于INetworkDefinition)

1. 写在前面

我们在前面已经讲过很多关于YOLOv10的一些知识点,也简单理了一下如何训练自己的数据。

现在本篇文章主要是讲解一下如何在TensorRT中部署YOLOv10,相信经过这一步,各位小伙伴已经能够无限的接近于将YOLOv10产品化了。

另一个需要说明的是,本文中所述的TensorRT部署并不是基于ONNX中转方法的部署,而是通过基于INetworkDefinition的手动构建网络进行部署,这种方式能过使得我们对网络有一个较为清晰的了解和认识。

登录https://github.com/tecsai/YOLOv10_TensorRT获取网络构建代码。

话不多数,开始吧。

2. YOLOv10的网络组成

本次以YOLOv10m为例,通过阅读"yolov10m.yaml"文件可以看到,整个网络的组成和YOLOv8极为相似,仅在一些局部位置有改动。

从网络结构配置文件yolov10m.yaml可以看出,YOLOv10有一些基本的模块组成,包括Conv、C2f、SCDown、C2fCIB、PSA、Upsample以及v10Detection组成。

其中,

Conv就是基本的卷积模块;

C2f照搬了YOLOv8中的C2f;

SCDown由两个基本的Conv模块组成,通过控制Kernel和Stride,实现特征图的两倍下采样;

C2fCIB就是特殊形式的C2f,其将C2f中的Bottleneck模块换成了CIB模块。

Upsample就是特征图上采样;

PSA是MHSA与FFN配合实现的Transformer结构(QKV自注意力);

3. 基本的Conv模块

Conv模块就是普通的卷积模块,如下所示为训练工程中的PyTorch版本的卷积。

可以看到,参数已经设置的非常全面了,包括输入输出通道(c1, c2),卷积核尺寸(k),stride(s),padding(p),group(g),dilation(d)以及是否使用激活函数(act)。

对应,我们在基于TensorRT的版本中依葫芦画瓢就可以了,参考如下。

这里有一个建议,后期大家可以将激活函数SILU换成LeakyReLU,这也算是在边缘端提速的一个技巧了。

4. C2f与C2fCIB

C2f可以看做是C3的优化,与C3单元相比,每一个Bottleneck的输入Tensor的Channel都只有上一级的0.5倍,因此计算量明显降低。从另一方面讲,梯度流的增加,也能够明显提升收敛速度和收敛效果。如下分别是C3和C2f的网络结构。

(C3)

(C2f)

基于上述图,我们构建的基于INetworkDefinition的C2f的代码如下。

其中,第89行和93行即是对Tensor进行了拆分操作,依次来完成一种类"CSP"的结构。

之后在103行和108行进行一个类"ELAN"的操作,减少计算量,但丰富了梯度流。

在前面我们说过,C2f和C2fCIB实际上是一样的,仅将Bottleneck结构换成了CIB结构,CIB又是Bottleneck的一种演化版本。

5. PSA

PSA模块本质上是引入了QKV机制的自注意力模块,实现了CNN与Transformer的结合。

PSA全称Partial Self-Attention,即将特征图Tensor的部分进行MHSA+FFN,另一部分则执行了Cross Stage,并与MHSA+FFN的输出进行了融合(Concatenate)。

如下是Attention模块和PSA模块。

Attention

PSA

6. 代码参考

完整的代码可登录https://github.com/tecsai/YOLOv10_TensorRT 获取。

相关推荐
Monkey的自我迭代8 分钟前
KNN算法实战:手写数字识别详解
人工智能·opencv·计算机视觉
Sherlock Ma1 小时前
字节跳动GR-3:可泛化、支持长序列复杂操作任务的机器人操作大模型(技术报告解读)
人工智能·计算机视觉·语言模型·机器人·大模型·aigc·具身智能
CoovallyAIHub2 小时前
数据集分享 | 电子元件检测数据集
深度学习·算法·计算机视觉
宇钶宇夕2 小时前
图像处理中级篇 [1]—— 彩色照相机的效果与预处理
人工智能·数码相机·计算机视觉
一起搞IT吧3 小时前
Camera相机人脸识别系列专题分析之十九:MTK ISP6S平台FDNode传递三方FFD到APP流程解析
android·图像处理·人工智能·数码相机·计算机视觉
8Qi83 小时前
深度学习(鱼书)day06--神经网络的学习(后两节)
人工智能·python·深度学习·神经网络
CoovallyAIHub3 小时前
原来工业 AI 异常检测只做了一半?AnomalyNCD 补上了“最关键一环”
深度学习·算法·计算机视觉
图灵学术计算机期刊速推3 小时前
AAAI‘26 | 聚焦人工智能前沿:西工大李学龙教授荣任赞助主席,论文取号逼近三万,精彩不容错过!
人工智能·机器学习·计算机视觉
CoovallyAIHub4 小时前
数据集分享 | PCB缺陷检测与玻璃缺陷实例分割数据集分享
深度学习·算法·计算机视觉
z are4 小时前
PyTorch 生态四件套:从图片、视频到文本、语音的“开箱即用”实践笔记
人工智能·pytorch·笔记·深度学习