Jetson nano部署剪枝YOLOv8

注意事项

一、2023/5/6更新

有看官反映在剪枝后的YOLOv8不要从yaml导入网络结构中，按照我提示的修改后发现并无任何改变，YOLOv8还是从yaml中导入，这里提供下另外一种方式吧！！！请参考第5小结内容

前言

写本文的目的是刚学习完了剪枝与重参的课程，想跟着梁老师复现下对YOLOv8模型进行简单的剪枝，熟悉下整个流程。剪枝完成后我们可以将剪枝后的YOLOv8部署到Jetson嵌入式端，进一步提高其检测速度。于是乎我又可以水一篇文章了(哈哈哈，机智如我，我真是个天才😏)

默认大家对模型剪枝和模型部署都有一定的了解，大家可以查看我之前的剪枝与重参第七课：YOLOv8剪枝和Jetson nano部署YOLOv8，很多细节就不在这里赘述了。考虑到nano的算力，这里采用yolov8s.pt，本文主要分享yolov8s.pt模型剪枝和jetson nano部署剪枝yolov8两方面的内容。若有问题欢迎各位看官批评指正!!!😘

一、YOLOv8模型剪枝训练

YOLOv8模型剪枝训练的基本流程如下：

首先我们获得一个预训练模型，用做benchmark方便后续对比，然后进行约束训练，主要对BN层加上L1正则化，获得约束训练的模型后我们就可以对其进行剪枝了，最后将剪枝后的模型进行微调即可。

1. Pretrain[option]

获取预训练模型其实非必要，博主在这里为了方便对比，故选择进行预训练。

1.1 项目的克隆

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

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

1.2 数据集

训练采用的VOC数据集，这里给出下载链接Baidu Drive[pwd:yolo]，本次训练并没有用到所有的数据，博主将train2007和val2007作为训练集，将test2007作为验证集，整个数据集文件夹内容如下图所示：

其中，images存放的内容是图片文件，labels存放的内容是YOLO格式的.txt标签文件，所有文件都可以从我分享的链接下载，大家可以按照上述方式将数据集进行整合。

1.3 训练

代码和数据集准备好后就可以进行训练了，训练修改的文件主要是两个即VOC.yaml用于指定数据集的相关路径和数据集包含的类别信息，default.yaml用于指定训练用到的权重和一些超参数，我们一个一个来修改。

VOC.yaml位于ultralytics/dataset下，其具体内容如下：

• 首先path路径指定为上面整合的数据集的绝对路径，路径中最好不要含中文，在Windows下训练时最好将路径中的\替换为\\或者/，防止\转义。
• train、val、test的内容就是VOC数据集下的用于训练、验证以及测试的图片
• names不用修改
• download内容全部删除即可

# Ultralytics YOLO 🚀, GPL-3.0 license
# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC by University of Oxford
# Example usage: yolo train data=VOC.yaml
# parent
# ├── ultralytics
# └── datasets
#     └── VOC  ← downloads here (2.8 GB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: D:/YOLO/yolov8-prune/ultralytics/datasets/VOC
train: # train images (relative to 'path')  16551 images
  - images/train2007
  - images/val2007
val: # val images (relative to 'path')  4952 images
  - images/test2007
test: # test images (optional)
  - images/test2007

# Classes
names:
  0: aeroplane
  1: bicycle
  2: bird
  3: boat
  4: bottle
  5: bus
  6: car
  7: cat
  8: chair
  9: cow
  10: diningtable
  11: dog
  12: horse
  13: motorbike
  14: person
  15: pottedplant
  16: sheep
  17: sofa
  18: train
  19: tvmonitor

default.yaml是一个配置文件，位于ultralytics/yolo/cfg下，其需要修改的内容如下：

• model即预训练模型，本次选用yolov8s.pt，可以从官网上下载好预训练模型放到v8/detect目录下，或者不用下载，你指定yolov8s.pt训练后会检测v8/detect路径下是否存在yolov8s.pt，不存在会直接去官网帮你下载(下载时需访问外网，如果没有代理还是手动下载好放到v8/detect目录下吧)
• data即训练配置文件路径，也就是上面配置的VOC.yaml的绝对路径
• epochs即训练迭代次数，这个得根据自己的显卡算力来，博主显卡不太顶用，100个epoch时间太长了，遭不住呀
• amp即自动混合精度，有非常多的好处(比如…)，但是在剪枝后需要finetune就不开启了，开启后需要很多地方设置，麻烦

