由谭杰夫、哈米德·肖贾纳泽里、吉塔·乔汗撰写

概述

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

  1. Jetson Inference 是 NVIDIA 提供的高级 API,它内置了对在 Jetson 平台上运行大多数常见计算机视觉模型的支持,这些模型可以使用 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

然后使用已预装 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 引擎文件格式的模型,但您可以通过遵循 Transfer Learning with PyTorch (for Jetson Inference)中的步骤来微调这些模型。

使用 TensorRT

TensorRT 是 NVIDIA 用于高性能推理的 SDK。Jetson Nano 通过包含在用于设置 Jetson Nano 的 SD 卡镜像中的 Jetpack SDK 支持 TensorRT。要确认 TensorRT 已安装在 Nano 中,请使用 run dpkg -l|grep -i tensorrt

理论上,TensorRT 可以用来“将训练好的 PyTorch 模型优化,使其在 NVIDIA GPU 上推理时运行更高效。”请按照笔记本中的说明和代码,了解如何通过 ONNX 在 torchvision Resnet50 模型中使用 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 的模型可能无法运行,可能会出现以下错误信息:

Error Code 1: Cuda Runtime (all CUDA-capable devices are busy or unavailable)

在将 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)

您还可以使用本节中描述的 Using 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 目录中。要使用实时摄像头而不是本地图像进行检测测试,请在运行 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. 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。

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

参考文献列表

Torch-TensorRT,PyTorch 的 TensorRT 编译器: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