跳转到主要内容
博客

在 Jetson Nano 上运行 PyTorch 模型

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

概述

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

  1. Jetson Inference 是更高级别的 NVIDIA API，内置支持运行大多数常见的计算机视觉模型，这些模型可以通过 PyTorch 在 Jetson 平台上进行迁移学习。
  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

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

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)。

您可以将 notebook 代码中的 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 演示应用程序中相同的测试文件，您可以比较在移动设备和 Jetson Nano 上运行 YOLOv5 PyTorch 模型生成的结果

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 应用程序，但您也可以选择在将 PyTorch 模型转换为 TensorRT 引擎文件格式后使用 TensorRT。

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

参考资料

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 的 MaskEraser 应用程序，直接通过 pip 安装：https://github.com/INTEC-ATI/MaskEraser#install-pytorch