PyTorch 入门指南
学习 PyTorch
图像和视频
音频
后端
强化学习
在生产环境中部署 PyTorch 模型
Profiling PyTorch
代码变换与FX
前端API
扩展 PyTorch
模型优化
并行和分布式训练
边缘端的 ExecuTorch
推荐系统
多模态

在 Raspberry Pi 4 上实现实时推理(30 帧/秒!)

作者: Tristan Rice

PyTorch 原生支持 Raspberry Pi 4。本教程将指导您如何为 Raspberry Pi 4 设置 PyTorch,并在 CPU 上实时(30 fps+)运行 MobileNet v2 分类模型。

所有测试均在 Raspberry Pi 4 Model B 4GB 上进行,但同样适用于 2GB 版本,以及在 3B 上运行时性能会有所降低。

https://user-images.githubusercontent.com/909104/153093710-bc736b6f-69d9-4a50-a3e8-9f2b2c9e04fd.gif

前提条件

要完成本教程,您需要准备一台 Raspberry Pi 4、一个兼容的摄像头以及所有其他标准配件。

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 将无法运行。

https://user-images.githubusercontent.com/909104/152866212-36ce29b1-aba6-4924-8ae6-0a283f1fca14.gif

安装过程至少需要几分钟,具体时间取决于您的网络速度和 SD 卡速度。完成后,界面应如下所示:

https://user-images.githubusercontent.com/909104/152867425-c005cff0-5f3f-47f1-922d-e0bbb541cd25.png

接下来,将 SD 卡插入 Raspberry Pi,连接摄像头并启动设备。

https://user-images.githubusercontent.com/909104/152869862-c239c980-b089-4bd5-84eb-0a1e5cf22df2.png

设备启动并完成初始设置后,您需要编辑 /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 系统上一样正常工作。

$pipinstalltorchtorchvisiontorchaudio
$pipinstallopencv-python
$pipinstallnumpy--upgrade

https://user-images.githubusercontent.com/909104/152874260-95a7a8bd-0f9b-438a-9c0b-5b67729e233f.png

我们现在可以检查是否所有内容都已正确安装:

$python-c"import torch; print(torch.__version__)"

https://user-images.githubusercontent.com/909104/152874271-d7057c2d-80fd-4761-aed4-df6c8b7aa99f.png

视频捕获

对于视频捕获,我们将使用 OpenCV 来流式传输视频帧,而不是更常见的 picamera。picamera 在 64 位 Raspberry Pi OS 上不可用,并且它比 OpenCV 慢得多。OpenCV 直接访问 /dev/video0 设备来获取帧。

我们使用的模型(MobileNetV2)接受的图像大小为 224x224,因此我们可以直接从 OpenCV 以 36fps 的帧率请求图像。我们的目标是让模型以 30fps 运行,但我们请求的帧率略高于此,以确保始终有足够的帧。

importcv2
fromPILimport 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 变换进行处理是相同的。

fromtorchvisionimport 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: 量化与即时编译

为了获得最佳性能,我们需要一个经过量化和融合的模型。量化意味着模型使用 int8 进行计算,这比标准的 float32 数学运算性能更高。融合意味着连续的运算尽可能被合并成一个更高效的版本。通常在推理过程中,像激活函数(ReLU)这样的操作可以合并到前一层(如 Conv2d)中。

aarch64 版本的 PyTorch 需要使用 qnnpack 引擎。

importtorch
torch.backends.quantized.engine = 'qnnpack'

在本例中,我们将使用由 torchvision 提供的预量化和融合版本的 MobileNetV2。

fromtorchvisionimport models
net = models.quantization.mobilenet_v2(pretrained=True, quantize=True)

然后我们希望将模型进行 JIT 编译,以减少 Python 的开销并融合任何操作。使用 JIT 后,帧率从大约 20fps 提升到了大约 30fps。

net = torch.jit.script(net)

整合所有部分

现在我们可以将所有部分整合在一起并运行它:

importtime

importtorch
importnumpyasnp
fromtorchvisionimport models, transforms

importcv2
fromPILimport 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 fps。

https://user-images.githubusercontent.com/909104/152892609-7d115705-3ec9-4f8d-beed-a51711503a32.png

这是在使用 Raspberry Pi OS 的所有默认设置的情况下。如果您禁用了默认启用的 UI 和其他后台服务,性能会更高效且更稳定。

如果我们查看 htop,会发现 CPU 利用率接近 100%。

https://user-images.githubusercontent.com/909104/152892630-f094b84b-19ba-48f6-8632-1b954abc59c7.png

为了验证它是否端到端正常工作,我们可以计算类别的概率,并使用 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 实时运行:

https://user-images.githubusercontent.com/909104/153093710-bc736b6f-69d9-4a50-a3e8-9f2b2c9e04fd.gif

检测橙子:

https://user-images.githubusercontent.com/909104/153092153-d9c08dfe-105b-408a-8e1e-295da8a78c19.jpg

检测杯子:

https://user-images.githubusercontent.com/909104/153092155-4b90002f-a0f3-4267-8d70-e713e7b4d5a0.jpg

故障排除:性能问题

默认情况下,PyTorch 会使用所有可用的核心。如果您在 Raspberry Pi 上运行了后台任务,可能会导致与模型推理的资源竞争,从而引发延迟峰值。为了缓解这个问题,您可以减少线程数量,这将会以轻微的性能损失为代价来降低峰值延迟。

torch.set_num_threads(2)

对于 shufflenet_v2_x1_5,使用 2 线程 而不是 4 线程 会将最佳情况下的延迟从 60 毫秒 增加到 72 毫秒,但消除了 128 毫秒 的延迟峰值。

下一步

您可以创建自己的模型或微调现有的模型。如果您在torchvision.models.quantized中的某个模型上进行微调,大部分融合和量化的工作已经为您完成,因此您可以直接在Raspberry Pi上部署并获得良好的性能。

了解更多:

  • 量化 了解更多关于如何量化和融合模型的信息。

  • 迁移学习教程 了解如何使用迁移学习在现有模型的基础上针对您的数据集进行微调。

本页目录