博客

在 Jetson Nano 上运行 PyTorch 模型

作者 2022年3月16日2024年11月15日暂无评论

概述

NVIDIA Jetson Nano 是 Jetson 系列产品或 Jetson 模组的一员，是一款小巧而强大的基于 Linux (Ubuntu) 的嵌入式计算机，配备 2GB/4GB GPU。使用它，您可以高效地运行许多 PyTorch 模型。本文档总结了我们在 Jetson Nano 上使用 3 种不同机制运行各种深度学习模型的经验。

  1. Jetson Inference 是更高级的 NVIDIA API，内置了对运行大多数常见计算机视觉模型的支持，这些模型可以在 Jetson 平台上通过 PyTorch 进行迁移学习。
  2. TensorRT 是一个来自 NVIDIA 的高性能推理 SDK，需要将 PyTorch 模型转换为 ONNX，然后再转换为 TensorRT 运行时可以运行的 TensorRT 引擎文件。
  3. PyTorch 使用直接的 PyTorch API torch.nn 进行推理。

设置 Jetson Nano

在此处购买 Jetson Nano 后，只需按照清晰的分步说明将 Jetson Nano 开发者套件的 SD 卡镜像下载并写入 microSD 卡，即可完成设置。设置完成后启动 Nano，您将看到标准的 Linux 提示符以及设置时使用的用户名和 Nano 名称。

要检查 Nano 上的 GPU 状态，请运行以下命令：

sudo pip3 install jetson-stats
sudo jtop

您将看到相关信息，包括：

您还可以查看已安装的 CUDA 版本：

$ ls -lt /usr/local
lrwxrwxrwx  1 root root   22 Aug  2 01:47 cuda -> /etc/alternatives/cuda
lrwxrwxrwx  1 root root   25 Aug  2 01:47 cuda-10 -> /etc/alternatives/cuda-10
drwxr-xr-x 12 root root 4096 Aug  2 01:47 cuda-10.2

若要在 Jetson Nano 上使用摄像头（例如 Arducam 8MP IMX219），请按照此处的说明操作，或在安装摄像头模块后运行以下命令：

cd ~
wget https://github.com/ArduCAM/MIPI_Camera/releases/download/v0.0.3/install_full.sh
chmod +x install_full.sh
./install_full.sh -m arducam

另一种方法是使用原生的 Jetson Nano 摄像头驱动程序：

sudo dpkg -r arducam-nvidia-l4t-kernel
sudo shutdown -r now

然后，使用 ls /dev/video0 来确认摄像头已被识别：

$ ls /dev/video0
/dev/video0

最后，使用以下命令查看摄像头实时画面：

nvgstcapture-1.0 --orientation=2

使用 Jetson Inference

NVIDIA Jetson Inference API 为在 Jetson Nano 上运行图像识别、目标检测、语义分割和姿态估计模型提供了最简单的方法。Jetson Inference 内置了 TensorRT，因此速度非常快。

要测试运行 Jetson Inference，首先克隆仓库并下载模型：

git clone --recursive https://github.com/dusty-nv/jetson-inference
cd jetson-inference

然后使用已经安装了 PyTorch 的预构建 Docker 容器来测试运行模型：

docker/run.sh --volume ~/jetson_inference:/jetson_inference

要在测试图像上运行图像识别、目标检测、语义分割和姿态估计模型，请使用以下命令：

cd build/aarch64/bin
./imagenet.py images/jellyfish.jpg /jetson_inference/jellyfish.jpg
./segnet.py images/dog.jpg /jetson_inference/dog.jpeg
./detectnet.py images/peds_0.jpg /jetson_inference/peds_0.jpg
./posenet.py images/humans_0.jpg /jetson_inference/pose_humans_0.jpg

将生成运行四个不同模型后的四个结果图像。退出 Docker 镜像即可查看它们：

$ ls -lt ~/jetson_inference/
-rw-r--r-- 1 root root  68834 Oct 15 21:30 pose_humans_0.jpg
-rw-r--r-- 1 root root 914058 Oct 15 21:30 peds_0.jpg
-rw-r--r-- 1 root root 666239 Oct 15 21:30 dog.jpeg
-rw-r--r-- 1 root root 179760 Oct 15 21:29 jellyfish.jpg
Using jest interface example 1
Using jest interface example 2
Using jest interface example 3
Using jest interface example 4

