目录

前言

开始YOLOv8的模型部署，目的是为大家推荐一个全新的tensorrt仓库https://github.com/shouxieai/infer，大家可以查看我之前的Jetson嵌入式系列模型部署教程，很多细节这里就不再赘述了。考虑到nano的算力，这里采用yolov8n.pt模型，本文主要分享yolov8模型训练和jetson nano部署yolov8两方面的内容。若有问题欢迎各位看官批评指正!!!😄

一、YOLOv8模型训练

先来欣赏下YOLOv8的网络结构😅

在这里插入图片描述

yolov8的代码风格和yolov5相差较大，训练流程也稍有差异。博主主要参考魔鬼面具的YOLOv8最强操作.

1. 项目的克隆和必要的环境依赖

1.1 项目的克隆

yolov8的代码是开源的可直接从github官网上下载，源码下载地址是https://github.com/ultralytics/ultralytics，由于yolov8刚发布不久一个固定版本都没有，故只能采用主分支进行模型的训练和部署工作(PS:由于代码更新频繁，可能大家会遇到不同的bug)。Linux下代码克隆指令如下

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

也可手动点击下载，点击右上角的Code按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击here[pwd:yolo]下载博主准备好的代码(注意代码下载于2023/3/14日，若有改动请参考最新)

在这里插入图片描述

1.2 项目代码结构整体介绍

将下载后的yolov8的代码解压，其代码目录如下图

在这里插入图片描述
只需要关注ultralytics这个文件夹的内容即可，下面对这个文件夹的整体目录做一个介绍

  • |-assets：存放测试图片
  • |-datasets：存放一些超参数的配置文件以及配置训练集和验证集路径的coco.yaml文件，如果需要修改自己的数据集，那么需要修改其中的yaml文件
  • |-hub：pytorch扩展模型
  • |-models：存放不同模型的yaml文件，包括v3、v5和v8
  • |-nn：存放yolov8整体网络模型搭建的py文件
  • |-tracker：存放yolov8目标跟踪的py文件
  • |-yolo：存放yolov8模型预测、训练、导出的py文件
    • cfg：存放yolov8的配置文件，包括训练时的参数指定如epoch、batch等以及超参数设置，所有的相关配置都可通过这个文件设置(重点关注cfg/default.yaml文件)
    • data：存放数据加载的py文件
    • engine：存放模型导出的py文件
    • utils：存放工具类函数，包括loss、metrics、plots函数等
    • v8：存放yolov8分类、检测、分割等不同任务的预测、训练以及验证的py文件(重点关注)
1.3 环境安装

关于深度学习的环境安装可参考炮哥的利用Anaconda安装pytorch和paddle深度学习环境+pycharm安装—免额外安装CUDA和cudnn(适合小白的保姆级教学)，这里不再赘述。如果之前配置过yolov5的环境，yolov8可直接使用。

2. 数据集和预训练权重的准备

2.1 数据集

这里采用的数据集是口罩识别，来源于B站UP主HamlinZheng的口罩识别数据集，这里给出下载链接Baidu Drive[password:yolo]，博主将原数据集整合了下，方便后续的训练，解压后整个数据集目录结构如下

VOCdevkit
   └─VOC2007
       ├─Annotations
       └─JPEGImages

其中JPEGImages中存放的图像文件，Annotations存放的是对应的XML标签文件。关于标签的制作可参考B站UP主霹雳吧啦Wz的PASCAL VOC2012数据集讲解与制作自己的数据集，由于labelimg标注的是VOC格式标签的XML文件，需要转化为YOLO格式标签的txt文件，关于转换的代码可参考炮哥的目标检测—数据集格式转化及训练集和验证集划分，下面给出VOC格式转YOLO格式的代码：

# voc2yolo.py
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random
from shutil import copyfile

# 1. 修改为自制数据集需检测的类别数
classes = ["have_mask", "no_mask"]

# 2. 训练集和验证集的比例
TRAIN_RATIO = 80


def clear_hidden_files(path):
    dir_list = os.listdir(path)
    for i in dir_list:
        abspath = os.path.join(os.path.abspath(path), i)
        if os.path.isfile(abspath):
            if i.startswith("._"):
                os.remove(abspath)
        else:
            clear_hidden_files(abspath)


