在 Jetson Nano 上运行 PyTorch 模型

作者：Jeff Tang, Hamid Shojanazeri, Geeta Chauhan

概述

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

Jetson Inference 是更高级的 NVIDIA API，内置支持在 Jetson 平台上运行大多数常见的计算机视觉模型，这些模型可以使用 PyTorch 进行迁移学习。
TensorRT 是 NVIDIA 的一个高性能推理 SDK，它需要将 PyTorch 模型转换为 ONNX 格式，然后再转换为 TensorRT 运行时可以运行的 TensorRT 引擎文件。
使用直接 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

您也可以使用上面提到的 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，Jetpack SDK 包含在用于设置 Jetson Nano 的 SD 卡镜像中。要确认 Nano 中已安装 TensorRT，请运行 run dpkg -l|grep -i tensorrt

理论上，TensorRT 可以用于“将训练好的 PyTorch 模型优化，使其在 NVIDIA GPU 上进行推理时运行得更有效率。”按照笔记本中的说明和代码，了解如何在 torchvision Resnet50 模型上通过 ONNX 将 PyTorch 与 TensorRT 结合使用

如何将模型从 PyTorch 转换为 ONNX；
如何将 ONNX 模型转换为 TensorRT 引擎文件；
如何使用 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 demo app 或 Android YOLOv5 demo app 中使用的相同测试文件，您可以比较在移动设备和 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. PyTorch YOLOv5 在 Jetson Nano 上运行。

PyTorch YOLOv5 on iOS, example with a dog

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

图 2. PyTorch YOLOv5 在 iOS 上运行。

PyTorch YOLOv5 on Android, example with a dog

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

图 3. PyTorch YOLOv5 在 Android 上运行。

总结

根据我们在 Jetson Nano 上运行不同 PyTorch 模型以构建潜在演示应用的经验，我们看到，即使是 Jetson 产品系列中的低端产品 Jetson Nano，也提供了强大的 GPU 和嵌入式系统，可以直接高效地运行一些最新的 PyTorch 模型，无论是预训练的还是迁移学习的。

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

但是，如果您只需要使用 NVIDIA 的 Jetson Inference 在 Jetson Nano 上运行一些常见的计算机视觉模型（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/ 包括

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

概述

设置 Jetson Nano

使用 Jetson Inference

使用 TensorRT

使用 PyTorch

总结

参考资料

文档

教程

资源