Jetson Nano目标检测手把手实战教程（pytorch训练、tensorrt推理，含完整代码和数据）

目录

一、概述

1.1 深度学习和边缘计算

1.2 Jetson Nano简介

二、深度学习环境安装

2.1 Pytorch框架

2.2 在Jetson Nano上安装Pytorch

三、算法原理

四、具体实现步骤

4.1 工程代码和环境准备

4.2 模型训练和推理

4.2.1数据集准备

4.2.2训练

4.2.3模型转换

4.2.4 实时推理

五、总结

参考文献：

---

一、概述

1.1 深度学习和边缘计算

2012年AlexNet模型在ImageNet数据集上取得了巨大成功，从此，深度学习技术开始进入黄金发展期。随后，深度学习在多个应用领域均大幅提升了性能指标，尤其在计算机视觉领域，超越SVM等一众传统机器学习方法。然而，深度学习的高准确性是以深度学习对计算和内存的高需求为代价的。由于输入数据可能具有很高的维数，并且需要对输入数据执行数百万次计算，因此训练和推理一个深度学习模型对空间和计算需求都很高。高准确性和高资源消耗是深度学习的特点。

​ 为了满足深度学习的计算需求，一种常见的方法是利用云计算。要使用云资源，必须将数据从网络边缘的数据源推送上云然后再在云上进行推理。这种将数据从源迁移到云的解决方案带来了几个问题：延迟大、服务器负荷重、隐私性弱。面对这3个棘手的问题，越来越多的研究学者和商业公司将目标瞄向边缘计算场景。

所谓边缘计算，简单来理解就是赋能传统终端设备（安卓、ios、嵌入式设备等）以深度学习推理能力，直接在前端就完成深度学习推理，而不需要将图像等数据通过web传至总服务器再处理，这样能极大的减少后端服务器压力。从深度学习的训练和推理角度来理解，目前的边缘计算更多的是指利用终端设备执行深度学习推理，而不是训练。整个训练还是采用高性能图形服务器来执行，训练完后得到模型参数，然后再部署到终端执行推理。

采用边缘计算设备，既能够利用深度学习弥补传统算法精度不高的问题，同时能够满足众多工业应用场景实时性、便捷性、小型独立性的需求。因此，边缘计算成为了现如今的主流研究方向。

1.2 Jetson Nano简介

Jetson Nano是一款体积小巧、功能强大的人工智能嵌入式开发板，于2019年3月由英伟达推出，搭载英伟达研发的128核Maxwell GPU，可以快速将AI技术落地并应用于多种边缘计算场景。其官网地址为：https://www.nvidia.com/en-us/autonomous-machines/embedded-systems/jetson-nano/。

下面详细列举一些Jetson Nano的优势：

(1) 体型小巧，性能强大，价格实惠。

(2) 专为AI而设计，搭载128核Maxwell GPU，可为机器人终端、工业视觉终端带来足够的AI算力。

(3) 可提供472 GFLOP，支持高分辨率传感器，可以并行处理多个传感器。

(4) 支持英伟达的CUDA，cuDNN和TensorRT软件库。

(5) 支持一系列流行的AI框架和算法，比如TensorFlow，PyTorch，Caffe / Caffe2，Keras，MXNet等，使得开发人员能够简单快速的将AI模型和框架集成到产品中，轻松实现图像识别，目标检测，姿势估计，语义分割，视频增强和智能分析等强大功能。

一般桌面PC或服务器级别的显卡（如英伟达1080Ti等）价格非常昂贵，不适合边缘计算需求，而且体积也过于庞大。因此，英伟达推出的这款嵌入式人工智能开发板Jetson Nano非常契合当前行业需求。最近，英伟达又推出了2G版本的jetson nano，售价只需要49美元（400元人民币左右），使得边缘计算开发成本大幅下降了。相信，未来深度学习赋能的嵌入式开发板会成为工业界的主流。

本文重点讲解如何利用Jetson Nano开发板来开发深度学习应用。关于Jetson Nano的更多介绍和基本使用方法请参考我的另一篇博客：https://blog.csdn.net/qianbin3200896/article/details/103760640，这篇文章里已经完整的讲解了Jetson Nano的基本使用方法，请不熟悉Jetson Nano的读者先阅读并且安装基本的开发环境，例如code-oss等，本文对这些内容不再重复讲解。

二、深度学习环境安装

2.1 Pytorch框架

很多教程仅仅简单的介绍如何在Jetson Nano中根据预先下载好的深度学习权重文件进行推理，例如SSD-lite等，但是这些教程没有完整的讲解如何通过训练得到这个权重？如何按照我们自己的任务需求来改造模型、训练模型？