def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def convert_annotation(image_id):
    in_file = open('VOCdevkit/VOC2007/Annotations/%s.xml' % image_id)
    out_file = open('VOCdevkit/VOC2007/YOLOLabels/%s.txt' % image_id, 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    in_file.close()
    out_file.close()


wd = os.getcwd()
wd = os.getcwd()
data_base_dir = os.path.join(wd, "VOCdevkit/")
if not os.path.isdir(data_base_dir):
    os.mkdir(data_base_dir)
work_sapce_dir = os.path.join(data_base_dir, "VOC2007/")
if not os.path.isdir(work_sapce_dir):
    os.mkdir(work_sapce_dir)
annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
    os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
if not os.path.isdir(image_dir):
    os.mkdir(image_dir)
clear_hidden_files(image_dir)
yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
if not os.path.isdir(yolo_labels_dir):
    os.mkdir(yolo_labels_dir)
clear_hidden_files(yolo_labels_dir)
yolov5_images_dir = os.path.join(data_base_dir, "images/")
if not os.path.isdir(yolov5_images_dir):
    os.mkdir(yolov5_images_dir)
clear_hidden_files(yolov5_images_dir)
yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
if not os.path.isdir(yolov5_labels_dir):
    os.mkdir(yolov5_labels_dir)
clear_hidden_files(yolov5_labels_dir)
yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
if not os.path.isdir(yolov5_images_train_dir):
    os.mkdir(yolov5_images_train_dir)
clear_hidden_files(yolov5_images_train_dir)
yolov5_images_test_dir = os.path.join(yolov5_images_dir, "val/")
if not os.path.isdir(yolov5_images_test_dir):
    os.mkdir(yolov5_images_test_dir)
clear_hidden_files(yolov5_images_test_dir)
yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
if not os.path.isdir(yolov5_labels_train_dir):
    os.mkdir(yolov5_labels_train_dir)
clear_hidden_files(yolov5_labels_train_dir)
yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
if not os.path.isdir(yolov5_labels_test_dir):
    os.mkdir(yolov5_labels_test_dir)
clear_hidden_files(yolov5_labels_test_dir)

train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
train_file.close()
test_file.close()
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'a')
list_imgs = os.listdir(image_dir)  # list image files
prob = random.randint(1, 100)
print("Probability: %d" % prob)
for i in range(0, len(list_imgs)):
    path = os.path.join(image_dir, list_imgs[i])
    if os.path.isfile(path):
        image_path = image_dir + list_imgs[i]
        voc_path = list_imgs[i]
        (nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
        (voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
        annotation_name = nameWithoutExtention + '.xml'
        annotation_path = os.path.join(annotation_dir, annotation_name)
        label_name = nameWithoutExtention + '.txt'
        label_path = os.path.join(yolo_labels_dir, label_name)
    prob = random.randint(1, 100)
    print("Probability: %d" % prob)
    if (prob < TRAIN_RATIO):  # train dataset
        if os.path.exists(annotation_path):
            train_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention)  # convert label
            copyfile(image_path, yolov5_images_train_dir + voc_path)
            copyfile(label_path, yolov5_labels_train_dir + label_name)
    else:  # test dataset
        if os.path.exists(annotation_path):
            test_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention)  # convert label
            copyfile(image_path, yolov5_images_test_dir + voc_path)
            copyfile(label_path, yolov5_labels_test_dir + label_name)
train_file.close()
test_file.close()

代码总共需要修改两处

  • 第11行，修改要检测的类别名称
  • 第14行，修改训练集和验证集的划分比例

整个目录结构如下，注意voc2yolo.py与VOCdevkit处于同一级目录

VOCdevkit
   └─VOC2007
       ├─Annotations
       └─JPEGImages
voc2yolo.py

注：目录结构一定要与博主的一致，因为程序已经将对应目录写死。

运行voc2yolo.py代码之后得到如下结果
在这里插入图片描述
可以看到目录下有一些新的文件生成，首先VOCdevkit文件夹下分别生成了images和labels文件夹，分别存放着图像和对应的yolo格式的标签文件，每个文件夹下分别包含train和val两个子文件夹，代表各自对应的训练集和验证集。VOC2007文件夹下生成了YOLOLabels文件夹，存放着对应yolo格式的标签文件。然后还生成了yolov5_train.txt以及yolov5_val.txt两个txt文件，存放着训练集和验证集图片的完整路径。yolov8的训练只需要VOCdevkit目录下的images和labels两个文件夹，其它均不需要，故最终的目录结构如下

VOCdevkit
    ├─images
    │  ├─train
    │  └─val
    └─labels
        ├─train
        └─val

至此，数据集的准备工作完毕。

2.2 预训练权重准备

yolov8预训练权重可以通过here找到下图的Models，然后点击不同大小的权重即可下载，博主也提供下载好的几个预训练权重Baidu Drive[pwd:yolo]，注意这是yolov8没固定版本的预训练权重，若后续有版本更新，记得替换。本次训练自己的数据集使用的预训练权重为yolov8n.pt

