Skip to main content

在 NVIDIA Jetson 上使用 TensorRT 部署 YOLOv8

note

本文档由 AI 翻译。如您发现内容有误或有改进建议,欢迎通过页面下方的评论区,或在以下 Issue 页面中告诉我们:https://github.com/Seeed-Studio/wiki-documents/issues

本指南介绍了如何将 YOLOv8 模型部署到 NVIDIA Jetson 平台,并使用 TensorRT 进行推理。在这里,我们使用 TensorRT 来最大化 Jetson 平台上的推理性能。

本指南将介绍不同的计算机视觉任务,例如:

  • 目标检测
  • 图像分割
  • 图像分类
  • 姿态估计
  • 目标跟踪

前置条件

  • Ubuntu 主机电脑(原生或使用 VMware Workstation Player 的虚拟机)
  • reComputer Jetson 或任何运行 JetPack 5.1.1 或更高版本的 NVIDIA Jetson 设备
note

本指南已在 reComputer J4012 和 reComputer 工业版 J4012 [https://www.seeedstudio.com/reComputer-Industrial-J4012-p-5684.html](由 NVIDIA Jetson Orin NX 16GB 模块提供支持)上测试和验证。

将 JetPack 刷写到 Jetson

现在,您需要确保 Jetson 设备已刷写 JetPack 系统。您可以使用 NVIDIA SDK Manager 或命令行将 JetPack 刷写到设备上。

有关 Seeed Jetson 支持设备的刷写指南,请参考以下链接:

note

请确保刷写 JetPack 版本 5.1.1,因为这是我们为本指南验证的版本。

一行代码将 YOLOv8 部署到 Jetson!

在您使用 JetPack 刷写 Jetson 设备后,只需运行以下命令即可运行 YOLOv8 模型。这将首先下载并安装必要的软件包、依赖项,设置环境,并从 YOLOv8 下载预训练模型以执行目标检测、图像分割、姿态估计和图像分类任务!

wget files.seeedstudio.com/YOLOv8-Jetson.py && python YOLOv8-Jetson.py
note

上述脚本的源代码可以在 这里 找到。

使用预训练模型

使用 YOLOv8 提供的预训练模型是快速入门的最佳方式。然而,这些是 PyTorch 模型,因此在 Jetson 上推理时仅能使用 CPU。如果您希望在 Jetson 上运行这些模型时利用 GPU 获得最佳性能,可以按照本 Wiki 的本节内容将 PyTorch 模型导出为 TensorRT。

YOLOv8 提供了 5 个用于目标检测的预训练 PyTorch 模型权重,这些模型在 COCO 数据集上以 640x640 的输入图像尺寸进行训练。您可以在下方找到它们的信息:

模型尺寸
(像素)
mAPval
50-95
速度
CPU ONNX
(毫秒)
速度
A100 TensorRT
(毫秒)
参数
(百万)
FLOPs
(十亿)
YOLOv8n64037.380.40.993.28.7
YOLOv8s64044.9128.41.2011.228.6
YOLOv8m64050.2234.71.8325.978.9
YOLOv8l64052.9375.22.3943.7165.2
YOLOv8x64053.9479.13.5368.2257.8

参考文档:https://docs.ultralytics.com/tasks/detect

您可以从上表中选择并下载所需的模型,然后执行以下命令对图像进行推理:

yolo detect predict model=yolov8n.pt source='https://ultralytics.com/images/bus.jpg' show=True

model 参数中,您可以更改为 yolov8s.pt、yolov8m.pt、yolov8l.pt 或 yolov8x.pt,它将下载相应的预训练模型。

您还可以连接一个摄像头并执行以下命令:

yolo detect predict model=yolov8n.pt source='0' show=True
note

如果在执行上述命令时遇到任何错误,请尝试在命令末尾添加 "device=0"

note

上述示例运行在 reComputer J4012/ reComputer Industrial J4012 上,使用了 YOLOv8s 模型,输入尺寸为 640x640,并采用 TensorRT FP16 精度。


使用 TensorRT 提高推理速度

如前所述,如果您希望在运行 YOLOv8 模型的 Jetson 上提高推理速度,首先需要将原始 PyTorch 模型转换为 TensorRT 模型。

按照以下步骤将 YOLOv8 PyTorch 模型转换为 TensorRT 模型。

note

这适用于我们之前提到的所有四种计算机视觉任务。

  • 步骤 1. 通过指定模型路径执行导出命令
yolo export model=<path_to_pt_file> format=engine device=0

例如:

yolo export model=yolov8n.pt format=engine device=0
note

如果遇到关于 cmake 的错误,可以忽略它。请耐心等待 TensorRT 导出完成,这可能需要几分钟。

在创建 TensorRT 模型文件(.engine)后,您将看到如下输出:

  • 步骤 2. 如果您想传递其他参数,可以参考下表:
描述
imgsz640图像大小,标量或 (h, w) 列表,例如 (640, 480)
halfFalseFP16 量化
dynamicFalse动态轴
simplifyFalse简化模型
workspace4工作区大小(GB)

例如,如果您希望将 PyTorch 模型转换为 FP16 量化的 TensorRT 模型,可以执行以下命令:

yolo export model=yolov8n.pt format=engine half=True device=0

一旦模型成功导出,您可以直接在 yolopredict 命令中使用 model= 参数替换此模型,用于运行检测、分类、分割、姿态估计的所有四种任务。

例如,对于对象检测:

yolo detect predict model=yolov8n.engine source='0' show=True

使用您自己的 AI 模型

数据收集与标注

如果您有特定的 AI 应用,并希望使用适合您应用的自定义 AI 模型,您可以收集自己的数据集,对其进行标注,然后使用 YOLOv8 进行训练。

如果您不想自己收集数据,也可以选择现成的公共数据集。您可以下载许多公开可用的数据集,例如 COCO 数据集Pascal VOC 数据集 等。Roboflow Universe 是一个推荐的平台,它提供了广泛的数据集,并且拥有 90,000+ 数据集和 66+ 百万张图像,可用于构建计算机视觉模型。此外,您还可以在 Google 上搜索开源数据集,并从各种可用的数据集中进行选择。

如果您有自己的数据集并希望对图像进行标注,我们推荐您使用 Roboflow 提供的标注工具。请参考 Wiki 的这一部分 了解更多信息。您也可以参考 Roboflow 的 这篇指南 了解标注相关内容。

训练

这里我们提供了三种训练模型的方法。

  1. 第一种方法是使用 Ultralytics HUB。您可以轻松将 Roboflow 集成到 Ultralytics HUB 中,这样您的所有 Roboflow 项目都可以直接用于训练。它提供了一个 Google Colab 笔记本,可以轻松启动训练过程,并实时查看训练进度。

  2. 第二种方法是使用我们创建的 Google Colab 工作空间,以简化训练过程。在这里,我们使用 Roboflow API 从 Roboflow 项目中下载数据集。

  3. 第三种方法是使用本地 PC 进行训练。在这种情况下,您需要确保拥有足够强大的 GPU,并且需要手动下载数据集。

这里我们使用 Ultralytics HUB 加载 Roboflow 项目,然后在 Google Colab 上进行训练。

  • 步骤 1. 访问 此 URL 并注册一个 Ultralytics 账户

  • 步骤 2. 使用新创建的账户登录后,您将看到以下仪表板

  • 步骤 3. 访问 此 URL 并注册一个 Roboflow 账户

  • 步骤 4. 使用新创建的账户登录后,您将看到以下仪表板

  • 步骤 5. 创建一个新的工作空间,并按照 我们准备的 Wiki 指南 在工作空间下创建一个新项目。您也可以 点击这里 从 Roboflow 官方文档中了解更多信息。

  • 步骤 6. 一旦您的工作空间中有几个项目,它将如下所示

  • 步骤 7. 转到 设置 并点击 Roboflow API
  • 步骤 8. 点击 复制 按钮以复制 私有 API 密钥
  • 步骤 9. 返回到 Ultralytics HUB 仪表板,点击 集成,将之前复制的 API 密钥粘贴到空白列中,然后点击 添加
  • 步骤 10. 如果您看到列出了您的工作空间名称,则说明集成成功
  • 步骤 11. 导航到 数据集,您将在这里看到所有的 Roboflow 项目
  • 步骤 12. 点击一个项目以了解更多关于数据集的信息。这里我选择了一个可以检测健康和受损苹果的数据集
  • 步骤 13. 点击 训练模型
  • 步骤 14. 选择 架构,设置一个 模型名称(可选),然后点击 继续。这里我们选择 YOLOv8s 作为模型架构
  • 步骤 15.高级选项 下,根据您的偏好配置设置,复制并粘贴 Colab 代码(稍后将粘贴到 Colab 工作空间中),然后点击 打开 Google Colab
  • 步骤 16. 如果您尚未登录,请登录您的 Google 账户
  • 步骤 17 导航到 Runtime > Change runtime type
  • 步骤 18Hardware accelerator 下选择 GPU,在 GPU type 下选择最高可用选项,然后点击 Save
  • 步骤 19 点击 Connect
  • 步骤 20 点击 RAM, Disk 按钮以检查硬件资源使用情况
  • 步骤 21 点击 Play 按钮运行第一个代码单元格
  • 步骤 22 将之前从 Ultralytics HUB 复制的代码单元粘贴到 Start 部分下并运行它以开始训练
  • 步骤 23 现在返回到 Ultralytics HUB,你会看到消息 Connected。点击 Done
  • 步骤 24 在模型使用 Google Colab 进行训练时,你将实时看到 Box Loss, Class Loss 和 Object Loss
  • 步骤 25 训练完成后,你将在 Google Colab 上看到以下输出
  • 步骤 26 现在返回到 Ultralytics HUB,进入 Preview 标签页并上传测试图像以检查训练模型的性能
  • 步骤 27 最后进入 Deploy 标签页并以你喜欢的格式下载训练好的模型以便使用 YOLOv8 进行推理。这里我们选择了 PyTorch。

现在你可以使用这个下载的模型完成我们之前在本 Wiki 中解释的任务。你只需要用你的模型替换模型文件即可。

例如:

yolo detect predict model=<your_model.pt> source='0' show=True

性能基准测试

准备工作

我们对运行在 reComputer J4012/ reComputer Industrial J4012(由 NVIDIA Jetson Orin NX 16GB 模块驱动)上的 YOLOv8 所支持的所有计算机视觉任务进行了性能基准测试。

在示例目录中包含一个名为 trtexec 的命令行包装工具。trtexec 是一个无需开发自己的应用程序即可使用 TensorRT 的工具。trtexec 工具有三个主要用途:

  • 基于随机或用户提供的输入数据对网络进行基准测试。
  • 从模型生成序列化引擎。
  • 从构建器生成序列化的时间缓存。

在这里,我们可以使用 trtexec 工具快速对具有不同参数的模型进行基准测试。但首先,你需要一个 ONNX 模型,我们可以通过使用 ultralytics yolov8 来生成这个 ONNX 模型。

  • 步骤 1. 使用以下命令生成 ONNX 模型:
yolo mode=export model=yolov8s.pt format=onnx
  • 步骤 2. 使用 trtexec 生成引擎文件,命令如下:
cd /usr/src/tensorrt/bin
./trtexec --onnx=<path_to_onnx_file> --saveEngine=<path_to_save_engine_file>

例如:

./trtexec --onnx=/home/nvidia/yolov8s.onnx --saveEngine=/home/nvidia/yolov8s.engine

这将输出性能结果,同时生成一个 .engine 文件。默认情况下,它会将 ONNX 转换为 FP32 精度的 TensorRT 优化文件,你可以看到如下输出:

如果你需要 FP16 精度(相比 FP32 提供更好的性能),可以执行以下命令:

./trtexec --onnx=/home/nvidia/yolov8s.onnx --fp16 --saveEngine=/home/nvidia/yolov8s.engine 

然而,如果你需要 INT8 精度(相比 FP16 提供更好的性能),可以执行以下命令:

./trtexec --onnx=/home/nvidia/yolov8s.onnx --int8 --saveEngine=/home/nvidia/yolov8s.engine 

结果

以下是我们在 reComputer J4012/ reComputer Industrial J4012 上运行的四个计算机视觉任务的结果总结。

额外演示:使用 YOLOv8 的运动检测与计数器

我们基于 YOLOv8-Pose 模型构建了一个用于运动检测和计数的姿态估计演示应用程序。你可以在 这里 查看该项目,了解更多关于此演示的信息,并将其部署到你自己的 Jetson 设备上!

手动设置 YOLOv8 于 NVIDIA Jetson

如果之前提到的一键脚本出现错误,您可以按照以下步骤逐步为 Jetson 设备准备 YOLOv8。

安装 Ultralytics 包

  • 步骤 1. 访问 Jetson 设备的终端,安装 pip 并升级
sudo apt update
sudo apt install -y python3-pip -y
pip3 install --upgrade pip
  • 步骤 2. 安装 Ultralytics 包
pip3 install ultralytics
  • 步骤 3. 升级 numpy 到最新版本
pip3 install numpy -U
  • 步骤 4. 重启设备
sudo reboot

卸载 Torch 和 Torchvision

上述 Ultralytics 的安装会自动安装 Torch 和 Torchvision。然而,通过 pip 安装的这两个包并不兼容运行在基于 ARM aarch64 架构 的 Jetson 平台。因此,我们需要手动安装预编译的 PyTorch pip wheel 并从源码编译/安装 Torchvision。

pip3 uninstall torch torchvision

安装 PyTorch 和 Torchvision

访问 此页面 获取所有 PyTorch 和 Torchvision 的链接。

以下是 JetPack 5.0 及以上版本支持的一些版本。

PyTorch v2.0.0

支持 JetPack 5.1 (L4T R35.2.1) / JetPack 5.1.1 (L4T R35.3.1),Python 3.8

文件名: torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl
URL: https://nvidia.box.com/shared/static/i8pukc49h3lhak4kkn67tg9j4goqm0m7.whl

PyTorch v1.13.0

支持 JetPack 5.0 (L4T R34.1) / JetPack 5.0.2 (L4T R35.1) / JetPack 5.1 (L4T R35.2.1) / JetPack 5.1.1 (L4T R35.3.1),Python 3.8

文件名: torch-1.13.0a0+d0d6b1f2.nv22.10-cp38-cp38-linux_aarch64.whl
URL: https://developer.download.nvidia.com/compute/redist/jp/v502/pytorch/torch-1.13.0a0+d0d6b1f2.nv22.10-cp38-cp38-linux_aarch64.whl

  • 步骤 1. 根据您的 JetPack 版本安装 torch,格式如下:
wget <URL> -O <file_name>
pip3 install <file_name>

例如,这里我们运行的是 JP5.1.1,因此选择 PyTorch v2.0.0

sudo apt-get install -y libopenblas-base libopenmpi-dev
wget https://nvidia.box.com/shared/static/i8pukc49h3lhak4kkn67tg9j4goqm0m7.whl -O torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl
pip3 install torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl
  • 步骤 2. 根据您安装的 PyTorch 版本安装对应的 torchvision。例如,我们选择了 PyTorch v2.0.0,因此需要选择 Torchvision v0.15.2
sudo apt install -y libjpeg-dev zlib1g-dev
git clone https://github.com/pytorch/vision torchvision
cd torchvision
git checkout v0.15.2
python3 setup.py install --user

以下是根据 PyTorch 版本需要安装的对应 torchvision 版本列表:

  • PyTorch v2.0.0 - torchvision v0.15
  • PyTorch v1.13.0 - torchvision v0.14

如果需要更详细的列表,请查看 此链接

安装 ONNX 并降级 Numpy

如果您需要将 PyTorch 模型转换为 TensorRT,则需要执行以下步骤:

  • 步骤 1. 安装 ONNX(这是一个必要的依赖)
pip3 install onnx
  • 步骤 2. 降级到较低版本的 Numpy 以修复错误
pip3 install numpy==1.20.3

资源

技术支持与产品讨论

感谢您选择我们的产品!我们致力于为您提供多种支持,以确保您使用我们的产品时获得尽可能顺畅的体验。我们提供多种沟通渠道,以满足不同的偏好和需求。

Loading Comments...