# Train settings -------------------------------------------------------------------------------------------------------
model: yolov8s.pt  # path to model file, i.e. yolov8n.pt, yolov8n.yaml
data: D:/YOLO/yolov8-prune/ultralytics/datasets/VOC.yaml # path to data file, i.e. coco128.yaml
epochs: 50  # number of epochs to train for
amp: False  # Automatic Mixed Precision (AMP) training, choices=[True, False], True runs AMP check

将前面的步骤完成后，在ultralytics/yolo/v8/detect文件夹下找到train.py文件点击运行即可开始模型的训练。

博主训练的模型为yolov8s.pt且使用单个GPU进行训练，显卡为RTX3060矿卡(😂)，操作系统为Windows10，pytorch版本为1.12.1，训练时长大概2个小时

训练完成后在detect文件夹下会生成一个runs文件夹，模型权重就保存在runs/detect/train/weights文件夹下，这里提供博主训练好的权重文件下载链接Baidu Driver[pwd:yolo]

2. Constraint training

约束训练主要是对模型进行BN层进行L1正则化，因此需要在trainer.py文件夹下添加BN层进行L1约束的代码，trainer.py文件位于ultralytics/yolo/engine文件夹下，添加的具体位置在327行，添加的具体内容如下：

# Backward
self.scaler.scale(self.loss).backward()

# ========== 新增 ==========
l1_lambda = 1e-2 * (1 - 0.9 * epoch / self.epochs)
for k, m in self.model.named_modules():
    if isinstance(m, nn.BatchNorm2d):
        m.weight.grad.data.add_(l1_lambda * torch.sign(m.weight.data))
        m.bias.grad.data.add_(1e-2 * torch.sign(m.bias.data))
# ========== 新增 ==========

# Optimize - https://pytorch.org/docs/master/notes/amp_examples.html
if ni - last_opt_step >= self.accumulate:
    self.optimizer_step()
    last_opt_step = ni

将代码修改好后，按照之前提到的Pretrain中将VOC.yaml和default.yaml修改好，点击train.py开始训练即可。这里提供博主训练好的权重文件下载链接Baidu Driver[pwd:yolo]，约束训练完成效果如下图所示：

3. Prune

我们拿到约束训练的模型后就可以开始剪枝了，开工👨‍🏭，本次剪枝使用的是约束训练中的last.pt模型(我们不使用best.pt，通过result.csv你会发现mAP最高的模型在第一个epoch，主要是因为模型在COCO数据集上的预训练泛化性比较强，所以开始的mAP很高，这显然是不真实的)，我们在根目录ultralytics-main文件夹下创建一个prune.py文件，用于我们的剪枝任务，同时将约束训练中的last.pt模型放到根目录下，prune.py文件的具体内容如下所示：

from ultralytics import YOLO
import torch
from ultralytics.nn.modules import Bottleneck, Conv, C2f, SPPF, Detect

# Load a model
yolo = YOLO("last.pt")  # build a new model from scratch
model = yolo.model

ws = []
bs = []

for name, m in model.named_modules():
    if isinstance(m, torch.nn.BatchNorm2d):
        w = m.weight.abs().detach()
        b = m.bias.abs().detach()
        ws.append(w)
        bs.append(b)
        # print(name, w.max().item(), w.min().item(), b.max().item(), b.min().item())
# keep
factor = 0.8
ws = torch.cat(ws)
threshold = torch.sort(ws, descending=True)[0][int(len(ws) * factor)]
print(threshold)

