Jetson nano部署剪枝YOLOv8
目录
- 前言
- 一、YOLOv8模型剪枝训练
- 1. Pretrain[option]
- 1.1 项目的克隆
- 1.2 数据集
- 1.3 训练
- 2. Constraint training
- 3. Prune
- 4. finetune
- 二、YOLOv8模型剪枝部署
- 1. 源码下载
- 2. 环境配置
- 2.1 trtexec环境变量设置
- 3. ONNX导出
- 3.1 Transpose节点的添加
- 3.2 Resize节点解析的问题
- 4. 运行
- 4.1 engine生成
- 4.2 源码修改
- 4.3 编译运行
- 4.4 拓展-摄像头检测
- 三、讨论
- 结语
- 下载链接
- 参考
前言
写本文的目的是刚学习完了剪枝与重参的课程，想跟着梁老师复现下对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模型剪枝训练部分完成，下面来开始部署剪枝模型🚀🚀🚀
二、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左右，不得不佩服杜老师呀😂
讨论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实战