本教程采用目前最受欢迎的Pytorch框架来实现模型训练和推理。当然，对于训练任务来说，目前更多的是采用台式电脑实现（带高端GPU显卡），这样训练速度会更快。而对于推理部分来说，就采用Jetson Nano来实现。之所以采用Pytorch是因为其实现简单，并且方便改动，同时它的动态图机制非常适合调试。另外，针对Pytorch的开源项目目前最多，基本每个应用领域都可以找到有效的开源代码，这为我们在jetson nano上开发各种应用提供了先天条件。基本上来说，只要完整的跑通一个模型（包括训练、推理、部署），后面其他的算法都可以采用这种方式实现终端AI应用开发（当然，因为资源受限，很多重量级模型均需要轻量化优化才能在Jetson Nano上运行）。

因为本教程面向pytorch，因此，需要读者有一定的pytorch基础，不熟悉pytorch框架的读者可以先自行学习pytorch。

2.2 在Jetson Nano上安装Pytorch

为了能够顺利的安装pytorch，建议使用英伟达官网推荐的方法安装指定版本的pytorch：https://forums.developer.nvidia.com/t/pytorch-for-jetson-nano-version-1-6-0-now-available/72048。该网址已经提供预先编译好的各版本二进制whl安装包：

本教程选择当前最新的Pytorch1.6.0版本。复制对应的链接地址，然后按照官方教程下载再安装，具体命令如下：

wget https://nvidia.box.com/shared/static/9eptse6jyly1ggt9axbja2yrmj6pbarc.whl -O https://nvidia.box.com/shared/static/9eptse6jyly1ggt9axbja2yrmj6pbarc.whl
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev 
pip3 install Cython
pip3 install numpy torch-1.6.0-cp36-cp36m-linux_aarch64.whl

在下载whl包时可能由于网络的原因暂时下载不了，这里可以先在普通的windows电脑上用迅雷将该whl文件下载下来，再拷贝到jetson nano中进行安装。安装完成后验证下pytorch是否安装成功：

接下来再安装torchvision,根官网介绍,pytorch1.6吻合的torch版本为0.7.0

因此，我们通过下面的命令下载和安装torch v0.7.0：

sudo apt-get install libjpeg-dev zlib1g-dev
git clone --branch v0.7.0 https://github.com/pytorch/vision torchvision
cd torchvision
export BUILD_VERSION=0.7.0
sudo python3 setup.py install 

最后验证下安装是否成功：

三、算法原理

目标检测，顾名思义，就是将图像中的感兴趣目标检测出来，这里的“检测出来”一般是指用矩形框框出目标，并给出目标所属的类型。也就是说我们既要知道感兴趣区域在图像中的具体坐标位置，同时还需要知道感兴趣区域到底是什么东西。例如在实际场景中，我们想做一个水果分拣机，那么我们就需要安装一个摄像头，然后逐帧的捕获图像，再对图像进行检测，检测图像中是否有需要分拣的水果，同时给出这个水果的类型（例如苹果自动分拣到1号流水线、橘子分拣到2号流水线等等），再调用机械设备进行分拣动作。所有这些的核心是高精度的目标检测算法作支撑。对水果作目标检测如下图所示：

考虑到Jetson Nano计算性能有限，因此，我们必须选择轻量级的目标检测算法。本文按照英伟达官方教程选择ssd-mobilenet算法来实现。

在基于“Proposal + Classification”的目标检测方法中，R-CNN 系列（R-CNN、 SPPnet、  Fast R-CNN以及 Faster R-CNN等）取得了非常好的结果，但是在速度方面离实时效果还比较远。在提高 mAP (Mean Average Precision) 的同时兼顾速度，逐渐成为神经网络目标检测领域未来的趋势。YOLO检测算法不仅能够达到实时的效果，而且mAP与前面面提到的 RCNN系列相比有很大的提升。 但是YOLO 有一些缺陷：每个网格只能预测一个物体，容易造成漏检；且对于物体的尺度相对比较敏感，面对尺度变化较大的物体时泛化能力较差。针对YOLO 中的这些不足，SSD(Single Shot MultiBox Detector)网络在这两方面都有所改进，同时兼顾了 mAP 和实时性的要求。

上图是ssd算法的模型结构图，具体的算法步骤如下：

1、输入一幅图片（300x300），将其输入到预训练好的分类网络中来获得不同大小的特征映射，修改传统的VGG16网络；

2、抽取Conv4_3、Conv7、Conv8_2、Conv9_2、Conv10_2、Conv11_2层的feature map，然后分别在这些feature map层上面的每一个点构造6个不同尺度大小的bbox，然后分别进行检测和分类，生成多个bbox，如下图所示；

3、将不同feature map获得的bbox结合起来，经过NMS（非极大值抑制）方法来抑制掉一部分重叠或者不正确的bbox，生成最终的bbox集合（即检测结果）；

