02-Jetson Xavier NX 训练自己的yolov4-tiny模型并利用TensorRT完成部署(Python API)

记录原因:

  • 因前面完成了yolov4-tiny源码的部署,可以跑到50FPS。但那个跑的是别人的模型,可以识别80类,但是在特定的情况下比如识别螺丝螺帽,只需要两种,在训练的时候直接把种类改掉就好了,但是在Jetson Xavier NX上应该怎么利用TensorRT进行加速的(未进行加速的推大概在5.XFPS,根本无法满足实际应用的需要),所以本文主要是解决“训练完成后得到pth如何在Jetson板子上利用TensorRT进行加速”
  • 我们在部署的时候一般的模型格式转化为:Pytorch -> onnx -> trt

一、pytorch转onnx

  • pytorch转onnx的部分,参考下面网址,博主在测试predict有pth转onnx的函数,咱们直接调用即可:
    bubbliiiing/yolov4-tiny-pytorch
  • 如果要查看转换完成的onnx网络,可使用下面的工具导入查看netron -onnx网络查看工具
  • 经函数测试,我们可以把onnx拷贝到jetson板子上。

二、onnx转trt

  • 这部分的话有两种方式:
    • 通过官方提供给我们的工具trtexec工具,此处转换是可以成功,但是我没有找到可以运行推理trt的代码,参考那些网上的一些trt推理的代码经常会报错;
    • 第二种方式当然就是大佬们的封装了,我参考了下面博主的代码:
      TensorRT部署 Jetson可用,这个工程文件里面提供了weight --> onnx,onnx–>trt的两种文件,同时也还有trt的推理程序。
因为大佬们已经做了很多封装工具,像我这种小菜鸡就直接用了。
  • 1.我们需要根据他的程序把onnx文件转为trt文件,但是这个时候是需要一个cfg配置文件参考的,所以我们需要根据我上一个教程中,把yolov4-tiny的cfg文件下载下来,然后还要修改 ,把种类classes换成我们需要识别的种类数,比如我在实验中检测的是螺丝螺帽,所以是两种,把classes改成2(在最后的地方,需要修改两处)

  • 2.运行下面指令进行trt模型的转换

python3 onnx_to_tensorrt.py -m models --verbose

其中,models表示的onnx的名称,省略了.onnx, --verbose表示我想看到具体的一些日志信息(如果没有的话,可以就会感觉等很久,但是实际也是运行的)

  • 3.需要种类名称,因为我们这个是只有2类模型,而原生的yolov4-tiny有80类,所以需要我们修改下名称:具体文件的名字路径为: tensorrt_demos/utils/yolo_classes
    [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-BMjpUhN9-1679477632087)(https://note.youdao.com/yws/res/b/WEBRESOURCE76c29df88bed2bec536163ace184482b)]

在这里插入图片描述

  • 4.检查一下可用摄像头,输入以下指令:
ls -l /dev/video*
  • 5.运行trt_yolo.py推理文件即可实现推理,5表示摄像头标号:
python3 trt_yolo.py --usb 5 -m models
  • 6.查看效果(帧率可到达55FPS左右):
    在这里插入图片描述
  • 2
    点赞
  • 6
    收藏
    觉得还不错? 一键收藏
  • 2
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值