当前位置：首页 > news >正文

【jetson】训练图形分类推理模型

news 2026/6/4 23:04:21

一，安装pytorch

1）安装准备

jetson环境： JetPack 5.1.6
PyTorch wheel包：torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl
torchvision源码包：torchvision.tar.gz
PyTorch、torchvision 与 JetPack 版本需保持匹配
PyTorch wheel包详解：
文件命名解析：torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl
- torch-2.0.0：PyTorch主版本号，此为2.0.0。
- +nv23.05：NVIDIA为Jetson平台定制的版本标识，包含针对Jetson的CUDA、TensorRT等优化。
- cp38-cp38：表示该wheel包适用于Python 3.8（cp38指CPython 3.8）。这是JetPack 5.1.6默认的Python版本，必须严格匹配。
- linux_aarch64：平台标识，代表适用于Linux操作系统、ARM64（aarch64）架构，即Jetson设备的CPU架构。
获取方式：此wheel包通常需要从NVIDIA官方论坛、开发者网站或JetPack配套资源中下载。确保来源可靠，避免版本不兼容。
版本匹配的重要性：JetPack 5.1.6 包含了特定版本的CUDA、cuDNN、TensorRT等底层库。+nv23.05后缀确保了PyTorch与这些底层驱动和库完全兼容。若使用不匹配的版本（如通用x86_64版本或其他JetPack版本的wheel），可能导致安装失败、CUDA不可用或运行时崩溃。
文件存放建议：下载后，建议将其放在用户目录下的Downloads文件夹（如/home/guest/Downloads/），便于后续安装命令引用。

2）安装PyTorch

执行以下命令安装PyTorch：

pip3install~/Downloads/torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl

3）安装torchvision

进入下载目录并解压源码：

tar-xzftorchvision.tar.gzcdtorchvision# 确认解压后的文件夹名称#安装必要依赖sudoapt-getupdate&&sudoapt-getinstall-y\libjpeg-dev\zlib1g-dev\libpython3-dev\libavcodec-dev\libavformat-dev\libswscale-dev\libopenblas-dev\libopenblas-base#设置环境变量并开始编译：过程较长可能较慢，耐心等待exportBUILD_VERSION=0.15.1 python3 setup.pyinstall--user#验证安装python3-c"import torch; print('PyTorch:', torch.__version__, 'CUDA:', torch.cuda.is_available()); import torchvision; print('TorchVision:', torchvision.__version__)"

4）验证安装

执行以下命令验证安装是否成功：

#执行python3-c"import torch; print('PyTorch:', torch.__version__, 'CUDA:', torch.cuda.is_available()); import torchvision; print('TorchVision:', torchvision.__version__)"#预期输出PyTorch:2.0.0+nv23.05 CUDA: True TorchVision:0.15.1

二，训练推理模型

1）训练流程与命令执行

准备训练图片集cat_dog 数据集放到如下目录：

pwd/home/guest/jetson-inference/python/training/classification/data

查看目录 ~/jetson-inference/python/training/classification 下的核心文件，确认训练脚本、数据集、模型保存目录齐全
- train.py：核心训练脚本
- data/cat_dog：猫狗分类数据集目录
- models/cat_dog：模型权重保存目录
训练命令

# --model-dir：训练完成的模型权重保存路径# data/cat_dog：训练数据集的路径python3 train.py --model-dir=models/cat_dog data/cat_dog

2）报错记录

报错信息：

报错原因：缺少 PyTorch 可视化工具 tensorboard 依赖，
此依赖安装过程耗时较长，耐心等待；安装完成后，重新执行训练命令即可

# 1. 更新软件源sudoaptupdate# 2. 安装apt预编译的grpcio（避免源码编译，大幅节省时间）sudoaptinstallpython3-grpcio-y# 3. 安装tensorboard依赖pip3installtensorboard

3）模型训练过程与结果解析

Val:[90/125]Time0.035(0.048)Loss4.4410e-01(4.3245e-01)Accuracy87.500(77.335)Val:[100/125]Time0.036(0.047)Loss6.3032e-01(4.2943e-01)Accuracy75.000(77.723)Val:[110/125]Time0.036(0.046)Loss3.8643e-01(4.2486e-01)Accuracy87.500(78.491)Val:[120/125]Time0.036(0.045)Loss2.3032e-01(4.2406e-01)Accuracy100.000(78.926)Val:[124/125]Time0.034(0.045)Loss3.8653e-01(4.2849e-01)Accuracy100.000(78.900)=>Epoch34* Train Loss5.0718e-01 * Train Accuracy74.0400* Val Loss4.2849e-01 * Val Accuracy78.9000saved checkpoint to: models/cat_dog/checkpoint.pth.tar#训练指标解析#Epoch 34：模型完成了 34 轮完整的数据集训练#Train Loss/Acc：训练集损失值、训练集准确率（74.04%）#Val Loss/Acc：验证集损失值、验证集准确率（78.9%）#模型保存：训练完成的权重文件自动保存至 models/cat_dog/checkpoint.pth.tar

训练结果说明

验证集准确率为 78.9%，模型具备猫狗图像分类能力
训练日志无报错、模型权重正常保存，代表训练流程成功

三，测试模型

基于jetson-inference 模型训练→ ONNX导出→ 推理运行完整步骤

1）训练完成后，导出ONNX模型文件

模型训练结束后，需要将训练好的模型转换为ONNX格式，用于后续推理。

cd~/jetson-inference/python/training/classification

执行ONNX导出脚本

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

报错解决（缺少onnx库）
执行导出时若提示ModuleNotFoundError: No module named 'onnx'，先安装依赖：

# 安装onnx库pip3install--useronnx# 验证安装是否成功python3-c"import onnx; print(onnx.__version__)"

重新执行导出命令，即可完成ONNX模型转换

模型训练完成后，models 目录下会自动生成以下文件：

2）运行ONNX模型进行图像分类

使用imagenet程序加载转换后的ONNX模型，对测试图片完成分类推理。

cd~/jetson-inference/build/aarch64/bin

执行分类推理命令

./imagenet\--model=/home/guest/jetson-inference/python/training/classification/models/cat_dog/resnet18.onnx\--input_blob=input_0\--output_blob=output_0\--labels=/home/guest/jetson-inference/python/training/classification/models/cat_dog/labels.txt\/home/guest/Downloads/test.jpg\/home/guest/Downloads/output.jpg