为了进一步加速，nvidia给出的ssd官方例子使用mobilenet来实现优化。在原来的cnn卷积中通过使用深度可分离卷积来加速和优化，关于深度可分离卷积的原理可以参考我的另一篇博客。

总结：SSD中的Defalut box和Faster-rcnn中的anchor机制很相似。就是预设一些目标预选框，后续通过softmax分类+bounding box regression获得真实目标的位置。对于不同尺度的feature map 上使用不同大小的Default boxes。具体可以仔细参考官方代码细细品读。

四、具体实现步骤

接下来我们按照官方教程给出Jetson Nano上目标检测推理的完整步骤。这里需要采用Nvidia给出的推理库：jetson-inference ，这个库包含图像分类、检测、语义分割等模块，更加重要的是这个库实现了完整的pytorch模型到英伟达tensorrt模型的转换以及如何应用tensorrt优化器调用GPU进行高速推理。

4.1 工程代码和环境准备

我们使用Jetson Nano预装的Python3.6。首先确保Jetson Nano上已经装了最新版的JetPack 4.4（安装镜像时可以选择），里面已经装好了cuda、cudnn和TensorRT7.1，同时按照前面的步骤安装好Pytorch1.6.0。注意上述版本的一致性，尤其是TensorRT，如果用了低版本的TensorRT，那么在加载ONNX模型的时候会出问题。

环境安装好以后就可以下载工程代码：

git config --global http.postBuffer 524288000
git clone --recursive https://github.com/dusty-nv/jetson-inference.git

然后进入该工程：

cd jetson-inference

建立build文件夹：

mkdir build

进入build文件夹并编译：

cd build
cmake ../
make -j$(nproc)

编译过程中需要下载一些预训练模型，会有提示选择框。实际下载过程中由于网路原因可能下载不成功，所以我们这里不下载任何模型。后期如果需要这些预训练模型，可以打开tools/download-models.sh文件，然后找到相关模型的下载路径，在windows电脑上科学上网下载下来后再拷贝到jetson-inference/data/networks中。

安装过程中把*都勾掉，一个都不要选。然后按回车确认即可。接下来会询问是否安装pytorch，这里跳过这个步骤。

编译完成后进行安装：

sudo make install 
sudo ldconfig

安装完以后进行测试：

如果没有错误则说明安装成功。

4.2 模型训练和推理

一般情况下，完整的开发和部署流程为：

（1）在GPU服务器上使用pytorch训练模型得到pth模型文件；

  (2) 将pth模型转化为onnx格式文件；

  (3) 在jetson nano上利用tensorrt加载onnx模型，实现快速推理；

完整的ssd-mobilenet代码位于下载工程的jetson-inference/python/training/detection/ssd目录下，因此，先切换到该目录：

cd jetson-inference/python/training/detection/ssd

然后安装必要的依赖环境：

pip3 install -v -r requirements.txt

4.2.1数据集准备

本小节我们使用Open Images数据集来进行讲解。该数据集包含600种物体类，在当前ssd工程下面提供了一个脚本文件open_images_downloader.py，使用该文件可以下载需要的数据类图片。这里由于我们想要训练一个水果目标检测器，因此，选择如下8种水果图片：

Apple,Orange,Banana,Strawberry,Grape,Pear,Pineapple,Watermelon
苹果， 橘子， 香蕉，  草莓，   葡萄， 梨，   菠萝，    西瓜

具体的，运行下述命令进行下载：

python3 open_images_downloader.py --class-names "Apple,Orange,Banana,Strawberry,Grape,Pear,Pineapple,Watermelon" --data=data/fruit

最终数据会被下载到当前data目录下面。如果想要下载其它的类，可以参考种类列表。以上8类水果图片一共是6360张，含27188个检测框。

最后，我们下载一个预训练模型：

wget https://nvidia.box.com/shared/static/djf5w54rjvpqocsiztzaandq1m3avr7c.pth -O models/mobilenet-v1-ssd-mp-0_675.pth

这个预训练模型是提前在PASCAL VOC数据集上训练得到的ssd-mobilenet模型，使用这个预训练模型可以加速我们后面的模型收敛速度，同时可以在一定程度上提升后续任务的精度，这也是典型的迁移学习方法。通过这种迁移学习，即使后面我们的数据集规模不大，也能取得不错的检测精度。

4.2.2训练

运行下面的命令即可开始训练：

python3 train_ssd.py --data=data/fruit --model-dir=models/fruit --batch-size=4 --epochs=30