def prune_conv(conv1: Conv, conv2: Conv):
    gamma = conv1.bn.weight.data.detach()
    beta = conv1.bn.bias.data.detach()
    keep_idxs = []
    local_threshold = threshold
    while len(keep_idxs) < 8:
        keep_idxs = torch.where(gamma.abs() >= local_threshold)[0]
        local_threshold = local_threshold * 0.5
    n = len(keep_idxs)
    # n = max(int(len(idxs) * 0.8), p)
    # print(n / len(gamma) * 100)
    # scale = len(idxs) / n
    conv1.bn.weight.data = gamma[keep_idxs]
    conv1.bn.bias.data = beta[keep_idxs]
    conv1.bn.running_var.data = conv1.bn.running_var.data[keep_idxs]
    conv1.bn.running_mean.data = conv1.bn.running_mean.data[keep_idxs]
    conv1.bn.num_features = n
    conv1.conv.weight.data = conv1.conv.weight.data[keep_idxs]
    conv1.conv.out_channels = n

    if conv1.conv.bias is not None:
        conv1.conv.bias.data = conv1.conv.bias.data[keep_idxs]

    if not isinstance(conv2, list):
        conv2 = [conv2]

    for item in conv2:
        if item is not None:
            if isinstance(item, Conv):
                conv = item.conv
            else:
                conv = item
            conv.in_channels = n
            conv.weight.data = conv.weight.data[:, keep_idxs]


def prune(m1, m2):
    if isinstance(m1, C2f):  # C2f as a top conv
        m1 = m1.cv2

    if not isinstance(m2, list):  # m2 is just one module
        m2 = [m2]

    for i, item in enumerate(m2):
        if isinstance(item, C2f) or isinstance(item, SPPF):
            m2[i] = item.cv1

    prune_conv(m1, m2)


for name, m in model.named_modules():
    if isinstance(m, Bottleneck):
        prune_conv(m.cv1, m.cv2)

seq = model.model
for i in range(3, 9):
    if i in [6, 4, 9]: continue
    prune(seq[i], seq[i + 1])

detect: Detect = seq[-1]
last_inputs = [seq[15], seq[18], seq[21]]
colasts = [seq[16], seq[19], None]
for last_input, colast, cv2, cv3 in zip(last_inputs, colasts, detect.cv2, detect.cv3):
    prune(last_input, [colast, cv2[0], cv3[0]])
    prune(cv2[0], cv2[1])
    prune(cv2[1], cv2[2])
    prune(cv3[0], cv3[1])
    prune(cv3[1], cv3[2])

for name, p in yolo.model.named_parameters():
    p.requires_grad = True

# yolo.val() # 剪枝模型进行验证 yolo.val(workers=0)
# yolo.export(format="onnx") # 导出为onnx文件
# yolo.train(data="VOC.yaml", epochs=100) # 剪枝后直接训练微调

torch.save(yolo.ckpt, "prune.pt")

print("done")

我们通过上述代码可以完成剪枝工作并将剪枝好的模型进行保存，用于finetune，有以下几点说明：

• 在本次剪枝中我们利用factor变量来控制剪枝的保留率

• 我们用来剪枝的模型一定是约束训练的模型，即对BN层加上L1正则化后训练的模型

• 约束训练后的b.min().item值非常小，接近于0或者等于0，可以依据此来判断加载的模型是否正确

• 我们可以选择将yolo.train()取消注释，在剪枝完成直接进入微调训练，博主在这里选择先保存剪枝模型

• 我们可以选择yolo.export()取消注释，将剪枝完成后的模型导出为ONNX，查看对应的大小和channels是否发生改变，以此确认我们完成了剪枝

• yolo.val()用于进行模型验证，建议取消注释进行相关验证，之前梁老师说yolo.val()验证的mAP值完全等于0是不正常的，需要检查下剪枝过程是否存在错误，最好是有一个值，哪怕非常小，博主剪枝后进行验证的结果如下图所示，可以看到mAP值真的是惨不忍睹(🤣)，所以后续需要finetune模型来恢复我们的精度

在我们打开yolo.val()进行模型剪枝任务时可能会出现如下问题，这个错误出现在Windows下面，原因在于Linux系统中可以使用多个子进程加载程序，而Windows则不能，解决办法就是将workers设置为0，将yolo.val()修改为如下代码，参考自解决RuntimeError: An attempt has been made to start a new process before…办法

yolo.val(workers=0)

我们拿到剪枝后的model后可以导出为ONNX，来看看剪枝前后模型的差异性，首先从模型大小来看，剪枝前的ONNX模型大小为42.6MB，剪枝后的ONNX模型大小为35.4MB，然后从ONNX模型对比图来看，channels发生了变化，具体可看下面的示例图，可以看到剪枝前后Conv的通道数发生了明显的变化。这里提供博主剪枝好的权重文件下载链接Baidu Driver[pwd:yolo]

