Jetson 上基于 PyTorch 重训 SSD-Mobilenet 目标检测模型:从 Open Images 数据下载到 TensorRT 部署全流程 人工智能计算机视觉深度学习微调【免费下载链接】jetson-inferenceHello AI World guide to deploying deep-learning inference networks and deep vision primitives with TensorRT and NVIDIA Jetson.项目地址https://gitcode.com/gh_mirrors/je/jetson-inference点击查看免费下载本指南以 jetson-inference 仓库中的 pytorch-ssd.md 为核心脉络完整讲解在 JetsonJetPack 4.4含 TensorRT 7.1上用 PyTorch 对 SSD-Mobilenet 做迁移学习transfer learning重训自定义目标检测模型的端到端流程包括训练环境准备、Open Images 数据集下载与裁剪、train_ssd.py训练、ONNX 导出以及最终用detectnet在 TensorRT 上对静态图片和摄像头实时流进行推理。读完本文你将掌握一套从自定义类别数据集到Jetson 上实时部署的完整可复现方案。背景为什么选择 SSD-Mobilenet 迁移学习SSD-Mobilenet 是移动端与嵌入式设备上非常流行的实时目标检测架构它把 SSD-300Single-Shot MultiBox Detector单阶段多框检测器与 MobileNet 骨干网络结合兼顾检测精度与推理速度非常适合 Jetson 这类算力受限的板卡。在本仓库的示例中目标是在 Open Images 数据集上训练一个能定位8 种水果Apple、Orange、Banana、Strawberry、Grape、Pear、Pineapple、Watermelon的自定义检测模型例如服务于水果采摘机器人。你完全可以替换成 Open Images 600 多个类别中的任意组合来训练自己的模型。版本前提重要训练环境需要JetPack 4.4 或更新版本并已为Python 3安装 PyTorch。JetPack 4.4 内置 TensorRT 7.1是能通过 ONNX 加载 SSD-Mobilenet 的最低 TensorRT 版本更新版本的 TensorRT 同样兼容。PyTorch 的安装方式见 pytorch-transfer-learning.md可以使用仓库 build 目录下的install-pytorch.sh脚本或直接使用 Docker 容器。训练环境准备SetupSSD-Mobilenet 的训练代码位于仓库python/training/detection/ssd目录下该目录在镜像仓库中对应上游的pytorch-ssd项目。如果你没有使用 Docker 容器需要手动完成两步准备# 仅在不使用容器时才需要执行 $ cd jetson-inference/python/training/detection/ssd $ wget https://nvidia.box.com/shared/static/djf5w54rjvpqocsiztzaandq1m3avr7c.pth -O models/mobilenet-v1-ssd-mp-0_675.pth $ pip3 install -v -r requirements.txt第一步下载基础模型权重mobilenet-v1-ssd-mp-0_675.pth到ssd/models目录。这个基础模型已经在另一个数据集PASCAL VOC上预训练过所以我们无需从零训练 SSD-Mobilenet那样会慢得多而是利用迁移学习对其微调fine-tune让它学会检测我们自定义的新类别。第二步安装训练所需的若干 Python 包如果使用容器则已预装。note:前提仍是 Jetson 上已具备 JetPack 4.4 以及针对 Python 3 安装的 PyTorch。内存优化建议训练会消耗大量额外内存。除了 Jetson AGX Xavier 之外建议挂载 4GB 交换空间并可在训练期间临时关闭桌面 GUI 释放约 800MBUnity/GNOME或约 250MBLXDE内存。具体操作详见 pytorch-transfer-learning.md 的 Mounting Swap 与 Disabling the Desktop GUI 两节先禁用 ZRAM 并创建/mnt/4GB.swap再把/mnt/4GB.swap none swap sw 0 0写入/etc/fstab使其开机自动生效训练期间可用sudo init 3关桌面、训练完sudo init 5恢复。下载训练数据Open Images 数据集Open Images 数据集包含超过 600 个物体类别可通过脚本open_images_downloader.py自动下载你指定的类别。note:使用的类别越少模型推理时运行得越快。同时Open Images 某些类别尤其是人和车辆类数据量极大250GB下载前务必先阅读下面限制数据量一节。本文示例选用的类别为Apple,Orange,Banana,Strawberry,Grape,Pear,Pineapple,Watermelon水果类共约 6500 张图规模适中。下载命令如下$ python3 open_images_downloader.py --class-names Apple,Orange,Banana,Strawberry,Grape,Pear,Pineapple,Watermelon --datadata/fruit ... 2020-07-09 16:20:42 - Starting to download 6360 images. 2020-07-09 16:20:42 - Downloaded 100 images. 2020-07-09 16:20:42 - Downloaded 200 images. ... 2020-07-09 16:32:12 - Task Done.默认情况下数据集会下载到jetson-inference/python/training/detection/ssd下的data/目录该目录会被自动挂载进 Docker 容器见 aux-docker.md 的 Mounted Data Volumes你也可以通过--dataPATH修改存储位置。数据量较大时可能有必要使用外部存储如果下载多套数据集应各自存放在独立子目录中。限制数据量Limiting the Amount of Data在选定自己的类别后、真正下载前强烈建议先用--stats-only参数运行下载脚本——它只统计各类别图片数量会下载约 1GB 的标注数据但不下图片不会实际下载图像$ python3 open_images_downloader.py --stats-only --class-names Apple,Orange,Banana,Strawberry,Grape,Pear,Pineapple,Watermelon --datadata/fruit ... 2020-07-09 16:18:06 - Total available images: 6360 2020-07-09 16:18:06 - Total available boxes: 27188 ------------------------------------- train set statistics ------------------------------------- Image count: 5145 Bounding box count: 23539 Bounding box distribution: Strawberry: 7553/23539 0.32 Orange: 6186/23539 0.26 Apple: 3622/23539 0.15 Grape: 2560/23539 0.11 Banana: 1574/23539 0.07 Pear: 757/23539 0.03 Watermelon: 753/23539 0.03 Pineapple: 534/23539 0.02 ... ------------------------------------- Overall statistics ------------------------------------- Image count: 6360 Bounding box count: 27188从统计可见水果数据集各类别分布并不均衡Strawberry 占比 0.32而 Pineapple 仅 0.02这也是后续需要用数据限制参数平衡的原因之一。实践中为控制训练时间与磁盘占用建议把总图片数控制在 10K 以内图片越多模型越准但训练越久。可以用以下两个参数限制下载量参数作用--max-images把整个数据集限制到指定图片数同时尽量保持各类别图片分布比例与原始数据集一致某类图片多比例仍大致不变--max-annotations-per-class把每个类别限制到指定的边界框数量若某类本身少于该数量则全部使用。适用于类别间数据分布不均衡的场景例如只想要 2500 张水果图片$ python3 open_images_downloader.py --max-images2500 --class-names Apple,Orange,Banana,Strawberry,Grape,Pear,Pineapple,Watermelon --datadata/fruit若既不设--max-images也不设--max-annotations-per-class默认会下载全部可用数据——所以下载前务必先用--stats-only查看数据量。虽然无法事先精确计算图片占用的磁盘空间但该数据集的经验估值约为每张图 ~350KB水果数据集约 2GB。训练性能参考下表为 SSD-Mobilenet 的近似训练吞吐在水果数据集上实测5145 张训练图、batch size 4可用于估算训练耗时设备吞吐Images/sec每轮耗时*Jetson Nano4.7717 分 55 秒Jetson Xavier NX14.655 分 50 秒*注在水果数据集5145 张训练图、batch size 4上测得。训练 SSD-Mobilenet 模型数据下载完成后运行train_ssd.py启动训练python3 train_ssd.py --datadata/fruit --model-dirmodels/fruit --batch-size4 --epochs30note:如果训练过程中内存不足或进程被 killed尝试 挂载 SWAP 并 关闭桌面 GUI。为节省内存还可以调小--batch-size默认 4与--workers默认 2。常用训练参数参数默认值说明--datadata/数据集所在位置--model-dirmodels/输出训练模型检查点checkpoint的目录--resume无已有检查点路径用于从中断处继续训练--batch-size4可依据可用内存酌情调大--epochs30理想情况下可到 100但会相应延长训练时间--workers2数据加载器线程数0 关闭多线程训练过程中应能看到 loss 逐步下降2020-07-10 13:14:12 - Epoch: 0, Step: 10/1287, Avg Loss: 12.4240, Avg Regression Loss 3.5747, Avg Classification Loss: 8.8493 2020-07-10 13:14:12 - Epoch: 0, Step: 20/1287, Avg Loss: 9.6947, Avg Regression Loss 4.1911, Avg Classification Loss: 5.5036 2020-07-10 13:14:13 - Epoch: 0, Step: 30/1287, Avg Loss: 8.7409, Avg Regression Loss 3.4078, Avg Classification Loss: 5.3332 ... 2020-07-10 13:19:26 - Epoch: 0, Validation Loss: 5.6730, Validation Regression Loss 1.7096, Validation Classification Loss: 3.9634 2020-07-10 13:19:26 - Saved model models/fruit/mb1-ssd-Epoch-0-Loss-5.672993580500285.pth日志中 Avg Loss 由 Avg Regression Loss回归损失负责边界框坐标回归与 Avg Classification Loss分类损失负责类别置信度两部分构成可以直观判断模型在框得准不准与认得对不对两个维度上的收敛情况。如果不想等全部 epoch 跑完可以随时按CtrlC终止训练之后用--resumeCHECKPOINT从断点继续。仓库还提供了已训练 100 epoch 的水果模型权重包位于 NVIDIA 官方分享链接原文档中可直接下载。将模型转换为 ONNX训练完成后需要把 PyTorch 模型转成 ONNX 格式才能交给 TensorRT 加载推理python3 onnx_export.py --model-dirmodels/fruit执行后会在jetson-inference/python/training/detection/ssd/models/fruit/下生成名为ssd-mobilenet.onnx的模型文件。用 TensorRT 处理静态图片接下来用detectnetC 程序或detectnet.pyPython 版加载自定义 SSD-Mobilenet ONNX 模型。运行命令时终端工作目录应保持在jetson-inference/python/training/detection/ssd/IMAGESpath-to-your-jetson-inference/data/images # 替换为你的 jetson-inference 实际路径 detectnet --modelmodels/fruit/ssd-mobilenet.onnx --labelsmodels/fruit/labels.txt \ --input-blobinput_0 --output-cvgscores --output-bboxboxes \ $IMAGES/fruit_*.jpg $IMAGES/test/fruit_%i.jpgnote:可以把上面的detectnet替换为detectnet.py来运行 Python 版本。几个关键参数说明参数含义说明--model模型文件指向导出的ssd-mobilenet.onnx--labels标签文件models/fruit/labels.txt每行一个类别名与训练类别顺序一致--input-blobinput_0输入层名称覆盖 detectNet 默认输入层名data--output-cvgscores置信度输出层覆盖 detectNet 默认输出层名coverage--output-bboxboxes边界框输出层覆盖 detectNet 默认输出层名bboxes这里需要对默认参数做一点源码级说明在 c/detectNet.h 中detectNet 面向传统 Caffe 模型定义了默认的输入/输出 blob 名——输入层默认dataDETECTNET_DEFAULT_INPUT输出层默认coverageDETECTNET_DEFAULT_COVERAGE与bboxesDETECTNET_DEFAULT_BBOX默认置信度阈值 DETECTNET_DEFAULT_CONFIDENCE_THRESHOLD 为 0.5。而 ONNX 导出的 SSD-Mobilenet 模型 blob 命名不同input_0、scores、boxes因此必须通过--input-blob、--output-cvg、--output-bbox显式指定否则 detectNet 将按默认 blob 名找不到对应层而无法加载。程序会把检测结果绘制到输出图片并保存到$IMAGES/test目录文件名中的%i会被替换为序号。运行实时摄像头程序同样地可以把重训后的模型跑在摄像头或视频流上detectnet --modelmodels/fruit/ssd-mobilenet.onnx --labelsmodels/fruit/labels.txt \ --input-blobinput_0 --output-cvgscores --output-bboxboxes \ csi://0其中csi://0是 Jetson 板载 CSI 摄像头。其他摄像头/视频源V4L2 摄像头、RTSP 流、视频文件等的使用方式详见 Camera Streaming and Multimedia。至此你就完成了从 Open Images 自定义数据集、PyTorch 迁移学习训练、ONNX 导出到 Jetson 上 TensorRT 静态图片与实时摄像头推理的完整闭环。后续可以继续阅读 pytorch-collect-detection.md收集你自己的检测数据集 来采集并标注完全自定义的数据进一步摆脱对公共数据集的依赖。赞分享人工智能计算机视觉深度学习微调【免费下载链接】jetson-inferenceHello AI World guide to deploying deep-learning inference networks and deep vision primitives with TensorRT and NVIDIA Jetson.项目地址https://gitcode.com/gh_mirrors/je/jetson-inference点击查看免费下载相关推荐基于 MMDetection 在 Open Images 大规模数据集上训练与评估目标检测模型基于 MMDetection 在 Open Images 大规模数据集上训练与评估目标检测模型 导读 本文以 configs/openimages/README人工智能计算机视觉深度学习模型评测Open Images数据集实战指南从下载到模型训练全流程Open Images数据集实战指南从下载到模型训练全流程 数据集概览 Open Images数据集是Google推出的超大规模计算机视觉数据集包含约900数据集计算机视觉Open Images数据集实战指南从下载到模型训练全流程Open Images数据集实战指南从下载到模型训练全流程 数据集概览 Open Images数据集是Google推出的超大规模计算机视觉数据集包含约900数据集计算机视觉上一篇MaxKey贡献者核心开发团队介绍下一篇Ink配置验证输入数据的完整性检查创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考