概述
NVIDIA Jetson Nano 是 Jetson 产品系列或 Jetson 模块中的一款小型而强大的基于 Linux(Ubuntu)的嵌入式计算机,配备 2/4GB GPU。使用它,您可以高效地运行许多 PyTorch 模型。本文档总结了我们在 Jetson Nano 上使用 3 种不同机制运行不同深度学习模型的经验:
-
Jetson Inference 是 NVIDIA 提供的高级 API,它内置了对在 Jetson 平台上运行大多数常见计算机视觉模型的支持,这些模型可以使用 PyTorch 在 Jetson 平台上迁移学习。
-
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,首先克隆仓库并下载模型:
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 引擎文件格式的模型,但您可以通过遵循 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:
-
如何将模型从 PyTorch 转换为 ONNX;
-
如何将 ONNX 模型转换为 TensorRT 引擎文件;
-
如何使用 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 模型生成的结果:


图 1. PyTorch YOLOv5 在 Jetson Nano 上。


图 2. PyTorch YOLOv5 在 iOS 上。


图 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/ 包括:
-
使用 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