在树莓派4上进行实时推理(30帧每秒!)¶
创建于:2022年2月08日 | 最后更新:2024年1月16日 | 最后验证:2024年11月05日
作者: Tristan Rice
PyTorch 对 Raspberry Pi 4 提供了开箱即用的支持。本教程将指导您如何设置 Raspberry Pi 4 以运行 PyTorch,并在 CPU 上实时(30 fps+)运行 MobileNet v2 分类模型。
这一切都是在Raspberry Pi 4 Model B 4GB上测试的,但也应该适用于2GB版本以及在3B上,尽管性能会有所降低。
树莓派4设置¶
PyTorch 仅为 Arm 64位 (aarch64) 提供 pip 包,因此您需要在 Raspberry Pi 上安装 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
然后重新启动。重新启动后,视频4linux2设备 /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。picamera在64位的Raspberry Pi OS上不可用,而且它比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 返回一个 numpy 数组,格式为 BGR,因此我们需要读取并进行一些调整,以将其转换为预期的 RGB 格式。
ret, image = cap.read()
# convert opencv output from BGR to RGB
image = image[:, :, [2, 1, 0]]
此数据读取和处理大约需要 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,因为它具有良好的性能和准确性。
树莓派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)中。
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后,我们的帧率从约20fps提升到了约30fps。
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
运行它显示我们大约在30帧每秒。
这是在树莓派操作系统中使用所有默认设置的情况。如果您禁用了默认启用的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 毫秒增加到72 毫秒,但消除了
128 毫秒的延迟峰值。
下一步¶
您可以创建自己的模型或微调现有的模型。如果您在torchvision.models.quantized中的一个模型上进行微调,大部分融合和量化的工作已经为您完成,因此您可以直接在Raspberry Pi上部署并获得良好的性能。
查看更多:
Quantization 了解更多关于如何量化和融合模型的信息。
Transfer Learning Tutorial 关于如何使用迁移学习来微调预训练模型以适应您的数据集。