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

一安装pytorch1安装准备jetson环境 JetPack 5.1.6PyTorch wheel包torch-2.0.0nv23.05-cp38-cp38-linux_aarch64.whltorchvision源码包torchvision.tar.gzPyTorch、torchvision 与 JetPack 版本需保持匹配PyTorch wheel包详解文件命名解析torch-2.0.0nv23.05-cp38-cp38-linux_aarch64.whltorch-2.0.0PyTorch主版本号此为2.0.0。nv23.05NVIDIA为Jetson平台定制的版本标识包含针对Jetson的CUDA、TensorRT等优化。cp38-cp38表示该wheel包适用于Python 3.8cp38指CPython 3.8。这是JetPack 5.1.6默认的Python版本必须严格匹配。linux_aarch64平台标识代表适用于Linux操作系统、ARM64aarch64架构即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执行以下命令安装PyTorchpip3install~/Downloads/torch-2.0.0nv23.05-cp38-cp38-linux_aarch64.whl3安装torchvision进入下载目录并解压源码tar-xzftorchvision.tar.gzcdtorchvision# 确认解压后的文件夹名称#安装必要依赖sudoapt-getupdatesudoapt-getinstall-y\libjpeg-dev\zlib1g-dev\libpython3-dev\libavcodec-dev\libavformat-dev\libswscale-dev\libopenblas-dev\libopenblas-base#设置环境变量并开始编译过程较长可能较慢耐心等待exportBUILD_VERSION0.15.1 python3 setup.pyinstall--user#验证安装python3-cimport torch; print(PyTorch:, torch.__version__, CUDA:, torch.cuda.is_available()); import torchvision; print(TorchVision:, torchvision.__version__)4验证安装执行以下命令验证安装是否成功#执行python3-cimport torch; print(PyTorch:, torch.__version__, CUDA:, torch.cuda.is_available()); import torchvision; print(TorchVision:, torchvision.__version__)#预期输出PyTorch:2.0.0nv23.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-dirmodels/cat_dog data/cat_dog2报错记录报错信息报错原因缺少 PyTorch 可视化工具 tensorboard 依赖此依赖安装过程耗时较长耐心等待安装完成后重新执行训练命令即可# 1. 更新软件源sudoaptupdate# 2. 安装apt预编译的grpcio避免源码编译大幅节省时间sudoaptinstallpython3-grpcio-y# 3. 安装tensorboard依赖pip3installtensorboard3模型训练过程与结果解析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-dirmodels/cat_dog报错解决缺少onnx库执行导出时若提示ModuleNotFoundError: No module named onnx先安装依赖# 安装onnx库pip3install--useronnx# 验证安装是否成功python3-cimport 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_blobinput_0\--output_bloboutput_0\--labels/home/guest/jetson-inference/python/training/classification/models/cat_dog/labels.txt\/home/guest/Downloads/test.jpg\/home/guest/Downloads/output.jpg使用前要校验模型文件、测试图片等资源的目录路径保证路径无误查看结果打开输出图片/home/guest/Downloads/output.jpg即可查看分类结果总结导出ONNX核心是安装onnx依赖 运行onnx_export.py脚本推理用imagenet程序加载ONNX模型指定输入输出节点、标签文件即可完成分类模型当前置信度不足需进一步优化训练