上一期我们教大家如何给新的JetsonNano 2GB烧录系统。这一期我们将教大家如何在JetsonNano上部署最新的Yolov5检测模型,并且采用TensorRT加速,看看我们的模型能否在JetsonNano这样的小设备上跑到实时。
首先我们来确认一下系统里面opencv对应的版本:
是最新的4.1,不错。
今天这是我们要实现的效果:
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-Us5L9WaW-1605327260590)(https://i.loli.net/2020/11/11/OeuA2y3MkF4fSbP.png)]
是的,你没有看错,我们将在JetsonNano上跑一个经过TensorRT加速的YoloV5算法,算法的功能是实现口罩佩戴的检测。从图中你可以看到,对于口罩佩戴的检测还是很准确的,而且基本上可以达到实时!这里要说明的是,这个模型的输入尺寸不大,并且在Yolov5的基础上经过了裁剪,所以能够做到这么快,但是如果你的模型直接拿官方的部署,可能无法跑到这个速度,最后速度带来的当然是精度的降低,不过在可以接受的范围之内!
还等什么,快来依照我们的教程来实现吧~
技术栈
本篇文章不涉及到Python或者任何训练方面的东西,纯部署,而且纯C++,因此教程里面可能会出现一些你不熟悉的库或者术语,没关心,你看到最后我们跑出来的结果就可以了,至于过程你可以通过加入我们的社群来获取想要的材料。大部分代码都在github开源。核心技术包括:
- C++;
- TensorRT相关基础知识;
- 对ONNX了解;
- 对Yolov5的基本网络结构了解。
准备工作
先做一些准备工作,什么准备工作呢?我们需要先安装一个C++下面比较常用的可视化库,thor:
git clone https://github.com/jinfagang/thor
大家可以进入到主页依照readme进行编译。为什么要做这个准备干工作?原因有两个:
- 它依赖一些库,如果你能够成功编译它,说明你的环境没有问题,如果不能,你的环境可能有问题了;
- 它可以很方便的可视化我们的检测输出结果。
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-MmRVptpR-1605327260592)(https://i.loli.net/2020/11/11/mHyaS6R2q4fJpDk.png)]
引用一张thor官方的图:
画检测结果既好看又方便!
编译onnx-tensorrt
我们给大家准备好了一个转好的onnx文件,如果你需要导出自己的onnx,过程有点复杂,你可以通过加入社群交流如何导出,我们没有现成的代码提供。
git clone --recursive https://github.com/onnx/onnx-tensorrt
编译onnx-tensorrt 在JetsonNano上可能会遇到一些问题。为什么呢?
第一个错误:
CMake Error at CMakeLists.txt:21 (cmake_minimum_required):
CMake 3.13 or higher is required. You are running version 3.10.2
-- Configuring incomplete, errors occurred!
然后你就需要重新编译一下CMake了:
git clone https://github.com/Kitware/CMake
说实话这个CMake也是异常的庞大了。clone下来100多M。然后开始编译CMake了:
./bootstrap.sh
make
sudo make install
cmake -V
编译过程还是挺慢的.
然后就可以接着编译onnx-tensorrt了。
编译完onnx-tensorrt之后,就可以用 onnx2trt
将模型转出到tensorrt 的engine了。
在JetsonNano上将ONNX转到engine文件
想必大家已经安装好了必备的工具,如果你在安装工具的时候遇到问题,不要灰心丧气,因为最终是可以work的。我们已经走到这一步了。
接下来就是拿一个onnx来做试验了。这里我推荐一个onnx文件,大小只有6M,我们将用它来实现前面的检测效果。
至于如何训练对应的模型,大家可以在神力平台找到对应的训练代码。
http://manaai.cn
我们成功的将onnx转到了engine:
onnx2trt best_sim.onnx -o a.trt -w 10909920
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-FlBdC4Ex-1605327260593)(https://i.loli.net/2020/11/14/kCdPOfblJW41cEz.png)]
这里需要注意的是这个-w
, 它指的是我们的workspace多大,你不要默认,也不要设置的太小,大概按照我这样的设置就可以了,默认的话就死机了,太小会报错。
对TensorRT Engine进行推理
接下来就要对这个Engine文件进行推理了。推理部分我们要实现load上面的tensorrt engine并实现口罩的后处理检测。对应的代码也可以在神力平台找到:
http://manaai.cn
最后我们这里贴出核心的推理代码:
YoloV5 model(inputParams, trtParams, v5Params);
if (!model.initSession(0)) {
std::cerr << "model init failed!\n";
exit(0);
}
string data_f = FLAGS_data_f;
if (thor::os::isfile(data_f)) {
if (thor::os::suffix(data_f) == "mp4" ||
thor::os::suffix(data_f) == "avi") {
std::cout << "inference on video.\n";
VideoCapture cap(data_f);
if (!cap.isOpened()) {
cout << "Error opening video stream or file" << endl;
return -1;
}
cap.set(CAP_PROP_POS_FRAMES, 60);
VideoWriter video("yolov5_result.mp4", CV_FOURCC('M', 'J', 'P', 'G'), 25,
Size(cap.get(3), cap.get(4)));
while (1) {
Mat image;
cap >> image;
if (image.empty()) break;
auto bboxes = model.predOneImage(image);
std::vector<thor::Detection> dets = toThorDetections(bboxes);
auto res =
thor::vis::VisualizeDetections(image, dets, CLASSES, &COLORS);
cv::imshow("yolov5", res);
video.write(res);
char c = (char)waitKey(25);
if (c == 27) break;
}
video.release();
}
由此可以得到我们上面给大家展示的效果。
总结
我们对JetsonNano 2GB进行了较为全面的评测,包括它的软件兼容性以及在一些极致优化的模型下的表现速度。事实上,如果用我们的512输入尺寸的Yolov5模型的TensorRT加速,它的最高速度可以达到25ms, 这基本上可以让你在一个嵌入式板子上将检测模型跑到realtime.
所以当你需要部署某个应用的时候,缺乏硬件考量?JetsonNano 2GB绝对是一个不错的首选。
哦对了对了,这次评测忘了一个很重要的点: JetsonNano 2GB的TensorCore支持fp16推理!! 这意味着在我们的评测速度上,JetsonNano 2GB还可以做到更快的速度!
本次教程的所有相关代码均为我们原创,转载请注明出处,谢谢!