4. finetune

拿到剪枝的模型后，我们需要先做两件事情

• 1.切记！！！在进行finetune之前需要将我们在trainer.py为BN层添加的L1正则化的代码注释掉(也就是我们在第2节添加的内容)
• 2.切记！！！剪枝后不要从yaml导入结构。如果我们直接将剪枝后的模型prune.pt放到v8/detect目录下修改default.yaml文件，然后点击train.py是会存在问题的，此时模型结构是通过yolov8.yaml加载的，而我们并没有修改yaml文件，因此finetune的模型其实并不是剪枝的模型

因此，正常finetune训练的步骤如下：

• 1.在yolo/engine/trainer.py中注释掉为BN层加L1正则化的代码

• 2.修改yolo/engine/model.py代码，让其不要从yaml导入网络结构，具体修改内容是BaseTrainer类中的setup_model方法中，代码大概在443行左右，新增一行代码即可，如下所示

  # ========== yolo/engine/trainer.py的443行 ==========
  self.model = self.get_model(cfg=cfg, weights=weights, verbose=RANK == -1)  # calls Model(cfg, weights)
  
  # ========== 新增该行代码 ==========
  self.model = weights
  
  return ckpt
  

• 3.将剪枝完保存的模型放到yolo/v8/detect文件夹下

• 4.修改default.yaml文件，主要修改model为prune.pt即剪枝完的模型，具体修改如下：

  model: prune.pt  # path to model file, i.e. yolov8n.pt, yolov8n.yaml
  

• 5.点击train.py开始训练即可，博主在这里选择的是微调50个epoch，大家根据自己的实际情况来，尽可能的多finetune几个epoch

微调50个epoch后模型的表现如下图所示，可以看到精度恢复得还可以，可以训练更多epoch使其精度更加稳定。这里提供博主训练好的权重文件下载链接Baidu Driver[pwd:yolo]

OK!至此，YOLOv8模型剪枝训练部分完成，下面来开始部署剪枝模型🚀🚀🚀

5. 剪枝后不从yaml导入结构(补充细节)

剪枝完成后的模型该如何正确的加载并训练呢？这里再提供另外一种方案，供大家借鉴参考，主要修改两个地方.

修改1：修改网络加载的地方，让其不要从yaml导入结构

• 具体修改代码在 ultralytics/yolo/engine/model.py
• 具体位置在 YOLO 类的 train 方法中，大概是 363 行的位置
• 修改代码如下：

# ===== ultralytics/yolo/engine/model.py 363行=====

# ===== 原代码 =====
if not overrides.get('resume'):  # manually set model only if not resuming
    self.trainer.model = self.trainer.get_model(weights=self.model if self.ckpt else None, cfg=self.model.yaml)
    self.model = self.trainer.model
            
# ===== 修改后代码 =====
if not overrides.get('resume'):  # manually set model only if not resuming
    # self.trainer.model = self.trainer.get_model(weights=self.model if self.ckpt else None, cfg=self.model.yaml)
    # self.model = self.trainer.model
    self.trainer.model = self.model.train()

修改2：自己新增一个my_train.py的训练代码

在主目录下新建一个 my_train.py 文件用于训练，该文件内容非常简单，如下所示：

from ultralytics import YOLO


if __name__ == "__main__":
    yolo = YOLO("prune.pt") # 加载剪枝后的模型

    yolo.train(data="D:/YOLO/yolov8-prune/ultralytics/datasets/VOC.yaml", epochs=50, amp=False, workers=8) # 训练

其中有以下几点值得注意：

• 加载剪枝后的模型，请修改为你自己的剪枝模型名称
• 关于训练参数的指定
  • data 表示之前训练时的 VOC.yaml 文件的绝对路径
  • epochs 迭代次数，根据自己实际需求设置
  • amp 混合精度设置为False
  • workers 工作核心数，根据自己的硬件设置，值越大训练越快
  • 其它超参数博主并未设置，有需求可自行修改

剪枝模型微调训练

