在 Jetson Nano 上运行 PyTorch 模型

作者：Jeff Tang, Hamid Shojanazeri, Geeta Chauhan

概述

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

Jetson Inference 是更高级别的 NVIDIA API，内置支持运行最常见的计算机视觉模型，这些模型可以在 Jetson 平台上使用 PyTorch 进行迁移学习。
TensorRT 是 NVIDIA 的高性能推理 SDK，需要将 PyTorch 模型转换为 ONNX，然后再转换为 TensorRT 引擎文件，以便 TensorRT 运行时可以运行。
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，首先克隆 repo 并下载模型

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

您也可以使用 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，请运行 dpkg -l|grep -i tensorrt

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

如何将模型从 PyTorch 转换为 ONNX；
如何将 ONNX 模型转换为 TensorRT 引擎文件；
如何使用 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 repo 用于在 Jetson Nano 上运行 PyTorch YOLOv5 模型。登录 Jetson Nano 后，请按照以下步骤操作

获取 repo 并安装所需内容

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 上的推理时间约为 140 毫秒，比 iOS 或 Android 上的推理时间（约 330 毫秒）快两倍以上。

如果您收到错误“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 Nano 这种 Jetson 产品系列的低端产品，也提供了强大的 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

总结

参考

文档

教程

资源