在这里插入图片描述

3. 训练模型

将准备好的数据集文件夹即VOCdevkit复制到yolov8项目环境中的ultralytics-main/ultralytics/datasets文件夹下，将准备好的预训练权重yolov8n.pt复制到yolov8项目环境中的ultralytics-main/ultralytics/yolo/v8/detect文件夹下。训练目标检测模型主要修改模型配置文件ultralytics-main/ultralytics/yolo/cfg/default.yaml以及数据配置文件ultralytics-main/ultralytics/datasets/coco128.yaml，项目结构如下所示

在这里插入图片描述

3.1 修改数据配置文件

修改ultralytics/datasets目录下相应的yaml文件，找到该目录下的coco128.yaml文件，主要修改如下：

  • 1.修改第11行的数据集的路径

  • 2.修改第12行的train的路径

  • 3.修改第13行的val路径

  • 4.修改第17行的需要检测的类别名称

  • 5.注释最后一行，取消coco数据集的下载

在这里插入图片描述

3.2 修改模型配置文件

修改ultralytics/yolo/cfg目录下对应的yaml文件，找到该目录下的default.yaml文件，主要修改如下：

  • 1.修改第8行的模型文件，指定yolov8n.pt

  • 2.修改第9行的数据文件路径，即3.1节的coco128.yaml

在这里插入图片描述

还有其它参数博主并未修改，如epoches训练轮数、batch批处理量等等，大家一定要根据自己的实际情况(如显卡算力等)指定不同的参数，如果你之前训练过yolov5，那我相信这对你来说应该是小case😄

3.3 训练模型

按照上面的步骤完成修改后可以找到yolo/v8/detect/train.py文件，点击运行即可开始训练。

博主训练的模型为yolov8n.pt且使用的是单个GPU进行训练，显卡为RTX3060，操作系统为Windows10，pytorch版本为1.12.1，训练时长大概1.6个小时。

在这里插入图片描述
训练完成后的模型权重保存在runs/detect/train/weights文件夹下，后续使用best.pt进行后续模型部署工作即可，这里提供博主训练好的权重文件下载链接Baidu Driver[pwd:yolo]

3.4 推理测试

利用项目中的yolo/v8/detect/predict.py文件进行测试，将需要推理的图片放入ultralytics/assets文件夹下，将刚训练的口罩识别权重文件放入yolo/v8/detect文件夹下，最后在predict.py文件第86行指定刚训练的模型权重文件best.pt如下所示，点击运行predict.py即可执行测试

model = 'best.pt'   # predict.py第86行修改

推理完成后会在detect/runs文件夹下生成一个detect/train目录，推理的结果保存在此目录下，推理结果如下所示

在这里插入图片描述
在这里插入图片描述

至此，yolov8模型训练已经完毕，下面开始jetson nano上的部署工作。

二、YOLOv8模型部署

Jetson nano上yolov8模型部署流程和yolov5、yolov7相差较大，主要换了一个部署框架，大家对jetson模型部署感兴趣的可以参考我之前发的Jetson嵌入式系列模型部署文章，这次带大家了解下一个全新的tensorrt封装，部署使用到的Github仓库是infer。该仓库通过trtexec工具完成模型的构建工作。对模型部署有疑问的可以参考Jetson嵌入式系列模型部署-1，想了解通过TensorRT的Layer API一层层完成模型的搭建工作可参考Jetson嵌入式系列模型部署-2，想了解通过TensorRT的ONNX parser解析ONNX文件来完成模型的搭建工作可参考Jetson嵌入式模型部署-3、Jetson nano部署YOLOv7。本文主要针对infer项目中的yolov8完成嵌入式模型部署，本文参考自infer的README.md，具体操作流程作者描述非常详细，这里再简单过一遍，本次训练的模型使用yolov8n，类别数为2，为口罩识别😷。

1. 源码下载

infer的代码可以直接从github官网上下载，源码下载地址是https://github.com/shouxieai/infer，由于infer部署框架刚发布不久一个固定的版本都没有，故只能采用主分支进行yolov8的部署工作(PS:由于代码更新频繁，可能大家会遇到不同的bug)。Linux下代码克隆指令如下

$ git clone https://github.com/shouxieai/infer.git

也可手动点击下载，点击右上角的Code按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击here[pwd:yolo]下载博主准备好的源代码(注意代码下载于2023/3/15日，若有改动请参考最新)

2. 环境配置

