基于Nvidia Jetson Nano边缘计算设备使用TensorRT部署YOLOv8模型实现目标检测推理

0、背景

最近拿到一台边缘计算设备，在部署YOLO模型的过程中遇到一些问题，特此记录。
在这里插入图片描述

设备介绍信息：NVIDIA Jetson Orin Nano T201Developer Kit 开发套件

开发者套件：Jetson Orin Nano T201 8GB开发套件

使用指南文档：使用指南PDF文档

YOLO模型Ultralytics官网对Nvidia Jetson设备的部署教程：https://docs.ultralytics.com/zh/guides/nvidia-jetson/

如果你的设备型号、Jetpack版本、Pytorch、torchvision和官方提供的指南文档一致，那么恭喜你，可以很方便的将YOLO模型的pt文件转为tensorRT文件直接进行推理，TensorRT转化的模型文件engine可能会比pt模型文件推理速度块快5倍左右。所以模型部署还是要进行tensorRT转化的。下面开始。

1、NVIDIA Jetson 是什么？

NVIDIA Jetson 是一系列嵌入式计算板，旨在为边缘设备带来加速 AI（人工智能）计算。这些小巧且功能强大的设备围绕NVIDIA 的GPU 架构构建，能够直接在设备上运行复杂的人工智能算法和深度学习模型，而无需依赖云计算资源。Jetson 板通常用于机器人、自动驾驶汽车、工业自动化和其他需要在本地以低延迟和高效率执行人工智能推理的应用。此外，这些板卡基于 ARM64 架构，与传统的GPU 计算设备相比，运行功耗更低。

在这里插入图片描述

2、查看设备Jetpack版本信息

使用以下命令：

jtop

如果你没有这个命令，请参考前面的设备使用指南PDF文档进行安装。
在这里插入图片描述
根据下面提升，按7是打开info信息，在info信息页可以看到当前设备各种组件版本信息。例如我的Jetpack版本是5.1.1，Ultralytics官方给的是6.2的版本安装教程，所以我的不能用，因此花费了好大功夫去安装适配我版本的pytorch等，下面一步步介绍怎么去寻找适合自己版本的安装包。
在这里插入图片描述

3、配置Cuda环境变量信息

确保 CUDA 路径已正确配置，在 ~/.bashrc 中添加以下行：

export CUDA_HOME=/usr/local/cuda
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
export PATH=/usr/local/cuda/bin:$PATH

生效配置：

source ~/.bashrc

如果以上环境变量设置成功的话，执行nvcc --version 查看cuda版本信息：
在这里插入图片描述

安装 CUDA 依赖库：

sudo apt-get install libopenblas-base libopenmpi-dev libjpeg-dev zlib1g-dev

4、安装Pytorch

根据以上输出信息，我的jetpack版本信息是5.1.1，cuda版本是11.4，系统是Ubuntu20.04 arm架构。如果直接安装pytorch官方版本我试了下是不行的。因为Jetson Nano 的 ARM 架构要求必须安装 NVIDIA 官方预编译的 PyTorch 包，而非 PyTorch 官方源的 x86 版本

所以经过我的一番辛苦查找，终于在英伟达官方找到了适配Jetson的Pytorch版本，地址为：https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform/index.html#overview__section_orin
在这里插入图片描述
在这里找到Jetson平台的Pytorch安装。
安装步骤如下：

1、Jetpack安装地址：

https://developer.nvidia.com/embedded/jetpack

2、Install system packages required by PyTorch

安装系统Pytorch依赖的包：

sudo apt-get -y update; 
sudo apt-get install -y  python3-pip libopenblas-dev;

3、安装Pytorch

根据以下命令安装Pytorch

export TORCH_INSTALL=https://developer.download.nvidia.cn/compute/redist/jp/v511/pytorch/torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl# 安装
python3 -m pip install --upgrade pip; python3 -m pip install numpy==’1.26.1’; python3 -m pip install --no-cache $TORCH_INSTALL

如果版本不对，可以切换版本，替换以下版本号

https://developer.download.nvidia.com/compute/redist/jp/v$JP_VERSION/pytorch/$PYT_VERSION

JP_VERSION：就是你的JetPack 版本，例如461就是JetPack 4.6.1，（原话：The major and minor version of JetPack you are using, such as 461 for JetPack 4.6.1 or 50 for JetPack 5.0.
）

PYT_VERSION：The released version of the PyTorch wheels, as given in the Compatibility Matrix.

其他不清楚的可以直接看官方文档。

验证Pytorch安装

运行以下代码确保 CUDA 可用：

import torch
print("PyTorch 版本:", torch.__version__)
print("CUDA 可用:", torch.cuda.is_available())
print("GPU 设备:", torch.cuda.get_device_name(0))

到这里应该CUDA就是可用状态了。

在这里插入图片描述

5、安装torchvision

安装了Pytorch之后还要安装torchvision，因为我们安装的pytorch是Nvidia官方预编译的，直接按照Pytorch官方源的torchvision肯定不行，会报错版本不匹配。例如：
我直接安装：pip install torchvision==0.15.1，就会报错：

Couldn’t load custom C++ ops. This can happen if your PyTorch and torchvision versions are incompatible, or if you had errors while compiling torchvision from source. For further information on the compatible versions, check https://github.com/pytorch/vision#installation for the compatibility matrix. Please check your PyTorch version with torch.version and your torchvision version with torchvision.version and verify if they are compatible, and if not please reinstall torchvision so that it matches your PyTorch install.

所以还得安装指定版本的torchvision

重新安装 torchvision

由于当前版本可能存在环境污染或编译问题，重新安装：

pip uninstall torchvision -y  # 卸载现有版本
git clone --branch v0.15.1 https://github.com/pytorch/vision
cd vision
python3 setup.py install --user  # 从源码编译安装

注意事项：

使用 --user 参数避免权限问题。
确保 PyTorch 已正确安装且版本与 torchvision 0.15.1 兼容。

这个编译过程可能会有点长，耐心等待一会儿。

验证安装

检查 torchvision 版本及功能：

import torchvision
print(torchvision.__version__)  # 预期输出: 0.15.1

到此，应该就算环境安装完成了，下面就可以执行YOLO模型的转换工作了。

6、在NVIDIA Jetson 上使用TensorRT

6.1 将模型转换为TensorRT 并运行推理

PyTorch 格式的 YOLO11n 模型被转换为TensorRT ，以便使用导出的模型进行推理。

还需要先安装 pip install onnx
然后执行转换：

# Export a YOLO11n PyTorch model to TensorRT format
yolo export model=yolo11n.pt format=engine  # creates 'yolo11n.engine'# Run inference with the exported model
yolo predict model=yolo11n.engine source='https://ultralytics.com/images/bus.jpg'