其中参数batch-size可以根据实际的显卡设备性能来调整。如果在jetson nano上训练，那么batch-size可以设置为4。如果训练过程中出现死机现象，说明是jetson nano的显卡内存不够了，可以尝试交换内存并且关闭图形界面来训练。对于Open Images数据集，在Jetson Nano上训练，当batch-size设置为4时，1个epoch需要将近17分钟的时间。（说明：一般情况下不推荐使用jetson nano来进行训练，因为这个终端设备的GPU性能不高，有条件的话还是选择高端的显卡加上高端的主机来进行训练，而jetson nano只用来推理）。

训练完成后训练好的模型位于models/fruit目录下面，这里直接给出训练好的模型的下载地址：https://nvidia.box.com/shared/static/gq0zlf0g2r258g3ldabl9o7vch18cxmi.gz

4.2.3模型转换

下载后的模型为pth文件，该文件是pytorch可认的深度学习参数文件，但是对于其它框架它并不可认，因此，需要将它转成一种其它深度学习框架通用的onnx格式。这里可能会有个疑问，为什么要转换文件呢，直接在jetson nano上通过pytorch进行推理就可以了。答案当然是可以的，但是这样做不好，至少在jetson nano这款终端设备上这样处理并不是最优的。因为我们通过pytorch训练出来的模型在底层是没有优化过的，直接在jetson nano上跑速度慢、显卡利用率低。为了能够最优的使用英伟达的显卡，英伟达针对深度学习的常用底层算子（例如卷积等），英伟达推出了TensorRT优化器。

TensorRT是一个高性能的深度学习推理（Inference）优化器，可以为深度学习应用提供低延迟、高吞吐率的部署推理。TensorRT可用于对超大规模数据中心、嵌入式平台或自动驾驶平台进行推理加速。TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架，将TensorRT和NVIDIA的GPU结合起来，能在几乎所有的框架中进行快速和高效的部署推理。

因此，模型转换的目的就是为了能够利用英伟达的这个tensorrt优化器进行推理加速。之前有人评测过，使用tensorrt后平均单张图片推理速度可以加速3倍以上，因此，掌握这个优化推理工具是非常重要的。

具体的，我们使用下面的命令来转换模型：

python3 onnx_export.py --model-dir=models/fruit

最终会生成一个名为ssd-mobilenet.onnx的参数文件，位于jetson-inference/python/training/detection/ssd/models/fruit目录下面。

注意，在模型转换的时候要在同目录下准备一个labels.txt文件，用来指明所有物体类：

BACKGROUND
Apple
Banana
Grape
Orange
Pear
Pineapple
Strawberry
Watermelon

其中第一个BACKGROUND即为背景类，这里也必须要将其作为1类来处理。

4.2.4 实时推理

我们采用USB免驱高清摄像头进行图像实时采集和推理。具体命令如下：

detectnet --model=models/fruit/ssd-mobilenet.onnx --labels=models/fruit/labels.txt \
          --input-blob=input_0 --output-cvg=scores --output-bbox=boxes \
          /dev/video0

第一次运行时由于需要将onnx转换为tensorrt的模型，因此，需要一定的转换时间。最终效果如下图所示：

可以看到成功的将手里的两个苹果认了出来，并且推理速度达到了24FPS，这为后面的性能优化提供了很大的改进空间，我们完全可以采用更复杂的模型来提高精度，然后降低一点实时性。

重要提示：执行AI相关操作时，一定要给Jetson Nano加装风扇来散热！！！同时Jetson Nano要采用外置适配器供电方式（最好配3A输出的电源）。

五、总结

本文详细讲解了如何在jetson nano上进行目标检测，具体包括算法原理介绍、模型训练、转换、部署，相信这篇博文能够给各位读者一定的启发。

最后，给自己刚上市的一本书打个软广：《Python Web开发从入门到实战》Django+Bootstrap，清华大学出版社出版，京东、淘宝、当当都可以买到，这本书比较注重实战，想要学习Python Web（Django）的朋友可以关注和支持一下。

完整的代码和数据集我放到百度网盘上（读者也可以自行在英伟达官方网站上下载）：

链接: https://pan.baidu.com/s/18sLaQX4AfWh_yc5tsLrPGw      提取码: te3a 

后记：关于模型原理部分本来想写仔细点、结合代码来写的，但是由于最近在忙着写一篇论文没时间弄了，后续等有机会再完善吧。

参考文献：

【1】英伟达官方教程：https://github.com/dusty-nv/jetson-inference

【2】Liu W, Anguelov D, Erhan D, et al. SSD: Single Shot MultiBox Detector[C]. european conference on computer vision, 2016: 21-37.

【3】He K, Zhang X, Ren S, et al. Deep Residual Learning for Image Recognition[C]. computer vision and pattern recognition, 2016: 770-778.

【4】Sandler M, Howard A, Zhu M, et al. MobileNetV2: Inverted Residuals and Linear Bottlenecks[C]. computer vision and pattern recognition, 2018: 4510-4520.