需要使用的软件环境有TensorRT、CUDA、CUDNN、OpenCV。所有软件环境在JetPack镜像中已经安装完成，只需要添加下trtexec工具的环境变量即可。博主使用的jetpack版本为JetPack4.6.1(PS:关于jetson nano刷机就不再赘述了，需要各位看官自行配置好相关环境😄，外网访问较慢，这里提供Jetson nano的JetPack镜像下载链接Baidu Drive[password:nano]【更新完毕!!!】几个(PS:提供4.6和4.6.1两个版本，注意4GB和2GB的区别，不要刷错了)，关于Jetson Nano 2GB和4GB的区别可参考链接Jetson NANO是什么？如何选？。(吐槽下这玩意上传忒慢了，超级会员不顶用呀，终于上传完了，折磨!!!)

博主的Jetpack对应的软件环境如下

在这里插入图片描述

2.1 trtexec环境变量设置

trtexec环境变量的添加主要参考这里，主要包含以下几步

1.打开bashrc文件

vim ~/.bashrc

2.按i进入输入模式，在最后一行添加如下语句

export PATH=/usr/src/tensorrt/bin:$PATH

3.按下esc，输入:wq!保存退出即可，最后刷新下环境变量

source ~/.bashrc

3. ONNX导出

关于静态batch和动态batch有以下几点说明，更多细节请查看视频

静态batch

  • 导出的onnx指定所有维度均为明确的数字，是静态shape模型
  • 在推理的时候，它永远都是同样的batch推理，即使你目前只有一个图推理，它也需要n个batch的耗时
  • 适用于大部分场景，整个代码逻辑非常简单

动态batch

  • 导出的时候指定特定维度为dynamic，也就是不确定状态
  • 模型推理时才决定所需推理的batch大小，耗时最优，但onnx复杂度提高了
  • 适用于如server有大量不均匀的请求时的场景
3.1 静态batch导出
3.1.1 Transpose节点的添加

将训练好的口罩识别权重best.pt放在ultralytics-main主目录下，新建导出文件export.py，内容如下，执行完成后将会在当前目录生成导出的best.onnx模型

from ultralytics import YOLO

model = YOLO("best.pt")

success = model.export(format="onnx", batch=1)

模型需要完成修改才能正确被infer框架使用，正常模型导出的输出为[1,6,8400]，其中1代表batch，6分别代表cx,cy,w,h,以及have_mask、no_mask两个类别分数，8400代表框的个数。首先infer框架的输出只支持[1,8400,6]这种形式的输出，因此我们需要在原始onnx的输出之前添加一个Transpose节点，infer仓库workspace/v8trans.py就是帮我们做这么一件事情，v8trans.py具体内容如下：

# v8trans.py
import onnx
import onnx.helper as helper
import sys
import os

def main():

    if len(sys.argv) < 2:
        print("Usage:\n python v8trans.py yolov8n.onnx")
        return 1

    file = sys.argv[1]
    if not os.path.exists(file):
        print(f"Not exist path: {file}")
        return 1

    prefix, suffix = os.path.splitext(file)
    dst = prefix + ".transd" + suffix

    model = onnx.load(file)
    node  = model.graph.node[-1]

    old_output = node.output[0]
    node.output[0] = "pre_transpose"

    for specout in model.graph.output:
        if specout.name == old_output:
            shape0 = specout.type.tensor_type.shape.dim[0]
            shape1 = specout.type.tensor_type.shape.dim[1]
            shape2 = specout.type.tensor_type.shape.dim[2]
            new_out = helper.make_tensor_value_info(
                specout.name,
                specout.type.tensor_type.elem_type,
                [0, 0, 0]
            )
            new_out.type.tensor_type.shape.dim[0].CopyFrom(shape0)
            new_out.type.tensor_type.shape.dim[2].CopyFrom(shape1)
            new_out.type.tensor_type.shape.dim[1].CopyFrom(shape2)
            specout.CopyFrom(new_out)

    model.graph.node.append(
        helper.make_node("Transpose", ["pre_transpose"], [old_output], perm=[0, 2, 1])
    )

    print(f"Model save to {dst}")
    onnx.save(model, dst)
    return 0

if __name__ == "__main__":
    sys.exit(main())

在命令行终端输入如下指令即可添加Transpose节点，执行完成之后在当前目录下生成best.transd.onnx模型，该模型添加了Transpose节点。

python v8trans.py best.onnx

下图对比了原始的best.onnx和best.transd.onnx之间的区别，从图中可以看出转换后的onnx模型在输出之前多了一个Transpose节点，且输出的1，2维度进行了交换，符合infer框架。

在这里插入图片描述

3.1.2 Resize节点解析的问题

先剧透下，当使用trtexec工具构建engine时会出现错误，我们一并解决，到时候可直接生成engine，错误信息如下图所示，大概意思就是说Resize_120这个节点的scales没有初始化(应该是这样理解的吧🤔)

在这里插入图片描述

我们先通过Netron工具打开best.transd.onnx模型查看下Resize_120这个节点的相关信息，在找找其它的Resize节点对比看看，如下图所示，左边是Resize_103节点的相关信息，右边是Resize_120节点的相关信息，可以看到其对应的Scales确实存在区别，Resize_120节点Scales没有initializer，没有明确的值。

在这里插入图片描述

下面来看解决方案，提供两种

解决方案1-onnxsim(推荐)

onnxoptimizer、onnxsim被誉为onnx的优化利器，其中onnxsim可以优化常量，onnxoptimizer可以对节点进行压缩，参考自onnxoptimizer、onnxsim使用记录。新建一个v8onnxsim.py文件，用于优化onnx文件，具体内容如下：

import onnx
from onnxsim import simplify

onnx_model = onnx.load("best.transd.onnx")
model_simp, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be Validated"
onnx.save(model_simp, "best.transd.sim.onnx")

运行后会在当前文件夹生成一个best.transd.sim.onnx模型，现在可以查看对应的Resize_120节点发生了改变

在这里插入图片描述

解决方案2-自己动手

俗话说自己动手丰衣足食，看了onnxsim优化的结果，突然有种自己动手操作的欲望(PS:大家看个热闹即可😙)。主要步骤如下：

  • 1.加载ONNX模型
  • 2.遍历模型的所有节点，找到要操作的Resize_120
  • 3.修改其输入从onnx::Resize_432变成onnx::Resize_431
  • 4.删除Identit_2节点
  • 5.保存模型

在当前目录下新建v8Resize.py文件，其内容如下：


import onnx
import onnx.helper as helper
import sys
import os

def main(file_path):

    # 1.加载ONNX模型文件
    model = onnx.load(file_path)
    resize_node = None

    # 2.遍历模型的所有节点，查找resize节点
    # Resize_103 Resize_120
    for node in model.graph.node:
        if node.op_type == "Resize" and node.name == "Resize_120":
            resize_node = node
            break

    if resize_node is None:
        raise ValueError("Resize node not found in the model.")

    # 3.修改输入
    resize_node.input[2] = "onnx::Resize_431"

    # 4.删除Identity_2节点
    for i in range(len(model.graph.node)):
        if model.graph.node[i].name == "Identity_2":
            del model.graph.node[i]
            break

    # 5.保存新模型
    onnx.save(model, "best.transd.my.onnx")
    print("new model save done.")

if __name__ == "__main__":
    file_path = "./best.transd.onnx"
    sys.exit(main(file_path))

运行该文件后会在当前目录下生成best.transd.my.onnx模型，现在可以查看对应的Resize_120节点发生了改变
在这里插入图片描述

至此，静态batch模型导出告一段落，后续通过trtexec工具可正常生成engine文件了，主要解决两个方面的问题，一是输出必须加一个Transpose转换，二是解决Resize节点无法解析问题。

3.2 动态batch导出
3.2.1 exporter.py修改

动态batch的导出需要修改ultralytics/yolo/engine/exporter.py文件第303行的内容，确保导出的batch指定为动态，其它维度不指定，修改的内容如下：

# -----exporter.py第303行-----
# dynamic = self.args.dynamic
# if dynamic:
#     dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}}  # shape(1,3,640,640)
#     if isinstance(self.model, SegmentationModel):
#         dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
#         dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
#     elif isinstance(self.model, DetectionModel):
#         dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)

# -----修改为如下代码-----
dynamic = self.args.dynamic
dynamic = True
if dynamic:
    dynamic = {'images': {0: 'batch'}}  # shape(1,3,640,640)
    if isinstance(self.model, SegmentationModel):
        dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
        dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
    elif isinstance(self.model, DetectionModel):
        dynamic['output0'] = {0: 'batch'}  # shape(1,25200,85)

然后将训练好的口罩识别权重best.pt放在ultralytics-main目录下，新建导出文件export.py，内容如下，执行完成后将会在当前目录下生成导出的best.onnx模型

from ultralytics import YOLO

model = YOLO("best.pt")

success = model.export(format='onnx', batch=1)

导出的onnx模型如下：

在这里插入图片描述

3.2.2 Transpose节点添加

由于infer框架只支持[1,8400,6]这种形式的输出，因此我们需要在原始onnx之前添加一个Transpose节点，infer仓库workspace/v8trans.py就是帮我们做这么一件事情，v8trans.py具体内容可参考3.1.1小节，这里不再赘述。

3.2.3 Resize节点解析的问题

同样的当使用trtexec工具构建engine时，会出现Resize节点解析错误

在这里插入图片描述

废话少说，直接上解决方案

解决方案1-onnxsim

import onnx
from onnxsim import simplify

onnx_model = onnx.load("best.onnx")
model_simp, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be Validated"
onnx.save(model_simp, "best.sim.onnx")

解决方案2-自己动手


import onnx
import onnx.helper as helper
import sys
import os

def main(file_path):

    # 1.加载ONNX模型文件
    model = onnx.load(file_path)
    resize_node = None

    # 2.遍历模型的所有节点，查找resize节点
    # Resize_106 Resize_123
    for node in model.graph.node:
        if node.op_type == "Resize" and node.name == "Resize_123":
            resize_node = node
            break

    if resize_node is None:
        raise ValueError("Resize node not found in the model.")

    # 3.修改输入
    resize_node.input[2] = "onnx::Resize_437"

    # 4.删除Identity_5节点
    for i in range(len(model.graph.node)):
        if model.graph.node[i].name == "Identity_5":
            del model.graph.node[i]
            break

    # 5.保存新模型
    onnx.save(model, "best.my.onnx")
    print("new model save done.")

if __name__ == "__main__":
    file_path = "./best.onnx"
    sys.exit(main(file_path))

至此，模型导出已经完毕，后续通过导出的模型完成在jetson nano上的部署工作，导出的模型文件可点击here[pwd:yolo]下载

4. 运行

4.1 engine生成

与tensorRT_Pro模型构建方式不同，infer框架直接通过trtexec工具生成engine，infer框架拥有一个全新的tensorrt封装，可轻易继承各类任务，相比于tensorRT_Pro优点如下：

  • 轻易实现各类任务的生产者和消费者模型，并进行高性能推理
  • 没有复杂的封装，彻底解开耦合!
  • 参考自如何高效使用TensorRT

将第3节导出的ONNX模型放入到infer/workspace文件夹下，然后在jetson nano终端执行如下指令(以导出的静态batch模型为例)

trtexec --onnx=workspace/best.transd.sim.onnx --saveEngine=workspace/best.transd.sim.engine

图解如下所示(PS:模型构架需要一段时间，请耐心等待)

在这里插入图片描述

模型构建完成后如下图所示，engine拿到手后就可以开干了👨‍🏭

在这里插入图片描述
注：导出动态batch模型执行的指令与静态batch不同！！！，具体可参考infer/workspace/build.sh文件中的内容，指令如下：

trtexec --onnx=best.transd.sim.onnx --minShapes=images:1x3x640x640 --maxShapes=images:16x3x640x640 --optShapes=images:1x3x640x640 --saveEngine=best.transd.sim.engine
4.2 源码修改

yolo模型的推理代码主要在src/main.cpp文件中，需要推理的图片放在workspace/inference文件夹中，源码修改较简单主要有以下几点：

  • 1.main.cpp 134，135行注释，只进行单张图片的推理

  • 2.main.cpp 104行 修改加载的模型为best.transd.sim.engine且类型为V8

  • 3.main.cpp 10行 新增mylabels数组，添加自训练模型的类别名称

  • 4.mian.cpp 115行 cocolabels修改为mylabels

具体修改如下

int main() {
  // perf();					//修改1 134 135行注释
  // batch_inference();			
  single_inference();
  return 0;
}

auto yolo = yolo::load("best.transd.sim.engine", yolo::Type::V8);	//  修改2

static const char *mylabels[] = {"have_mask", "no_mask"};	// 修改3 新增mylabels数组

auto name = mylabels[obj.class_label]		// 修改4 cocolabels修改为mylabels
4.3 编译运行

编译用到的Makefile文件需要修改，修改后的Makefile文件如下，详细的Makefile文件的分析可查看Mkaefile实战

cc        := g++
nvcc      = /usr/local/cuda-10.2/bin/nvcc

cpp_srcs  := $(shell find src -name "*.cpp")
cpp_objs  := $(cpp_srcs:.cpp=.cpp.o)
cpp_objs  := $(cpp_objs:src/%=objs/%)
cpp_mk	  := $(cpp_objs:.cpp.o=.cpp.mk)

cu_srcs	  := $(shell find src -name "*.cu")
cu_objs   := $(cu_srcs:.cu=.cu.o)
cu_objs	  := $(cu_objs:src/%=objs/%)
cu_mk	  := $(cu_objs:.cu.o=.cu.mk)

include_paths := src        \
			/usr/include/opencv4 \
			/usr/include/aarch64-linux-gnu \
			/usr/local/cuda-10.2/include

library_paths := /usr/lib/aarch64-linux-gnu \
			/usr/local/cuda-10.2/lib64

link_librarys := opencv_core opencv_highgui opencv_imgproc opencv_videoio opencv_imgcodecs \
			nvinfer nvinfer_plugin nvonnxparser \
			cuda cublas cudart cudnn \
			stdc++ dl

empty		  :=
export_path   := $(subst $(empty) $(empty),:,$(library_paths))

run_paths     := $(foreach item,$(library_paths),-Wl,-rpath=$(item))
include_paths := $(foreach item,$(include_paths),-I$(item))
library_paths := $(foreach item,$(library_paths),-L$(item))
link_librarys := $(foreach item,$(link_librarys),-l$(item))

cpp_compile_flags := -std=c++11 -fPIC -w -g -pthread -fopenmp -O0
cu_compile_flags  := -std=c++11 -g -w -O0 -Xcompiler "$(cpp_compile_flags)"
link_flags        := -pthread -fopenmp -Wl,-rpath='$$ORIGIN'

cpp_compile_flags += $(include_paths)
cu_compile_flags  += $(include_paths)
link_flags        += $(library_paths) $(link_librarys) $(run_paths)

ifneq ($(MAKECMDGOALS), clean)
-include $(cpp_mk) $(cu_mk)
endif

pro	   := workspace/pro
expath := library_path.txt

library_path.txt : 
	@echo LD_LIBRARY_PATH=$(export_path):"$$"LD_LIBRARY_PATH > $@

workspace/pro : $(cpp_objs) $(cu_objs)
		@echo Link $@
		@mkdir -p $(dir $@)
		@$(cc) $^ -o $@ $(link_flags)

objs/%.cpp.o : src/%.cpp
	@echo Compile CXX $<
	@mkdir -p $(dir $@)
	@$(cc) -c $< -o $@ $(cpp_compile_flags)

objs/%.cu.o : src/%.cu
	@echo Compile CUDA $<
	@mkdir -p $(dir $@)
	@$(nvcc) -c $< -o $@ $(cu_compile_flags)

objs/%.cpp.mk : src/%.cpp
	@echo Compile depends CXX $<
	@mkdir -p $(dir $@)
	@$(cc) -M $< -MF $@ -MT $(@:.cpp.mk=.cpp.o) $(cpp_compile_flags)
	
objs/%.cu.mk : src/%.cu
	@echo Compile depends CUDA $<
	@mkdir -p $(dir $@)
	@$(nvcc) -M $< -MF $@ -MT $(@:.cu.mk=.cu.o) $(cu_compile_flags)

run   : workspace/pro
		  @cd workspace && ./pro

clean :
	@rm -rf objs workspace/pro
	@rm -rf library_path.txt
	@rm -rf workspace/Result.jpg

# 导出符号，使得运行时能够链接上
export LD_LIBRARY_PATH:=$(export_path):$(LD_LIBRARY_PATH)

OK!源码也修改好了，Makefile文件也搞定了，可以编译运行了，直接在终端执行如下指令即可

make run

图解如下所示：

在这里插入图片描述

编译运行后的将在worksapce下生成Result.jpg为推理后的图片，如下所示，可以看到效果还是比较OK的。

在这里插入图片描述

在这里插入图片描述

4.4 拓展-摄像头检测

简单写了一个摄像头检测的demo，主要修改以下几点：

  • 1.main.cpp 新增yolo_video_demo()函数，具体内容参考下面

  • 2.main.cpp 新增调用yolo_video_demo()函数代码，具体内容参考下面

static void yolo_video_demo(const string& engine_file){		// 修改1 新增函数
  auto yolo = yolo::load(engine_file, yolo::Type::V8);
  if (yolo == nullptr)  return;
  
  // auto remote_show = create_zmq_remote_show();

  cv::Mat frame;
  cv::VideoCapture cap(0);
  if (!cap.isOpened()){
    printf("Engine is nullptr");
    return;
  }

  while(true){
    cap.read(frame);
    auto objs = yolo->forward(cvimg(frame));
    
    for(auto &obj : objs) {
      uint8_t b, g, r;
      tie(b, g, r) = yolo::random_color(obj.class_label);
      cv::rectangle(frame, cv::Point(obj.left, obj.top), cv::Point(obj.right, obj.bottom),
                    cv::Scalar(b, g, r), 5);
      
      auto name = mylabels[obj.class_label];
      auto caption = cv::format("%s %.2f", name, obj.confidence);
      int width = cv::getTextSize(caption, 0, 1, 2, nullptr).width + 10;
      cv::rectangle(frame, cv::Point(obj.left - 3, obj.top - 33),
                    cv::Point(obj.left + width, obj.top), cv::Scalar(b, g, r), -1);
      cv::putText(frame, caption, cv::Point(obj.left, obj.top - 5), 0, 1, cv::Scalar::all(0), 2, 16);
    }
      imshow("frame", frame);
      // remote_show->post(frame);
      int key = cv::waitKey(1);
      if (key == 27)
          break;
  }

  cap.release();
  cv::destroyAllWindows();
  return;
}

int main() {	// 修改2 调用该函数
  // perf();
  // batch_inference();
  // single_inference();
  yolo_video_demo("best.transd.sim.engine");
  return 0;
}

修改完成后执行make run即可看到对应的画面显示了

4.5 拓展-trtexec参数

trtexec是NVIDIA TensorRT SDK中的一个实用工具，它允许用户从命令行轻松运行和测试TensorRT引擎。trtexec命令行工具可以使用以下参数：

./trtexec [-h] [--uff model.uff [model.uff ...]] [--onnx model.onnx] [--model=model.plan] [--deploy=<deploy.prototxt>] [--output=<output_name>] [--batch=N] [--device=N] [--workspace=N] [--fp16] [--int8] [--calib=<dir>] [--useDLA=N] [--allowGPUFallback] [--iterations=N] [--avgRuns=N] [--verbose] [--nshapes=N] [--optShapes=NxN ... NxN] [--minShapes=NxN ... NxN] [--maxShapes=NxN ... NxN] [--shapeInput={0,1,2,...}] [--tacticSources] [--loadEngine=<filename>] [--saveEngine=<filename>] [--plugins=<XML file>] [--dumpOutput=<filename>] [--excludePlugin=<name>] [--start=<tag>] [--streams=N] [--batchTile] [--engine=<filename>] [--uffInput=input_name,input_shape] [--uffNHWC] [--uffNCHW] [--uffHW=<value>] [--workspaceSize=<size>] [--buildOnly] [--engineFormat=<format>] [--refit] [--saveRefinedEngine=<filename>] [--calibrator=<classname>] [--detach] [--check] [--fp16Char] [--int8Calib=<dir>] [--int8IO] [--disableTensorCores] [--useSpinWait] [--shapes=<shape>] [--maxBatch=<size>]

其中一些重要的参数如下：

  • --uff：指定输入为UFF模型，后面跟上模型文件的路径。
  • --onnx：指定输入为ONNX模型，后面跟上模型文件的路径。
  • --model：指定输入为序列化的引擎文件，后面跟上文件路径。
  • --deploy：指定输入为Caffe deploy文件的路径。
  • --output：指定输出Tensor名称。
  • --batch：指定执行推理时每个batch的大小，默认为1。
  • --device：指定执行推理的设备编号，默认为0。
  • --workspace：指定GPU内存的最大使用量，默认为1GB。
  • --fp16：启用FP16精度，可提高推理性能和减少内存使用。
  • --int8：启用INT8精度，可进一步提高推理性能和减少内存使用。
  • --calib：指定INT8校准数据集的路径。
  • --useDLA：指定使用哪个DLA，以及在DLA上运行哪些层。
  • --allowGPUFallback：如果使用DLA，当某些层无法在DLA上运行时，是否允许将其回退到GPU。
  • --iterations：指定测试迭代次数。
  • --avgRuns：指定平均运行次数。
  • --verbose：打印更详细的输出信息。
  • --loadEngine：指定加载的TensorRT引擎文件，后面跟上文件路径
  • --saveEngine：指定生成的TensorRT引擎文件，后面跟上文件路径

结语

本篇博客介绍了关于yolov8模型训练的流程，以及在jetson nano嵌入式上的部署工作，并安利了一个全新的推理部署框架infer。博主在这里只做了最基础的演示，如果有更多的需求需要各位看官自己去挖掘啦😄。感谢各位看到最后，创作不易，读后有收获的看官请帮忙点个👍⭐️。

下载链接

参考

Logo
GitCode 开源社区

旨在为数千万中国开发者提供一个无缝且高效的云端环境，以支持学习、使用和贡献开源项目。

更多推荐

cover

GitTalk | 使用面向业务的狮偶编程语言提升开发效率

avatar GitCode 开源社区
cover

GitTalk | DevUI Suits 场景解决方案

avatar GitCode 开源社区
cover

GitTalk | DevUI Admin 前端项目构建

avatar GitCode 开源社区