完成上述两点修改后，点击 my_train.py 即可开始剪枝模型的微调训练，训练后的文件会保存在 runs/detect/train 中，在训练之前请确保已经将 trainer.py 为 BN 层添加的 L1 正则化的代码注释掉！！！

二、YOLOv8模型剪枝部署

Jetson nano上yolov8的部署使用到的Github仓库是infer。想了解通过TensorRT的Layer API一层层完成模型的搭建工作可参考Jetson嵌入式系列模型部署-2，想了解通过TensorRT的ONNX parser解析ONNX文件来完成模型的搭建工作可参考Jetson嵌入式模型部署-3、Jetson nano部署YOLOv7。本文主要利用infer来对剪枝后的yolov8完成部署，本文参考自Jetson nano部署YOLOv8，具体流程该文描述非常详细，这里再简单过一遍，本次训练的模型使用yolov8s.pt，类别数为20，数据集是VOC数据集，部署的模型是经过剪枝后finetune的模型。

1. 源码下载

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

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

也可手动点击下载，点击右上角的Code按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击here[pwd:infer]下载博主准备好的源代码(注意代码下载于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是什么？如何选？。(吐槽下这玩意上传忒慢了，超级会员不顶用呀，终于上传完了，折磨!!!)

2.1 trtexec环境变量设置

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

1.打开bashrc文件

vim ~/.bashrc

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

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

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

source ~/.bashrc

3. ONNX导出

• 训练的模型使用yolov8n.pt，torch版本1.12.1，onnx版本1.13.1
• 参考自YoloV8的动态静态batch如何理解和使用

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

静态batch

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

动态batch

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

说明：本次为了方便仅使用静态batch，关于动态batch的使用可参考Jetson nano部署YOLOv8

3.1 Transpose节点的添加

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

from ultralytics import YOLO

yolo = YOLO("finetune_best.pt")

yolo.export(format="onnx", batch=1)

模型需要完成修改才能正确被infer框架使用，正常模型导出的输出为[1,24,8400]，其中1代表batch，24分别代表cx,cy,w,h以及VOC中20个类别分数，8400代表框的个数。首先infer框架的输出只支持[1,8400,24]这种形式的输出，因此我们需要再原始的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节点，执行完成之后在当前目录下生成finetune_best.transd.onnx模型，该模型添加了Transpose节点。

python v8trans.py finetune_best.onnx

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

3.2 Resize节点解析的问题

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

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

下面来看解决方案onnxsim

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

import onnx
from onnxsim import simplify

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

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

至此，模型导出已经完毕，后续通过导出的模型完成在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/finetune_best.transd.sim.onnx --saveEngine=workspace/finetune_best.transd.sim.engine

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

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

trtexec --onnx=finetune_best.transd.sim.onnx --minShapes=images:1x3x640x640 --maxShapes=images:16x3x640x640 --optShapes=images:1x3x640x640 --saveEngine=finetune_best.transd.sim.engine

4.2 源码修改

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

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

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

• 3.main.cpp 10行 新增voclabels数组，添加voc的类别名称

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

具体修改如下

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

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

static const char *voclabels[] = {"aeroplane",   "bicycle", "bird",   "boat",       "bottle",
                                  "bus",         "car",     "cat",    "chair",      "cow",
                                  "diningtable", "dog",     "horse",  "motorbike",  "person",
                                  "pottedplant",  "sheep",  "sofa",   "train",      "tvmonitor"};	// 修改3 新增voclabels数组

auto name = voclabels[obj.class_label]		// 修改4 cocolabels修改为mylabels

4.3 编译运行

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

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 = voclabels[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("finetune_best.transd.sim.engine");
  return 0;
}

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

三、讨论

讨论1：我们来看看剪枝前后模型的差异性，首先，先来看mAP指标，剪枝前mAP50为0.808，剪枝后mAP50为0.789(注：mAP略微下降，可能是迭代次数epoch太少)。

再来看检测速度，从两方面来看：

第一方面从trtexec工具构建的模型来看，我们只需要关注一个数字，就是latency中的mean，它代表平均推理一张图的耗时，但是不包含前后处理，不进行剪枝的模型使用trtexec构建工具，其latency中的mean为161.974ms(下图1)；进行剪枝的模型使用trtexec构建工具，其latency中的mena为133.107ms(下图2)；剪枝后的模型快了将近30ms，对比可知模型剪枝后的推理速度明显提高了

第二方面从infer中的perf()函数来看，我们只需要关注[BATCH1]的耗时，因为我们没有设置动态batch，且静态batch设置为1，它也代表平均一张图的耗时，但是是包含前后处理的，不进行剪枝的模型使用perf()函数测试性能，其BATCH1的耗时为166ms(下图3)，进行剪枝的模型使用perf()函数测试性能，其BATCH1的耗时为137ms(下图4)；剪枝后的模型快了将近30ms，从这里对比也可知模型剪枝后的推理速度提高了。

额外补充一句，将perf()函数测试的耗时减去trtexec工具测试的耗时就可以得出前后处理的时间，大概计算下，在Jetson nano上YOLOv8的前后处理时间仅仅耗时4ms左右，不得不佩服杜老师呀😂

图1 不进行剪枝模型的trtexec工具耗时(161.974ms)

图2 进行剪枝模型的trtexec工具耗时(133.107ms)

图3 不进行剪枝模型的perf函数推理耗时(BATCH1为166ms)

图4 进行剪枝模型的perf函数推理耗时(BATCH1为137ms)

讨论2：在之前的剪枝代码中，我们始终保持着channels大于8，小于8的时候我们就降低阈值，选择更多的通道数。可以看到剪枝后的channels千奇百怪，没有规律，而原始的ONNX的通道数都是8的倍数。我记得梁老师说过，对于保留的channels，它应该整除n才是最合适的，也就是说我们在prune剪枝的时候应该控制下channels的数量，让它能整除n，因为NVIDIA的硬件加速更加合适。n该如何选择呢？一般FP16模型，n选取8；INT8模型，n选取16。

罗里吧嗦一大堆，就是想表达下如果我们需要对剪枝后的模型进行量化加速时(比如利用NVIDIA的tensorRT)，是不是不应该向上面剪枝剪得那么随意呢？🤔，不然channels无法整除8或者16呀，而当量化生成FP16或者INT8模型时，硬件加速就不太顶了呀

讨论3：在约束训练添加的代码中，我们将L1正则化的约束设置为1e-2，在剪枝即prune.py的代码中，我们将剪枝保留率设置为0.8，这些参数其实都是超参数，需要考虑模型复杂度和性能的平衡，对于不同的网络以及不同的需求这些超参数可能是不一致的。

结语

本篇博客简单重新实现了下之前梁老师讲解的YOLOv8模型剪枝，然后将剪枝后的模型部署到Jetson nano上，在mAP值略微下降的前提下(mAP下降0.019)，大大提高了检测速度(推理速度快了30ms)，也算把流程都简单过了一遍吧。博主在这里只做了最简单的实现，并没有做原理分析，具体原理和细节那就需要各位看官自行去了解啦😄。感谢各位看到最后，创作真心不易，读后有收获的看官请帮忙点个👍⭐️。

下载链接

• 原始VOC数据集下载链接[pwd:yolo]

• VOC权重文件下载链接[pwd:yolo]

• JetPack镜像下载链接Baidu Drive[pwd:nano]【更新完毕!!!】

参考

• 剪枝与重参第七课：YOLOv8剪枝

• Jetson nano部署YOLOv8

• YOLOv8

• 解决RuntimeError: An attempt has been made to start a new process before…办法

• infer

• Jetson嵌入式系列模型部署-1

• Jetson嵌入式系列模型部署-2

• Jetson嵌入式系列模型部署-3

• Jetson nano部署YOLOv7

• Jetson nano部署YOLOv8

• Jetson NANO是什么？如何选？

• trtexec添加环境变量

• YoloV8的动态静态batch如何理解和使用【视频】

• Netron可视化工具

• onnxoptimizer、onnxsim使用记录

• 如何高效使用TensorRT【视频】

• Makefile实战

• Jetson nano部署YOLOv8

• Jetson NANO是什么？如何选？

• trtexec添加环境变量

• YoloV8的动态静态batch如何理解和使用【视频】

• Netron可视化工具

• onnxoptimizer、onnxsim使用记录

• 如何高效使用TensorRT【视频】

• Makefile实战