Jetson nano上部署自己的Yolov5模型（TensorRT加速）onnx模型转engine文件

Jetson nano上部署自己的Yolov5模型（TensorRT加速）onnx模型转engine文件

背景

在主机上训练自己的Yolov5模型，转为TensorRT模型并部署到Jetson Nano上，用DeepStream运行。

硬件环境：

RTX 2080TI主机

Jetson Nano 4G B01

软件环境：

Jetson Nano:

Ubuntu 18.04

Jetpack 4.5.1

DeepStream 5.1

主机：

Ubuntu 18.04

CUDA 10.2

yolov5 5.0

训练模型（主机上）

yolov5项目链接https://github.com/ultralytics/yolov5

克隆yolov5官方的代码

git clone https://github.com/ultralytics/yolov5.git

官方训练教程详见https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data

预备环境

准备python3.8以上环境，可用conda创建一个虚拟环境，安装yolov5项目下yolov5/requirements.txt里的依赖

pip install -r requirements.txt

pytorch建议按pytorch官网教程的方式安装PyTorch

如CUDA 10.2 conda环境下安装

conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch

conda安装太慢可用mamba

准备数据集

手上有一个用labelImg（https://github.com/tzutalin/labelImg）打好标签的voc格式数据集，数据集里只有两种类型，鸭子duck和马桶抽sucker。（可自行用其他方式打voc格式数据集或用labelImg直接打yolo格式数据集）

用下面代码将voc格式转成yolo格式数据集，生成images文件夹（存放所有图片），labels文件夹（存放打好的标签），test.txt（测试集），train.txt（训练集），val.txt（验证集）

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
sets = ['train', 'test','val']
classes = ["duck", "sucker"]
 
def convert(size, box):
    dw = 1. / size[0]