在树莓派 4 上进行实时推理（30 帧/秒！）#
创建日期：2022 年 2 月 8 日 | 最后更新：2024 年 1 月 16 日 | 最后验证：2024 年 11 月 5 日
作者： Tristan Rice
PyTorch 对树莓派 4 原生支持。本教程将指导您如何设置树莓派 4 以运行 PyTorch，并在 CPU 上实时（30 帧/秒以上）运行 MobileNet v2 分类模型。
本教程使用树莓派 4 Model B 4GB 版本进行测试，但也能在 2GB 版本上运行，在 3B 版本上性能会降低。
先决条件#
要学习本教程，您需要一台树莓派 4、为其配备的摄像头以及所有其他标准配件。
散热片和风扇（可选但推荐）
5V 3A USB-C 电源
SD 卡（至少 8GB）
SD 卡读写器
树莓派 4 设置#
PyTorch 仅为 Arm 64 位（aarch64）提供 pip 包，因此您需要在树莓派上安装 64 位操作系统的版本。
您可以从 https://downloads.raspberrypi.org/raspios_arm64/images/ 下载最新的 arm64 树莓派操作系统，并使用 rpi-imager 进行安装。
32 位树莓派操作系统将无法工作。
安装过程至少需要几分钟，具体取决于您的网速和 SD 卡速度。完成后，它应该会显示如下内容：
现在是将 SD 卡插入树莓派、连接摄像头并启动。
启动后完成初始设置，您需要编辑 /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 版本，因此您可以像在任何其他 Linux 系统上一样通过 pip 安装它们。
$ pip install torch torchvision torchaudio
$ pip install opencv-python
$ pip install numpy --upgrade
现在我们可以检查所有安装是否正常。
$ python -c "import torch; print(torch.__version__)"
视频捕获#
对于视频捕获，我们将使用 OpenCV 来流式传输视频帧，而不是使用更常见的 picamera。 picamera 在 64 位树莓派操作系统上不可用，并且比 OpenCV 慢得多。OpenCV 直接访问 /dev/video0 设备来抓取帧。
我们使用的模型（MobileNetV2）需要 224x224 的图像尺寸，因此我们可以直接从 OpenCV 请求该尺寸，帧率为 36fps。我们的目标是模型的帧率为 30fps，但我们请求的帧率略高于此，以便始终有足够的帧。
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，因为它具有良好的性能和准确性。
树莓派 4 基准测试结果
模型 |
帧率 |
总时间（毫秒/帧） |
模型时间（毫秒/帧） |
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 的性能，而没有 JIT 编译则约为 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。
这是使用树莓派操作系统中的所有默认设置。如果您禁用了 UI 和默认启用的所有其他后台服务，它会更具性能和稳定性。
如果我们检查 htop，我们会看到 CPU 利用率接近 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 默认会使用所有可用的核心。如果您在树莓派上运行任何后台程序，它可能会与模型推理发生争用，导致延迟峰值。为了缓解这个问题，您可以减少线程数，这会以轻微的性能损失来降低峰值延迟。
torch.set_num_threads(2)
对于 shufflenet_v2_x1_5，使用 2 个线程 而不是 4 个线程，最佳延迟从 60 毫秒 提高到 72 毫秒，但消除了 128 毫秒 的延迟峰值。
后续步骤#
您可以创建自己的模型或微调现有模型。如果您从 torchvision.models.quantized 中选择的模型进行微调，那么融合和量化的大部分工作已经为您完成，因此您可以直接在树莓派上以良好的性能进行部署。
了解更多
