Raspberry Pi 4 上的实时推断(30 帧/秒!)¶
Created On: Feb 08, 2022 | Last Updated: Jan 16, 2024 | Last Verified: Nov 05, 2024
作者:Tristan Rice
PyTorch 原生支持 Raspberry Pi 4。本教程将指导您如何设置 Raspberry Pi 4 以运行 PyTorch,并在 CPU 上实时运行 MobileNet v2 分类模型(30 帧/秒以上)。
所有测试均使用 Raspberry Pi 4 Model B 4GB,但也可以在 2GB 版本及性能较低的 3B 上运行。

前置条件¶
要跟随本教程,您需要 Raspberry Pi 4、一台相机以及所有其他标准配件。
散热片和风扇(可选但推荐)
5V 3A USB-C 电源适配器
SD 卡(至少 8GB)
SD 卡读卡器
Raspberry Pi 4 设置¶
PyTorch 仅提供适用于 Arm 64 位 (aarch64) 的 pip 包,因此您需要在 Raspberry Pi 上安装 64 位版本的操作系统。
您可以从 https://downloads.raspberrypi.org/raspios_arm64/images/ 下载最新的 arm64 Raspberry Pi OS,并通过 rpi-imager 安装它。
32 位的 Raspberry Pi OS 无法工作。

安装至少需要几分钟,具体取决于您的网络速度和 SD 卡速度。安装完成后应该看起来像这样:

现在是时候将您的 SD 卡插入 Raspberry Pi,连接相机并启动了。

启动后并完成初始设置后,您需要编辑 /boot/config.txt
文件以启用摄像头。
# This enables the extended features such as the camera.
start_x=1
# This needs to be at least 128M for the camera processing, if it's bigger you can just leave it as is.
gpu_mem=128
# You need to commment/remove the existing camera_auto_detect line since this causes issues with OpenCV/V4L2 capture.
#camera_auto_detect=1
然后重启。重启后,video4linux2 设备 /dev/video0
应该会存在。
安装 PyTorch 和 OpenCV¶
PyTorch 及我们需要的所有其他库都有 ARM 64 位/aarch64 变体,因此您可以通过 pip 安装并像在其他 Linux 系统上一样工作。
$ pip install torch torchvision torchaudio
$ pip install opencv-python
$ pip install numpy --upgrade

我们现在可以检查是否所有内容正确安装:
$ python -c "import torch; print(torch.__version__)"