您也可以使用该 Docker 镜像运行 PyTorch 模型，因为该镜像已安装了 PyTorch、torchvision 和 torchaudio。

# pip list|grep torch
torch (1.9.0)
torchaudio (0.9.0a0+33b2469)
torchvision (0.10.0a0+300a8a4)

尽管 Jetson Inference 已包含转换为 TensorRT 引擎格式的模型，但您可以通过按照此处“使用 PyTorch 进行迁移学习（针对 Jetson Inference）”的步骤进行微调。

使用 TensorRT

TensorRT 是 NVIDIA 推出的一个高性能推理 SDK。Jetson Nano 通过 Jetpack SDK 支持 TensorRT，该 SDK 包含在用于设置 Jetson Nano 的 SD 卡镜像中。要确认 TensorRT 已安装在 Nano 上，请运行 dpkg -l|grep -i tensorrt：

理论上，TensorRT 可用于“获取训练好的 PyTorch 模型并对其进行优化，以便在 NVIDIA GPU 上进行更高效的推理”。请按照笔记本 (notebook) 中的说明和代码，了解如何在 torchvision Resnet50 模型上通过 ONNX 将 PyTorch 与 TensorRT 结合使用。

  1. 如何将模型从 PyTorch 转换为 ONNX；
  2. 如何将 ONNX 模型转换为 TensorRT 引擎文件；
  3. 如何使用 TensorRT 运行时运行引擎文件以提升性能：推理时间从原始的 31.5ms/19.4ms（FP32/FP16 精度）提高到 6.28ms（TensorRT）。

您可以将笔记本代码中的 Resnet50 模型替换为其他 PyTorch 模型，按照上述转换过程操作，并使用 TensorRT 运行时运行最终转换后的 TensorRT 引擎文件，以查看优化后的性能。但请注意，由于 Nano GPU 内存限制，超过 100MB 的模型很可能会运行失败，并出现以下错误信息：

错误代码 1：Cuda Runtime（所有支持 CUDA 的设备均忙碌或不可用）

在将 PyTorch 模型转换为 ONNX 模型时，您也可能会遇到错误，这可以通过将以下代码：

torch.onnx.export(resnet50, dummy_input, "resnet50_pytorch.onnx", verbose=False)

替换为：

torch.onnx.export(model, dummy_input, "deeplabv3_pytorch.onnx", opset_version=11, verbose=False)

使用 PyTorch

首先，要在 Nano 上下载并安装 PyTorch 1.9，请运行以下命令（更多信息请参阅此处）：

wget https://nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl -O torch-1.8.0-cp36-cp36m-linux_aarch64.whl -O torch-1.9.0-cp36-cp36m-linux_aarch64.whl
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev 
pip3 install Cython
pip3 install numpy torch-1.9.0-cp36-cp36m-linux_aarch64.whl

要在 Nano 上下载并安装 torchvision 0.10，请运行以下命令：

https://drive.google.com/uc?id=1tU6YlPjrP605j4z8PMnqwCSoP6sSC91Z
pip3 install torchvision-0.10.0a0+300a8a4-cp36-cp36m-linux_aarch64.whl

完成上述步骤后，运行此命令进行确认：

$ pip3 list|grep torch
torch (1.9.0)
torchvision (0.10.0)

您也可以使用“使用 Jetson Inference”章节中描述的 Docker 镜像（该镜像也安装了 PyTorch 和 torchvision），从而跳过上述手动步骤。

官方 YOLOv5 仓库被用于在 Jetson Nano 上运行 PyTorch YOLOv5 模型。登录 Jetson Nano 后，按照以下步骤操作：

  • 获取仓库并安装所需依赖：
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt
  • 运行 python3 detect.py，默认使用 PyTorch yolov5s.pt 模型。您应该会看到类似以下内容：
detect: weights=yolov5s.pt, source=data/images, imgsz=[640, 640], conf_thres=0.25, iou_thres=0.45, max_det=1000, device=, view_img=False, save_txt=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=runs/detect, name=exp, exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, half=False
YOLOv5 🚀 v5.0-499-g48b00db torch 1.9.0 CUDA:0 (NVIDIA Tegra X1, 3956.1015625MB)

