Jetson nano上使用ncnn部署yolox

序言

之前一篇文章写yolox的训练,这篇文章写一下关于部署,yolox之所以刚出来就这么火爆,不仅是精度高、速度快,很大程度上还是因为直接把部署的代码叶开源了,可谓是从训练到部署一条龙服务,很良心,yolox里面开源的部署代码有四种:ncnn、openvino、onnx和tensorrt,刚好手上有一个nano的板子,所以打算试试yolox使用ncnn在nano上的效果,写一篇文章记录一下过程。

yolox相关博文:

  • YOLOX自定义数据集训练(抢先踩坑)
  • 使用libtorch部署yolox

一、模型转换pytorch->onnx

使用ncnn部署,模型转换的步骤还是:pytorch ->onnx -> ncnn,所以第一步需要将pytorch框架下训练的模型先转成onnx,然后在通过onnx转成ncnn进行部署。部署的时候在两台机器上操作,一是我的电脑主机,用于pytorch到onnx的模型转换及onnx模型简化,二是nano,用来编译ncnn及一些cpp代码修改和最后的编译代码工作。

转onnx的步骤很简单,因为代码已经在仓库中提供了,所以我们只需要在终端运行指令即可:

python tools/export_onnx.py -n yolox-s -c weights/best_ckpt.pth.tar

这里我加载了上次训练好的权重,权重放在weights文件夹下,运行后模型默认输出的是yolox.onnx,在主目录下可以找到。

特别要注意的是,因为这里默认构建的模型是按照coco数据集来的,所以分类数是80,因为我训练的网络分类数只有1类,所以在运行之前要修改exps/default/yolox_s.py文件,将类别数改为1:
Jetson nano上使用ncnn部署yolox_第1张图片
然后再重新运行上面的命令,得到yolox.onnx文件,得到onnx文件后我们需要使用onnx-simplifier来简化一下模型,可以一定程度上减少模型里面出现的一些胶水OP,所以需要在终端中运行如下代码:

pip install onnx-simplifier                  # pip安装,如果已经安装了可以不用执行这一步
python -m onnxsim yolox.onnx yolox_sim.onnx      # 直接运行

Jetson nano上使用ncnn部署yolox_第2张图片
这里就已经成功了,得到一个新的yolox_sim.onnx简化后的onnx模型。

二、nano上编译ncnn环境

这一步也比较简单,网上的教程也比较多,我这里就直接贴出运行的命令即可,关于nano的刷机以及opencv的编译可以参考我之前的文章:Jetson nano系统烧录及pytorch环境配置,然后在这个基础之上直接编译ncnn:

https://github.com/Tencent/ncnn.git                  # git ncnn源码
cd ncnn
git submodule update --depth=1 --init
mkdir build && cd build
cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/jetson.toolchain.cmake -DNCNN_VULKAN=ON -DCMAKE_BUILD_TYPE=Release ..
make -j4
sudo make install

编译完成后可以测试一下,把ncnn主目录下子文件夹中的param文件全部复制到刚才的build/benchmark目录中,然后进到ncnn/build/benchmark目录中执行:
Jetson nano上使用ncnn部署yolox_第3张图片
看到如上输出则为编译成功。

三、模型转换onnx->ncnn

编译成功好需要将onnx模型转换成ncnn模型,所以将刚才生成的yolox_sim.onnx模型拷贝到build/tools/onnx文件夹中,然后执行如下语句:

./onnx2ncnn yolox_sim.onnx yolox.param yolox.bin

Jetson nano上使用ncnn部署yolox_第4张图片
得到如上输出,不过不要紧,这是focus模块转换的报错,ncnn中已经实现了focus,只需要按照如下修改,打开yolox.param文件,修改这两处:
Jetson nano上使用ncnn部署yolox_第5张图片
其中圈起来的Crop操作即为focus模块,因为ncnn中已经实现了,所以我们需要将这十层删除,然后添加一层新的focus层,同样的网络层数也要随之改变,由235-10+1 = 226层,修改后的param文件为:

YoloV5Focus              focus                    1 1 images 503

Jetson nano上使用ncnn部署yolox_第6张图片
后面的503是Concat层的输出,记得要对应。修改完后重新用用ncnnoptimize对模型进行优化,顺便转为 fp16 存储减小模型体积,还是在onnx目录下:

../ncnnoptimize yolox.param yolox.bin yolox.param yolox.bin 65536

前面两个参数是优化前的模型,后面两个是优化后的模型,我这里直接覆盖之前的模型保存。
Jetson nano上使用ncnn部署yolox_第7张图片
运行后得到新的yolox.param和yolox.bin文件,这时候模型转换就完成了。

四、编译yolox.cpp

模型转换完成之后,可以直接把yolox.cpp文件复制到ncnn工程的examples目录下,然后在ncnn/examples/CMakeLists.txt中做如图修改,添加yolox工程:
Jetson nano上使用ncnn部署yolox_第8张图片
然后需要修改一下yolox.cpp文件,将这几个地方做出对应的修改:
Jetson nano上使用ncnn部署yolox_第9张图片
没编译vulkan,所以选择false,并且将ncnn模型和权重文件地址改为自己的绝对路径地址,因为我的模型是在640上效果更好一些,所以尺寸改成了640,保存后接着可以直接重新编译一遍ncnn源码:

cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/jetson.toolchain.cmake -DNCNN_VULKAN=ON -DCMAKE_BUILD_TYPE=Release ..
make -j4

Jetson nano上使用ncnn部署yolox_第10张图片
可以看到已经编译成功了,这时候打开examples目录可以看到yolox的可执行文件,我将一张1.jpg的图片拷贝到该目录中测试一下,执行:

./yolox 1.jpg

Jetson nano上使用ncnn部署yolox_第11张图片
看下实际的效果图,和之前训练出来的基本上一致,到这里就搞完了:

你可能感兴趣的:(深度学习,ncnn,pytorch)