使用tensorrt在jetson nx中部署yolov5等目标检测模型（超详细）

目录

• 0.前言
• 1.jetson nx 刷机
• 2.下载必要软件
• 3.训练yolov5.5
• 4.下载tensorrtx并转换模型
• 4.5.int8量化模型
• 5.cmake部署
• 6.总结

---

0.前言

Jetson Xavier NX是nvidia开发的体型超小（只有70mm*45mm）的AI超级计算机，适用于嵌入式系统和边缘系统。
NVIDIA® Jetson Xavier™ NX 是外形小巧的模组系统 (SOM)，可为边缘系统提供超级计算机性能。借助高达 21 TOPS 的加速计算能力，它可以并行运行现代神经网络并处理来自多个高分辨率传感器的数据，这是完整 AI 系统的要求。Jetson Xavier NX 是生产就绪型产品，可支持所有热门 AI 框架。

而jetson nx自带推理加速工具tensorrt，可以很方便的部署深度学习模型。

1.jetson nx 刷机

• 首先参考这个链接
• 如果你要刷18.04的版本需要使用18.04及以前的主机并选择jetpeck4.x的版本，如果要刷20.04的版本则不限制主机版本并选择jetpeck5.x的版本。

要注意点几个点:第一次刷机只刷了系统到sdk里，然后使用装上固态将内容拷贝到固态中，然后设置开机从硬盘启动。之后刷环境时虽然这个博客里没有勾选系统，但博主遇到的情况是单独装cuda等环境是装不上的，需要重新勾选系统，所以系统烤完插上屏幕进入系统后需要重新设置一遍硬盘，然后再紧接着烤环境。

• 所以我推荐一开始直接选择系统和环境而不是只选择系统，sdkmanger刷完系统会弹出一个窗口，在这个弹窗时去插屏幕把系统烤到固态并设置从硬盘启动，然后再回到sdkmanger的弹窗，设置选项后继续烤cuda等环境。

2.下载必要软件

• 包括且不限于jtop，cmake，vscode。

• vscode可以在官网下载arm64的deb包，注意:要下老版本的vscode，新版本会出现依赖项问题

• jtop和cmake的下载参考
  其他软件在网上搜jetson nx安装***就能找到参考。

3.训练yolov5模型并转换

• 首先要注意的是我们应该选择的模型版本，tensorrtx目前支持的版本可自行在tensorrtx中查看
• 博主使用的是yolov5.5.0模型，当然你们也可以使用tensorrtx中支持的其他模型。
• 这里是tensorrtx的链接，你可以打开自行选择需要用的模型。
• 这里是yolov5.50模型的链接
• 自行训练自己的模型并导出.pt文件，yolov5训练过程可以参考这个博客

4.tensorrt部署

用tensorrt部署首先要下载tensorrt，但如果你使用了sdkmanger将jetson nx刷机了，你可以通过jtop查看是否有tensorrt（大概率cuda等一套都刷好了），如果没有刷好，移步1.jetson nx刷机

• 现在我默认你已经完成了jetson nx的刷机，电脑有了jtop vscode，cmake等软件，并且已经训练出来自己的.pt模型。

4.1.pytorch model–>TensorRT engine

Pytorch训练出来的模型为XXX.pt或XXX.pth，我们以Yolov5举例，他训练出来的模型为.pt文件， 我们要将其通过tensorrt部署，就需要将pt文件转换为tensorrt支持的engine格式，转换方式有两种，可以用onnx也可以直接用github上的大佬开源的项目tensorrtx，此项目集模型转换与部署为一体，可以直接调用。

• 首先，下载tensorrtx，将 /tensorrtx-yolov5-v5.0/yolov5 中的gen_wts.py文件拷贝到yolov5-5.0目录中，并将 best.pt模型文件也拷贝到此目录
  

运行

python gen_wts.py -w best.pt

生成需要用的best.wts文件

• 然后进入 /tensorrtx-yolov5-v5.0/yolov5修改yololayer.h文件中的CLASS_NUM为你训练的模型的标签类别数
• 然后编译

mkdir build
cd build
cmake ..
make

• 生成可执行文件yolov5然后将best.wts拷贝到build文件夹中，并执行

./yolov5 -s best.wts best.engine s

• 注意这个命令最后一个字母是你的模型类型，这里博主的类型是s所以这样写，具体能用哪些可以在tensorrtx中yolov5的readme中查看。
  生成best.engine文件，这便是我们tensort部署所需要的模型文件啦

4.5.int8量化模型

了解什么是int8量化可以看这个博客。
简而言之就是让模型跑的更快，但精度会有微弱下降，实测推理速度会提高一倍左右。

具体操作步骤可以去看tensorrtx中yolov5的readme，博主在下面进行了简述。

• 首先下载coco_calib,将数据集解压。
• 然后进入 /tensorrtx-yolov5-v5.0/yolov5修改yololayer.h文件中的CLASS_NUM为你训练的模型的标签类别数，然后要修改src/config.h里前几行里的USE_FP16为USE_INT8
• 然后编译

mkdir build
cd build
cmake ..
make

• 生成可执行文件yolov5然后将best.wts和刚才下载的coco量化数据集拷贝到build文件夹中，并执行

./yolov5 -s best.wts best_int8.engine s

• 注意这个命令最后一个字母是你的模型类型，这里博主的类型是s所以这样写，具体能用哪些可以在tensorrtx中yolov5的readme中查看。
  生成best_int8.engine文件，这便是int8量化后获得的部署文件啦。
• 部署时记得将生成的int8文件和量化数据集放在项目工程的Yolo文件夹中。

5.cmake部署

至于cmake部署可以参考别人的demo工程Mask_Detector（检测人是否佩戴口罩，若佩戴则输出绿框，未佩戴则输出红框），只需要修改engine_name便可以实现不同的目标检测，记得修改yololar.h中的class_num，其他需要修改的可以打开yolo中的yolov5.cpp自行修改，包括但不限于输出置信度等。

• 注意，在cmake时可能会报错，说找不到libcudart.so,我们可以在usr/local/中搜索一下，然后拷贝到寻找的那个路径中。

6.总结

博主在jetson nx上尝试使用c++onnxruntime未果后，转而尝试使用tensorrt部署，结果发现效果还行。最大处理速度可以到120帧。博主写下这个完整流程留给之后要在jetson 上部署yolo模型的朋友一个参考，也方便之后自己忘记之后查看。

我是一只咸鱼_，文章对你有帮助的话就点赞收藏吧！