视频捕获¶
对于视频捕获,我们将使用 OpenCV 流式传输视频帧,而不是更常见的 picamera
。64 位 Raspberry Pi OS 上不支持 picamera,且其速度比 OpenCV 慢得多。OpenCV 直接访问 /dev/video0
设备以抓取帧。
我们使用的模型(MobileNetV2)接收 224x224
的图像尺寸,因此我们可以直接从 OpenCV 请求此大小的帧,帧率为 36fps。我们目标是让模型在 30 帧每秒的速率运行,但请求稍高的帧率以确保始终有足够的帧。
import cv2
from PIL import Image
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 224)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 224)
cap.set(cv2.CAP_PROP_FPS, 36)
OpenCV 返回一个 BGR 格式的 numpy
数组,因此我们需要读取并做些调整以将其转换为所需的 RGB 格式。
ret, image = cap.read()
# convert opencv output from BGR to RGB
image = image[:, :, [2, 1, 0]]
此数据读取和处理大约需要 3.5 毫秒
。
图像预处理¶
我们需要将帧转换为模型所需的格式。这与任何机器上标准 torchvision 转换的处理方式相同。
from torchvision import transforms
preprocess = transforms.Compose([
# convert the frame to a CHW torch tensor for training
transforms.ToTensor(),
# normalize the colors to the range that mobilenet_v2/3 expect
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(image)
# The model can handle multiple images simultaneously so we need to add an
# empty dimension for the batch.
# [3, 224, 224] -> [1, 3, 224, 224]
input_batch = input_tensor.unsqueeze(0)
模型选择¶
您可以选择许多模型以满足不同的性能特性。并非所有模型都提供 qnnpack
预训练变体,因此出于测试目的,建议选择一个具有此功能的模型,但如果您训练并量化了自己的模型,可以使用其中任何一个。
本教程中我们使用的是 mobilenet_v2
,因为它具有良好的性能和准确性。
Raspberry Pi 4 基准结果:
模型 |
帧率(FPS) |
每帧总时间(毫秒) |
每帧模型时间(毫秒) |
qnnpack 预训练 |
---|---|---|---|---|
mobilenet_v2 |
33.7 |
29.7 |
26.4 |
True |
mobilenet_v3_large |
29.3 |
34.1 |
30.7 |
True |
resnet18 |
9.2 |
109.0 |
100.3 |
False |
resnet50 |
4.3 |
233.9 |
225.2 |
False |
resnext101_32x8d |
1.1 |
892.5 |
885.3 |
False |
inception_v3 |
4.9 |
204.1 |
195.5 |
False |
googlenet |
7.4 |
135.3 |
132.0 |
False |
shufflenet_v2_x0_5 |
46.7 |
21.4 |
18.2 |
False |
shufflenet_v2_x1_0 |
24.4 |
41.0 |
37.7 |
False |
shufflenet_v2_x1_5 |
16.8 |
59.6 |
56.3 |
False |
shufflenet_v2_x2_0 |
11.6 |
86.3 |
82.7 |
False |
MobileNetV2:量化和 JIT¶
为了获得最佳性能,我们希望使用一个量化和融合的模型。量化意味着它使用 int8 进行计算,这比标准 float32 计算更高效。融合意味着将连续的操作融合为更高效的版本(如果可能)。例如,推理期间,通常会将激活函数(ReLU
)与前一层(Conv2d
)合并。
aarch64 版本的 pytorch 需要使用 qnnpack
引擎。
import torch
torch.backends.quantized.engine = 'qnnpack'
对于本示例,我们将使用由 torchvision 提供的、量化和融合的 MobileNetV2 预训练版本。
from torchvision import models
net = models.quantization.mobilenet_v2(pretrained=True, quantize=True)
然后我们希望通过 jit 减少 Python 开销并融合任何操作。Jit 使我们达到 ~30fps,而未使用时则仅为 ~20fps。
net = torch.jit.script(net)
将其整合在一起¶
现在我们可以将所有部分整合在一起并运行:
import time
import torch
import numpy as np
from torchvision import models, transforms
import cv2
from PIL import Image
torch.backends.quantized.engine = 'qnnpack'
cap = cv2.VideoCapture(0, cv2.CAP_V4L2)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 224)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 224)
cap.set(cv2.CAP_PROP_FPS, 36)
preprocess = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
net = models.quantization.mobilenet_v2(pretrained=True, quantize=True)
# jit model to take it from ~20fps to ~30fps
net = torch.jit.script(net)
started = time.time()
last_logged = time.time()
frame_count = 0
with torch.no_grad():
while True:
# read frame
ret, image = cap.read()
if not ret:
raise RuntimeError("failed to read frame")
# convert opencv output from BGR to RGB
image = image[:, :, [2, 1, 0]]
permuted = image
# preprocess
input_tensor = preprocess(image)
# create a mini-batch as expected by the model
input_batch = input_tensor.unsqueeze(0)
# run model
output = net(input_batch)
# do something with output ...
# log model performance
frame_count += 1
now = time.time()
if now - last_logged > 1:
print(f"{frame_count / (now-last_logged)} fps")
last_logged = now
frame_count = 0
运行结果显示我们维持在 ~30fps 左右。

这是在树莓派操作系统中使用所有默认设置的表现。如果您禁用了用户界面和所有其他默认启用的后台服务,性能会更好且更稳定。
如果我们检查 htop
,会发现几乎有100%的利用率。

为了验证其端到端工作,我们可以计算类别的概率,并使用`ImageNet类别标签 <https://gist.github.com/yrevar/942d3a0ac09ec9e5eb3a>`_来打印检测结果。
top = list(enumerate(output[0].softmax(dim=0)))
top.sort(key=lambda x: x[1], reverse=True)
for idx, val in top[:10]:
print(f"{val.item()*100:.2f}% {classes[idx]}")
实时运行的``mobilenet_v3_large``:

检测橙子:

检测杯子:

故障排除:性能问题¶
PyTorch会默认使用所有可用的核心。如果在树莓派上有任何后台运行的程序,可能会与模型推断冲突,从而导致延迟波动。为了解决此问题,您可以减少线程数,这将以小幅性能损失换取降低最高延迟。
torch.set_num_threads(2)
对于 shufflenet_v2_x1_5
,使用``2个线程``而不是``4个线程``,虽然将最佳延迟从``60毫秒``增加到``72毫秒``,但消除了``128毫秒``的延迟波动。
下一步¶
您可以创建自己的模型或微调现有模型。如果您使用`torchvision.models.quantized <https://pytorch.org/vision/stable/models.html#quantized-models>`_中的模型进行微调,大部分用于融合和量化的工作已为您完成,因此您可以直接部署到树莓派上并获得良好的性能。
查看更多: