Jetson Nano 模型训练和推理流程

深度学习的训练和推理流程，是先采用高性能图形服务器使用深度学习框架来训练（Training）机器学习算法，研究大量的数据来学习一个特定的场景，完成后得到模型参数，再部署到终端执行机器学习推理（Inference），以训练好的模型从新数据中得出结论。

一般的深度学习项目，训练时为了加快速度，会使用多 GPU 分布式训练。但在部署推理时，为了降低成本，往往使用单个 GPU 机器甚至嵌入式平台进行部署。Jetson Nano 可以完成整个训练和推理流程，但基于 Jetson Nano 的低算力，不推荐在 Jetson Nano 上进行复杂训练，而仅用于推理。

• 在GPU服务器上使用pytorch训练模型得到pth模型文件；
• 将pth模型转化为onnx格式文件；
• Jetson Nano上利用tensorrt加载onnx模型，实现快速推理。

以下是根据这篇文章的摘要，介绍使用 ssd-mobilenet 实现目标检测网络的训练、部署和推理。

环境准备

前提：基于 Jetson Inference 完整可运行状态，见《Jetson Nano 2GB 深入 Jetson Inference 项目》。

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

# 安装必要的依赖环境boto3、pandas和urllib3
pip3 install -v -r requirements.txt

数据集准备

在 ssd-mobilenet 工程中提供了一个脚本文件 open_images_downloader.py，使用该文件可以从Open Image 图像数据集下载需要的数据类图片。我们想要训练一个水果目标检测器，因此，选择下载8种水果图片（更多各类），一共是6360张图片（Image count: 6360），含27188个检测框（Bounding box count: 27188）。https://github.com/dusty-nv/pytorch-ssd/blob/master/open_images_classes.txt

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

Google Open Image 图像数据集（Open Images Dataset V6 + Extensions）链接了大约 600 个类别（包括扩展为 6000 类）的 900 万张图像，比同类 ImageNet 图像数据集（1000 个类别）包含更多贴近实际生活的实体，足够保证从头训练一个深度网络模型。

ILSVRC（ImageNet Large Scale Visual Recognition Challenge）是近年来机器视觉领域最受追捧也是最具权威的学术竞赛之一，代表了图像领域的最高水平。ImageNet 数据集是ILSVRC竞赛使用的是数据集，由斯坦福大学李飞飞教授主导，包含了超过1400万张全尺寸的有标记图片。

再下载一个预训练模型：

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

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

训练

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

在 Jetson Nano 上训练，batch-size 设置为 4，1个 epoch 需要近 17 分钟的时间。训练完成后训练好的模型位于 models/fruit 目录。

训练好的模型下载地址：
https://nvidia.box.com/shared/static/gq0zlf0g2r258g3ldabl9o7vch18cxmi.gz

模型转换

训练或下载后的模型为pth文件，该文件是pytorch可认的深度学习参数文件，但是对于其它框架它并不可认，因此，需要将它转成一种其它深度学习框架通用的onnx格式。

虽然在 Jetson Nano 直接通过 pytorch 进行推理，但通过 pytorch 训练出来的模型（pth文件）在底层没有优化，占用 Jetson Nano 资源过多。英伟达推出了 TensorRT 优化器，可以为深度学习应用提供低延迟、高吞吐率的部署推理加速。据评测，使用 tensorrt 后平均单张图片推理速度可以加速3倍以上。

TensorRT 现已能支持 TensorFlow、Caffe、Mxnet、Pytorch 等几乎所有的深度学习框架，将 TensorRT 和 NVIDIA 的 GPU 结合起来，能在几乎所有的框架中进行快速和高效的部署推理。

在 models/fruit 目录下准备一个 labels.txt 文件，用来指明所有物体类：

BACKGROUND
Apple
Banana
Grape
Orange
Pear
Pineapple
Strawberry

开始转换模型：

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

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

实时推理

采用USB免驱高清摄像头进行图像实时采集和推理。如使用 CSI 摄像头，/dev/video0 换为 csi://0

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的模型，因此，需要一定的转换时间。

via - https://blog.csdn.net/qianbin3200896/article/details/108949723

---

Jetson Nano 的资源分享：

链接：https://pan.baidu.com/s/1E95CKisbCLTGKGI6c-JydA
提取码：culm

本文由 CulmartPlay 整理发布，参考 CC-BY-SA 3.0 协议共享，欢迎转载、引用或改编。