在 Raspberry Pi 4 和 5 上进行实时推理(40 fps!)#
创建日期:2022 年 2 月 8 日 | 最后更新:2025 年 9 月 30 日 | 最后验证:2024 年 11 月 5 日
作者: Tristan Rice
PyTorch 开箱即用地支持 Raspberry Pi 4 和 5。本教程将指导您如何配置 Raspberry Pi 以运行 PyTorch,并在 CPU 上实时(30-40 fps)运行 MobileNet v2 分类模型。
以上内容已在 Raspberry Pi 4 Model B 4GB 上进行了测试,但同样适用于 2GB 版本,而在 3B 上性能会有所降低。
先决条件#
要学习本教程,您需要一台 Raspberry Pi 4 或 5、配套的摄像头以及所有其他标准配件。
散热片和风扇(可选但建议配备)
5V 3A USB-C 电源
SD 卡(至少 8GB)
SD 卡读写器
Raspberry Pi 配置#
PyTorch 仅提供 Arm 64 位 (aarch64) 的 pip 包,因此您需要在 Raspberry Pi 上安装 64 位版本的操作系统。
您需要安装 官方 rpi-imager 来安装 Raspberry Pi OS。
32 位 Raspberry Pi OS 不适用。
安装至少需要几分钟,具体取决于您的网速和 SD 卡速度。完成后,界面应如下所示
现在将 SD 卡插入 Raspberry Pi,连接摄像头并启动系统。
Raspberry Pi 4 配置#
如果您使用的是 Raspberry Pi 4,则需要进行一些额外的配置更改。这些更改在 Raspberry Pi 5 上不是必需的。
系统启动并完成初始设置后,您需要编辑 /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
然后重启。
安装 PyTorch 和 picamera2#
PyTorch 以及我们需要的所有其他库都有 ARM 64 位/aarch64 版本,因此您可以直接通过 pip 安装,像在任何其他 Linux 系统上一样运行。
$ sudo apt install -y python3-picamera2 python3-libcamera
$ pip install torch torchvision --break-system-packages
现在我们可以检查一切是否安装正确。
$ python -c "import torch; print(torch.__version__)"
视频采集#
首先通过在终端运行 libcamera-hello 来测试摄像头是否工作正常。
对于视频采集,我们将使用 picamera2 来获取视频帧。
我们使用的模型(MobileNetV2)接收 224x224 大小的图像,因此我们可以直接向 picamera2 请求该尺寸,并设置为 36fps。我们的模型目标是 30fps,但我们请求略高的帧率以确保始终有足够的帧。
from picamera2 import Picamera2
picam2 = Picamera2()
# print available sensor modes
print(picam2.sensor_modes)
config = picam2.create_still_configuration(main={
"size": (224, 224),
"format": "BGR888",
}, display="main")
picam2.configure(config)
picam2.set_controls({"FrameRate": 36})
picam2.start()
要捕获帧,我们可以调用 capture_image,它会返回一个我们可以与 PyTorch 一起使用的 PIL.Image 对象。
# read frame
image = picam2.capture_image("main")
# show frame for testing
image.show()
此数据读取和处理过程大约耗时 3.5 ms。
图像预处理#
我们需要获取帧并将它们转换为模型预期的格式。这与您在任何机器上使用标准 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 |
真 |
mobilenet_v3_large |
29.3 |
34.1 |
30.7 |
真 |
resnet18 |
9.2 |
109.0 |
100.3 |
假 |
resnet50 |
4.3 |
233.9 |
225.2 |
假 |
resnext101_32x8d |
1.1 |
892.5 |
885.3 |
假 |
inception_v3 |
4.9 |
204.1 |
195.5 |
假 |
googlenet |
7.4 |
135.3 |
132.0 |
假 |
shufflenet_v2_x0_5 |
46.7 |
21.4 |
18.2 |
假 |
shufflenet_v2_x1_0 |
24.4 |
41.0 |
37.7 |
假 |
shufflenet_v2_x1_5 |
16.8 |
59.6 |
56.3 |
假 |
shufflenet_v2_x2_0 |
11.6 |
86.3 |
82.7 |
假 |
MobileNetV2:量化和 JIT#
为了获得最佳性能,我们需要一个经过量化和融合的模型。量化意味着它使用 int8 进行计算,这比标准的 float32 数学运算性能高得多。融合意味着连续的操作在可能的情况下被合并为性能更高的版本。通常,激活函数(如 ReLU)可以在推理期间合并到前一层(如 Conv2d)中。
PyTorch 的 aarch64 版本需要使用 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
from torchvision import models, transforms
from picamera2 import Picamera2
torch.backends.quantized.engine = 'qnnpack'
picam2 = Picamera2()
# print available sensor modes
print(picam2.sensor_modes)
config = picam2.create_still_configuration(main={
"size": (224, 224),
"format": "BGR888",
}, display="main")
picam2.configure(config)
picam2.set_controls({"FrameRate": 36})
picam2.start()
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
image = picam2.capture_image("main")
# 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 ...
print(output.argmax())
# 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
运行结果显示,在 Raspberry Pi 4 上约为 30 fps,在 Raspberry Pi 5 上约为 41 fps。
这是在 Raspberry Pi OS 默认设置下的结果。如果您禁用了 UI 以及默认启用的所有其他后台服务,性能和稳定性会更好。
如果我们查看 htop,会发现几乎 100% 的利用率。
为了验证其端到端的工作情况,我们可以计算类别的概率并 使用 ImageNet 类标签 来打印检测结果。
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 将使用所有可用的核心。如果 Raspberry Pi 上有其他后台程序运行,可能会与模型推理产生竞争,导致延迟抖动。为了缓解这种情况,您可以减少线程数,这会降低峰值延迟,但会以牺牲少量性能为代价。
torch.set_num_threads(2)
对于 shufflenet_v2_x1_5,使用 2 个线程而不是 4 个线程,将最佳情况下的延迟从 60 ms 增加到了 72 ms,但消除了 128 ms 的延迟抖动。
下一步#
您可以创建自己的模型或微调现有模型。如果您对 torchvision.models.quantized 中的模型之一进行微调,大部分融合和量化的工作已经为您完成,因此您可以直接部署并在 Raspberry Pi 上获得良好的性能。
了解更多