jetson-inference 迁移学习实战:在 Jetson 上用 PyTorch 微调 ResNet-18 与 SSD-MobileNet

发布时间:2026/9/25 11:03:20
jetson-inference 迁移学习实战:在 Jetson 上用 PyTorch 微调 ResNet-18 与 SSD-MobileNet 人工智能计算机视觉深度学习微调【免费下载链接】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点击查看免费下载迁移学习Transfer Learning是复用预训练 DNN 权重、针对自定义数据集进行微调fine-tune的技术相比从零训练大幅节省时间与算力。本指南基于 jetson-inference 的 Hello AI World 教程docs/pytorch-transfer-learning.md讲解如何在 Jetson 设备上安装 PyTorch、配置训练环境swap 与桌面 GUI并以此为入口进入 Cat/Dog、PlantCLEF 等分类数据集的完整重训练流程。读完本文你将掌握一键安装 PyTorch 与 torchvision、验证 CUDA 可用性、训练前的系统调优以及从数据下载到 TensorRT 推理的完整迁移学习链路。什么是迁移学习为什么在 Jetson 上也能训练 DNN传统 DNN 训练通常需要大规模数据集与多张独立 GPU因而一般在 PC、服务器或云实例上完成。迁移学习的思路是基于 ImageNet 预训练权重起步只针对新数据集微调网络而不是从随机初始化开始训练。由于要学习的参数量大幅减少、迭代轮数缩短即使是 Jetson 这类嵌入式平台也能胜任重新训练任务。在 jetson-inference 中迁移学习示例围绕两个网络展开ResNet-18用于图像分类/识别Classification对应本指南主体SSD-MobileNet用于目标检测Object Detection对应 docs/pytorch-ssd.md。除了这两个网络你也可以通过训练脚本的--arch参数尝试其他网络结构。训练框架使用PyTorch仓库同时提供了示例数据集、训练脚本以及基于相机的数据集采集与标注工具camera-capture见 docs/pytorch-collect.md。阅读前提在开始之前请先完成 docs/jetpack-setup-2.mdJetPack 安装与 docs/building-repo-2.md项目构建其中构建章节也包含 PyTorch 安装选项。安装 PyTorch如果按以下任一方式操作PyTorch 应已随环境预装可跳过本步使用 Docker 容器方式运行项目见 docs/aux-docker.md容器基于l4t-pytorch基础镜像构建已内置训练与迁移学习所需的 PyTorch 支持在 docs/building-repo-2.md 的 Installing PyTorch 步骤中选择了安装 PyTorch。否则可以在不使用容器的情况下手动安装$ cd jetson-inference/build $ ./install-pytorch.sh安装脚本的底层机制该脚本仓库源码见 tools/install-pytorch.sh是一个带dialog图形化菜单的交互式安装工具其核心逻辑为检查 L4T 版本读取/etc/nv_tegra_release或回退到dpkg-query查询nvidia-l4t-core得到 L4T Release/Revision 并据此确定可安装的 PyTorch 版本展示可选包根据 JetPack 版本弹出 checklist允许选择 PyTorch 对应的 Python 版本如 JetPack 4.x 的 Python 2.7/3.6下载并安装 wheel用wget下载 NVIDIA 预编译的aarch64wheel然后通过pip/pip3 install安装构建 torchvision临时移走/usr/bin/ffmpeg避免 torchvision 构建时误用git clone对应版本的 torchvision 源码后python3 setup.py install完成后恢复 ffmpeg脚本中的move_ffmpeg/restore_ffmpeg函数。脚本按 JetPack 版本映射了多套 PyTorch 版本组合见 tools/install-pytorch.sh 中install_pytorch函数平台/版本L4T / JetPackPyTorch 版本PythontorchvisionJetPack 4.2 / 4.3L4T R32.2~R32.31.4.02.7 / 3.60.5.0JetPack 4.4 DPL4T R32.4.21.4.02.7 / 3.60.5.0JetPack 4.4 GA 及以上L4T R32.4.31.6.03.60.7.0JetPack 5.xL4T R34 / R352.0可选 1.123.80.15.1JetPack 6.xL4T R362.13.100.16.1注意自动化安装工具要求JetPack 4.2 或更新版本脚本对 L4T Release 32 的设备会直接提示升级或参考 Jetson Zoo 从源码构建。若打算做目标检测训练docs/pytorch-ssd.md应使用JetPack 4.4 或更新版本并为Python 3.6安装 PyTorch。该脚本在无交互模式BUILD_INTERACTIVE ! YES下会自动跳过安装。验证 PyTorch在终端运行python或python3进入交互式 shell逐条执行以下命令确认 PyTorch 正确安装且能识别 GPU import torch print(torch.__version__) print(CUDA available: str(torch.cuda.is_available())) a torch.cuda.FloatTensor(2).zero_() print(Tensor a str(a)) b torch.randn(2).cuda() print(Tensor b str(b)) c a b print(Tensor c str(c))再验证 torchvision import torchvision print(torchvision.__version__)按 docs/pytorch-transfer-learning.md 所述JetPack 4.4 环境下 torch 版本应报告为1.6.0、torchvision 为0.7.0在更新的 JetPack5.x/6.x上则会显示脚本对应安装的版本如 2.0/2.1以实际安装为准。挂载 Swap 交换空间除 Jetson AGX Xavier 外训练会占用大量额外内存因此建议挂载 4GB swap。在 Jetson 上容器外执行以下命令先禁用 ZRAM 并创建 swap 文件sudo systemctl disable nvzramconfig sudo fallocate -l 4G /mnt/4GB.swap sudo mkswap /mnt/4GB.swap sudo swapon /mnt/4GB.swap然后向/etc/fstab末尾追加以下行使配置在重启后持久生效/mnt/4GB.swap none swap sw 0 0此后重启会自动挂载 swap。可用swapon -s或tegrastats查看使用情况。注意禁用 ZRAM内存压缩交换能释放物理内存但需要重启才能生效。训练时禁用桌面 GUI如果训练时内存紧张可临时关闭 Ubuntu 桌面 GUI释放窗口管理器与桌面占用的内存Unity/GNOME 约 800MBLXDE 约 250MB。临时关闭、跑完训练再恢复桌面$ sudo init 3 # 停止桌面 # 重新登录控制台 # 运行 PyTorch 训练脚本 $ sudo init 5 # 重启桌面如需跨重启持久化可用 systemd 修改开机默认目标$ sudo systemctl set-default multi-user.target # 开机禁用桌面 $ sudo systemctl set-default graphical.target # 开机启用桌面重启后桌面将按所设默认值保持启用或禁用。训练数据集总览迁移学习系列教程提供以下路径分类/识别ResNet-18在 Cat/Dog 数据集上重训练在 PlantCLEF 数据集上重训练采集自己的分类数据集目标检测SSD-MobileNet重训练 SSD-MobileNet采集自己的检测数据集各数据集与训练时间概览如下训练时间为 Jetson Nano 上的近似值类型数据集大小类别数训练图像每轮耗时*训练总时长**分类Cat/Dog800MB25,000~7-8 分钟~4 小时分类PlantCLEF1.5GB2010,475~15 分钟~8 小时检测Fruit2GB86,375~15 分钟~8 小时表示在 Jetson Nano 上完整遍历数据集一轮的近似时间** 表示按 35 epochs 训练模型的近似总时长。完整迁移学习链路以 Cat/Dog 分类为例下文以 docs/pytorch-cat-dog.md 为例串起数据下载 → 训练 → ONNX 导出 → TensorRT 推理的完整流程PlantCLEF 流程与此完全一致仅数据集与命令中的路径不同。下载数据教程约定将数据集存放在宿主机jetson-inference/python/training/classification/data下——这是容器启动时自动挂载的数据卷之一容器关闭后数据集不会丢失。$ cd jetson-inference/python/training/classification/data $ wget https://nvidia.box.com/shared/static/o577zd8yp3lmxf5zhm38svrbrv45am3y.gz -O cat_dog.tar.gz $ tar xvzf cat_dog.tar.gz数据集内含 5000 张训练图像、1000 张验证图像、200 张测试图像猫/狗类别各半。训练集用于迁移学习验证集在训练期间评估分类精度测试集供训练完成后验证。网络只直接接触训练集从不针对验证集与测试集训练。数据集涵盖多种犬种/猫种含老虎、美洲狮等大型猫科部分图片含人类——模型会被训练为将其作为背景忽略、聚焦猫狗内容。重训练 ResNet-18训练脚本位于仓库python/training/classification/目录本镜像中该目录内容见python/training/classification路径训练脚本train.py、onnx_export.py即存放于此。脚本不绑定特定数据集各示例数据集共用同一套 PyTorch 代码默认训练 ResNet-18可用--arch切换其他网络。$ cd jetson-inference/python/training/classification $ python3 train.py --model-dirmodels/cat_dog data/cat_dog提示训练中若内存不足或进程被 killed请先尝试挂载 SWAP和禁用桌面 GUI。为节省内存还可降低--batch-size默认 8与--workers默认 2。训练启动后控制台输出大致如下Use GPU: 0 for training dataset classes: 2 [cat, dog] using pre-trained model resnet18 reshaped ResNet fully-connected layer with: Linear(in_features512, out_features2, biasTrue) Epoch: [0][ 0/625] Time 0.932 ( 0.932) Data 0.148 ( 0.148) Loss 6.8126e-01 (6.8126e-01) Acc1 50.00 ( 50.00) Acc5 100.00 (100.00) Epoch: [0][ 10/625] Time 0.085 ( 0.163) Data 0.000 ( 0.019) Loss 2.3263e01 (2.1190e01) Acc1 25.00 ( 55.68) Acc5 100.00 (100.00)从输出可见迁移学习的核心操作using pre-trained model resnet18加载预训练权重并将 ResNet 最后的全连接层重塑为Linear(in_features512, out_features2)把 1000 类 ImageNet 分类头替换为 2 类新分类头——这正是对预训练模型进行微调的标准做法。随时可按CtrlC停止训练之后可用--resume与--epoch-start参数从中断点恢复训练不必等完整训练结束即可先测试模型。运行python3 train.py --help可查看全部可用选项包括--arch可尝试的其他网络。训练指标解读训练过程中的统计输出含义如下Epoch对数据集完整遍历一轮称为一个 epoch。Epoch: [N]表示当前第 N 轮默认训练 35 epochs可用--epochsN修改。[N/625]当前 epoch 进行到第几个图像批次。训练图像按 mini-batch 处理以提升性能默认 batch size 为 8可用--batchN设置。将括号内数字乘以 batch size 即得图像数如 batch[100/625]→ 图像[800/5000]。Time当前批次的处理时间秒。Data当前批次的磁盘加载时间秒。Loss模型累积误差期望输出 vs 预测输出。Acc1Top-1 分类精度——模型恰好预测出正确类别。Acc5Top-5 分类精度——正确类别位于模型预测的 Top 5 输出中。Cat/Dog 只有 2 个类Top-5 恒为 100%PlantCLEF 等类别数超过 5 的数据集中该指标才有区分意义。模型精度在 5000 张图像的数据集上Jetson Nano 上 ResNet-18 每轮约需 7-8 分钟训练 35 轮约 4 小时可达 80% 分类精度约在第 30 轮达到 80%第 65 轮收敛于 82.5% 精度。想进一步提升精度可扩大数据集见下文生成更多数据或尝试更复杂的模型。训练好的模型保存在jetson-inference/python/training/classification/models/cat_dog/下包括最新一轮的 checkpoint 与精度最高的最佳模型。该classification/models目录同样被自动挂载进容器容器关闭后训练成果依然保留。仓库还提供了一份训练满 100 轮的完整模型可供下载测试见 docs/pytorch-cat-dog.md 原文。将模型转换为 ONNX要让 TensorRT 加载并运行重训练后的 ResNet-18需要先把 PyTorch 模型导出为ONNX开放模型格式支持 PyTorch、TensorFlow、TensorRT 等主流框架间的模型迁移。PyTorch 内置 ONNX 导出能力执行仓库提供的onnx_export.py即可python3 onnx_export.py --model-dirmodels/cat_dog这会在jetson-inference/python/training/classification/models/cat_dog/下生成resnet18.onnx。用 TensorRT 处理图像使用imagenetC 程序或imagenet.pyPython 程序的扩展命令行参数加载自定义模型。运行命令时终端工作目录应保持在jetson-inference/python/training/classification/NETmodels/cat_dog DATASETdata/cat_dog # C imagenet --model$NET/resnet18.onnx --input_blobinput_0 --output_bloboutput_0 --labels$DATASET/labels.txt $DATASET/test/cat/01.jpg cat.jpg # Python imagenet.py --model$NET/resnet18.onnx --input_blobinput_0 --output_bloboutput_0 --labels$DATASET/labels.txt $DATASET/test/cat/01.jpg cat.jpg对狗图像同理# C imagenet --model$NET/resnet18.onnx --input_blobinput_0 --output_bloboutput_0 --labels$DATASET/labels.txt $DATASET/test/dog/01.jpg dog.jpg # Python imagenet.py --model$NET/resnet18.onnx --input_blobinput_0 --output_bloboutput_0 --labels$DATASET/labels.txt $DATASET/test/dog/01.jpg dog.jpg关键参数说明--model指定 ONNX 模型路径TensorRT 会将其构建为推理引擎--input_blob/--output_blob模型输入/输出张量名导出脚本采用input_0/output_0命名--labels类别标签文件路径。批量处理全部测试图像数据集含 200 张猫狗测试图像也可换成自己的图片。可将整个目录作为输入批量处理mkdir $DATASET/test_output_cat $DATASET/test_output_dog imagenet --model$NET/resnet18.onnx --input_blobinput_0 --output_bloboutput_0 --labels$DATASET/../labels.txt \ $DATASET/test/cat $DATASET/test_output_cat imagenet --model$NET/resnet18.onnx --input_blobinput_0 --output_bloboutput_0 --labels$DATASET/../labels.txt \ $DATASET/test/dog $DATASET/test_output_dog此处全部图像从test/目录读取结果保存到test_output/目录。图像序列的读取/保存规则详见 docs/aux-streaming.md。运行实时摄像头推理把重训练模型跑在实时视频流上# C (MIPI CSI) imagenet --model$NET/resnet18.onnx --input_blobinput_0 --output_bloboutput_0 --labels$DATASET/labels.txt csi://0 # Python (MIPI CSI) imagenet.py --model$NET/resnet18.onnx --input_blobinput_0 --output_bloboutput_0 --labels$DATASET/labels.txt csi://0支持的视频流协议csi://、/dev/video0等详见 docs/aux-streaming.md。生成更多数据可选Cat/Dog 数据集是从约 22.5GB 的 ILSVRC12 子集中通过tools/cat-dog-dataset.sh脚本随机抽取生成的。该脚本位于仓库tools/目录可用它重新生成更大、更鲁棒的数据集修改要点将解压后的ilsvrc12_subset路径填入IMAGENET_DIR变量新建空目录存放 cat_dog 数据填入OUTPUT_DIR变量通过NUM_TRAIN、NUM_VAL、NUM_TEST变量调整数据集规模。脚本会在OUTPUT_DIR下创建 train/val/test 子目录并各填充指定数量的图像。之后按前述方式重训练模型可用--resume/--epoch-start接着上次进度训练注意重训练后需重新导出 ONNX。在 PlantCLEF 数据集上训练 20 类植物模型第二组示例针对 PlantCLEF 2017 的子集约 1.5GB含 10,475 张训练图像、1,155 张验证图像、30 张测试图像覆盖 20 类植物与树木ash、beech、cat-tail、cedar、clover、cyprus、daisy、dandelion、dogwood、elm、fern、fig、fir、juniper、maple、poison ivy、sweetgum、sycamore、trout lily、tulip tree这些类别在原数据集中均有至少 500 张训练图像。下载数据$ cd jetson-inference/python/training/classification/data $ wget https://nvidia.box.com/shared/static/vbsywpw5iqy7r38j78xs0ctalg7jrg79.gz -O PlantCLEF_Subset.tar.gz $ tar xvzf PlantCLEF_Subset.tar.gz启动训练与 Cat/Dog 共用同一脚本$ cd jetson-inference/python/training/classification $ python3 train.py --model-dirmodels/plants data/PlantCLEF_Subset启动后同样会打印dataset classes: 20 [...]、using pre-trained model resnet18与reshaped ResNet fully-connected layer with: Linear(in_features512, out_features20, biasTrue)训练指标含义与前例一致详见 docs/pytorch-cat-dog.md 的训练指标章节。在 10,475 张图像上Jetson Nano 每轮约 15 分钟35 epochs 约 8 小时约在第 30 轮达到 75% Top-5 精度第 65 轮收敛于 85% Top-5 精度与 Cat/Dog 模型出现稳定与收敛的时机相近。由于 PlantCLEF 图像中植物多样且内容挑战大单张图多株重叠、大量难以区分的叶片与树干特写模型的 Top-1 精度为 55%实践中表现已相当有效。导出 ONNX 与 TensorRT 推理工作目录保持在python/training/classification/python3 onnx_export.py --model-dirmodels/plantsNETmodels/plants DATASETdata/PlantCLEF_Subset # C imagenet --model$NET/resnet18.onnx --input_blobinput_0 --output_bloboutput_0 --labels$DATASET/labels.txt $DATASET/test/cattail.jpg cattail.jpg # Python imagenet.py --model$NET/resnet18.onnx --input_blobinput_0 --output_bloboutput_0 --labels$DATASET/labels.txt $DATASET/test/cattail.jpg cattail.jpg批量处理全部测试图像与实时摄像头运行方式同前可参考 docs/pytorch-plants.md 中的 elm、juniper 等示例。注意本示例中labels.txt位于$DATASET下而批量脚本中使用了$DATASET/../labels.txt的路径写法请以实际目录结构为准。采集自己的分类数据集不想使用现成数据集时可用仓库提供的camera-capture工具从实时视频采集并标注图像。工具源码位于tools/camera-capture/目录构建后安装到aarch64/bin与/usr/local/bin/命令参数与 docs/aux-streaming.md 中列出的输入 URI 一致$ camera-capture csi://0 # 使用默认 MIPI CSI 摄像头 $ camera-capture /dev/video0 # 使用 V4L2 摄像头 /dev/video0工具会按训练脚本期望的目录结构自动创建数据集‣ train/ • class-A/ • class-B/ • ... ‣ val/ • class-A/ • class-B/ • ... ‣ test/ • class-A/ • class-B/ • ...其中class-A、class-B等子目录名与类别标签文件通常叫labels.txt中的标签名一致。标签文件每行一个类名、按字母序排列保证标签顺序与磁盘子目录顺序一致。示例 5 类labels.txtbackground brontosaurus tree triceratops velociraptor采集时先选择数据集路径并加载标签文件工具会自动创建上述目录结构在Current Class下拉框选择当前类别点Capture按钮或按空格键拍照图像按当前选择的 train/val/test 集保存到对应类子目录。若使用容器请将数据集存放在挂载目录下以免容器关闭后丢失。完整说明见 docs/pytorch-collect.md。下一步在 Cat/Dog 数据集上重训练第一个 2 类分类实战在 PlantCLEF 数据集上重训练20 类植物分类重训练 SSD-MobileNet将迁移学习扩展到目标检测采集自己的分类数据集 与 采集自己的检测数据集用相机构建自定义数据集。训练环境准备就绪后即可按上述任意路径完成数据采集 → PyTorch 微调 → ONNX 导出 → TensorRT 部署的完整闭环把 jetson-inference 的推理能力拓展到自己的业务场景。赞分享人工智能计算机视觉深度学习微调【免费下载链接】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 上基于 PlantCLEF 数据集迁移学习重训 ResNet-18 植物识别模型并部署到 TensorRT 的完整实战指南在 Jetson 上基于 PlantCLEF 数据集迁移学习重训 ResNet 18 植物识别模型并部署到 TensorRT 的完整实战指南 本篇技术指南以 j人工智能计算机视觉深度学习微调jetson-inference 实战在 NVIDIA Jetson 上用 segNet 实时运行语义分割摄像头 Demojetson inference 实战在 NVIDIA Jetson 上用 segNet 实时运行语义分割摄像头 Demo 本篇技术指南基于 jetson i人工智能计算机视觉深度学习微调jetson-inference 实战用 imagenet / imagenet.py 在 Jetson 上运行实时摄像头图像分类jetson inference 实战用 imagenet / imagenet.py 在 Jetson 上运行实时摄像头图像分类 imagenet came人工智能计算机视觉深度学习微调上一篇Winhance中文版终极Windows系统优化大师一键提升电脑性能 下一篇3大核心功能深度解析MAA明日方舟助手如何帮你节省90%游戏时间创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考