Fusing layers... 
Model Summary: 224 layers, 7266973 parameters, 0 gradients
image 1/5 /home/jeff/repos/yolov5-new/yolov5/data/images/bus.jpg: 640x480 4 persons, 1 bus, 1 fire hydrant, Done. (0.142s)
...

在 Jetson Nano GPU 上的推理时间约为 140ms，比在 iOS 或 Android（约 330ms）上的推理时间快两倍以上。

如果遇到错误 “ImportError: The _imagingft C module is not installed.”，则需要重新安装 pillow：

sudo apt-get install libpng-dev
sudo apt-get install libfreetype6-dev
pip3 uninstall pillow
pip3 install --no-cache-dir pillow

成功运行 python3 detect.py 后，位于 data/images 中的测试图像的目标检测结果将保存在 runs/detect/exp 目录中。要使用实时摄像头而非本地图像进行检测测试，请在运行 python3 detect.py 时使用 --source 0 参数。

~/repos/yolov5$ ls -lt runs/detect/exp10
total 1456
-rw-rw-r-- 1 jeff jeff 254895 Oct 15 16:12 zidane.jpg
-rw-rw-r-- 1 jeff jeff 202674 Oct 15 16:12 test3.png
-rw-rw-r-- 1 jeff jeff 217117 Oct 15 16:12 test2.jpg
-rw-rw-r-- 1 jeff jeff 305826 Oct 15 16:12 test1.png
-rw-rw-r-- 1 jeff jeff 495760 Oct 15 16:12 bus.jpg

使用 PyTorch iOS YOLOv5 演示应用或 Android YOLOv5 演示应用中使用的相同测试文件，您可以比较在移动设备上运行 YOLOv5 PyTorch 模型与在 Jetson Nano 上运行的结果。

PyTorch YOLOv5 on Jetson Nano, example with a dog
PyTorch YOLOv5 on Jetson Nano, example with a horse and a rider

图 1. Jetson Nano 上的 PyTorch YOLOv5。

PyTorch YOLOv5 on iOS, example with a dog
PyTorch YOLOv5 on iOS, example with a horse and a rider

图 2. iOS 上的 PyTorch YOLOv5。

PyTorch YOLOv5 on Android, example with a dog
PyTorch YOLOv5 on Android, example with a horse and a rider

图 3. Android 上的 PyTorch YOLOv5。

总结

根据我们在 Jetson Nano 上运行各种 PyTorch 模型以用于潜在演示应用的经验，我们看到即使是 Jetson 系列产品中入门级的 Jetson Nano，也提供了强大的 GPU 和嵌入式系统，能够直接高效地运行一些最新的预训练或迁移学习过的 PyTorch 模型。

在 Jetson Nano 上构建 PyTorch 演示应用与在 Linux 上构建类似，但您也可以选择在将 PyTorch 模型转换为 TensorRT 引擎文件格式后使用 TensorRT。

但是，如果您只需要在 Jetson Nano 上使用支持图像识别、目标检测、语义分割和姿态估计模型的 NVIDIA Jetson Inference 运行一些常见的计算机视觉模型，那么这是最简单的方法。

参考文献

Torch-TensorRT，一个通过 TensorRT 实现的 PyTorch 编译器：https://github.com/NVIDIA/Torch-TensorRT/

Jetson Inference Docker 镜像详情：https://github.com/dusty-nv/jetson-inference/blob/master/docs/aux-docker.md

NVIDIA Jetson Nano 上使用 TensorRT 的指南：https://docs.donkeycar.com/guide/robot_sbc/tensorrt_jetson_nano/（包括）

  1. 使用 Jetson 作为便携式 GPU 设备来运行 NN 象棋引擎模型：https://medium.com/@ezchess/jetson-lc0-running-leela-chess-zero-on-nvidia-jetson-a-portable-gpu-device-a213afc9c018
  2. 一个使用 PyTorch 和 torchvision 并直接通过 pip 安装的 MaskEraser 应用：https://github.com/INTEC-ATI/MaskEraser#install-pytorch