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

• 人工智能
• 计算机视觉
• 深度学习
• 微调

迁移学习（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.md（JetPack 安装）与 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图形化菜单的交互式安装工具，其核心逻辑为：

1. 检查 L4T 版本：读取/etc/nv_tegra_release（或回退到dpkg-query查询nvidia-l4t-core），得到 L4T Release/Revision 并据此确定可安装的 PyTorch 版本；
2. 展示可选包：根据 JetPack 版本弹出 checklist，允许选择 PyTorch 对应的 Python 版本（如 JetPack 4.x 的 Python 2.7/3.6）；
3. 下载并安装 wheel：用wget下载 NVIDIA 预编译的aarch64wheel，然后通过pip/pip3 install安装；
4. 构建 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函数）：

注意：自动化安装工具要求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；在更新的 JetPack（5.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 约 800MB，LXDE 约 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 上的近似值）：

• 表示在 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-dir=models/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_features=512, out_features=2, bias=True) Epoch: [0][ 0/625] Time 0.932 ( 0.932) Data 0.148 ( 0.148) Loss 6.8126e-01 (6.8126e-01) Acc@1 50.00 ( 50.00) Acc@5 100.00 (100.00) Epoch: [0][ 10/625] Time 0.085 ( 0.163) Data 0.000 ( 0.019) Loss 2.3263e+01 (2.1190e+01) Acc@1 25.00 ( 55.68) Acc@5 100.00 (100.00)

从输出可见迁移学习的核心操作：using pre-trained model 'resnet18'（加载预训练权重），并将 ResNet 最后的全连接层重塑为Linear(in_features=512, out_features=2)（把 1000 类 ImageNet 分类头替换为 2 类新分类头）——这正是对预训练模型进行微调的标准做法。

随时可按Ctrl+C停止训练；之后可用--resume与--epoch-start参数从中断点恢复训练，不必等完整训练结束即可先测试模型。运行python3 train.py --help可查看全部可用选项（包括--arch可尝试的其他网络）。

训练指标解读

训练过程中的统计输出含义如下：

• Epoch：对数据集完整遍历一轮称为一个 epoch。Epoch: [N]表示当前第 N 轮；默认训练 35 epochs，可用--epochs=N修改。
• [N/625]：当前 epoch 进行到第几个图像批次。训练图像按 mini-batch 处理以提升性能；默认 batch size 为 8，可用--batch=N设置。将括号内数字乘以 batch size 即得图像数（如 batch[100/625]→ 图像[800/5000]）。
• Time：当前批次的处理时间（秒）。
• Data：当前批次的磁盘加载时间（秒）。
• Loss：模型累积误差（期望输出 vs 预测输出）。
• Acc@1：Top-1 分类精度——模型恰好预测出正确类别。
• Acc@5：Top-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-dir=models/cat_dog

这会在jetson-inference/python/training/classification/models/cat_dog/下生成resnet18.onnx。

用 TensorRT 处理图像

使用imagenet（C++ 程序）或imagenet.py（Python 程序）的扩展命令行参数加载自定义模型。运行命令时，终端工作目录应保持在jetson-inference/python/training/classification/：

NET=models/cat_dog DATASET=data/cat_dog # C++ imagenet --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_0 --labels=$DATASET/labels.txt $DATASET/test/cat/01.jpg cat.jpg # Python imagenet.py --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_0 --labels=$DATASET/labels.txt $DATASET/test/cat/01.jpg cat.jpg

对狗图像同理：

# C++ imagenet --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_0 --labels=$DATASET/labels.txt $DATASET/test/dog/01.jpg dog.jpg # Python imagenet.py --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_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_blob=input_0 --output_blob=output_0 --labels=$DATASET/../labels.txt \ $DATASET/test/cat $DATASET/test_output_cat imagenet --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_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_blob=input_0 --output_blob=output_0 --labels=$DATASET/labels.txt csi://0 # Python (MIPI CSI) imagenet.py --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_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-dir=models/plants data/PlantCLEF_Subset

启动后同样会打印dataset classes: 20 [...]、using pre-trained model 'resnet18'与reshaped ResNet fully-connected layer with: Linear(in_features=512, out_features=20, bias=True)，训练指标含义与前例一致（详见 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-dir=models/plants

NET=models/plants DATASET=data/PlantCLEF_Subset # C++ imagenet --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_0 --labels=$DATASET/labels.txt $DATASET/test/cattail.jpg cattail.jpg # Python imagenet.py --model=$NET/resnet18.onnx --input_blob=input_0 --output_blob=output_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.txt：

background 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 的推理能力拓展到自己的业务场景。

• 人工智能
• 计算机视觉
• 深度学习
• 微调