目录

实时推理在树莓派 4 上(30 fps!)

创建日期: 2022年2月8日 | 最后更新日期: 2024年1月16日 | 最后验证日期: 2024年11月5日

作者: 特里斯坦·赖斯

PyTorch 对 Raspberry Pi 4 提供了开箱即用的支持。本教程将指导你如何为运行 PyTorch 设置 Raspberry Pi 4,并在 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、一个摄像头以及所有其他标准配件。

树莓派 4 设置

PyTorch 只提供 Arm 64 位(aarch64)的 pip 包,因此你需要在 Raspberry Pi 上安装 64 位版本的操作系统。

您可以从 https://downloads.raspberrypi.org/raspios_arm64/images/ 下载最新的 arm64 版 Raspberry Pi 操作系统,并通过 rpi-imager 进行安装。

32位Raspberry Pi操作系统无法使用。

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

然后重新启动。重新启动后,视频4linux2设备 /dev/video0 应该存在。

安装 PyTorch 和 OpenCV

PyTorch 及我们所需的所有其他库都有 ARM 64 位/aarch64 版本,因此你可以通过 pip 安装它们,并使其像其他任何 Linux 系统一样正常工作。

$ pip install torch torchvision torchaudio
$ pip install opencv-python
$ pip install numpy --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来流式传输视频帧,而不是更常见的picamerapicamera在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 基准测试结果:

模型

FPS

总时间(毫秒/帧)

模型时间(毫秒/帧)

qnnpack 预训练模型

mobilenet_v2

33.7

29.7

26.4

mobilenet_v3_large

29.3

34.1

30.7

ResNet 18模型

9.2

109.0

100.3

假的

ResNet 50模型

4.3

233.9

225.2

假的

resnext101_32x8d

1.1

892.5

885.3

假的

inception_v3

4.9

204.1

195.5

假的

谷歌网络

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
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 帧每秒左右。

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

这是在树莓派操作系统中的所有默认设置下进行的。如果你禁用了UI以及所有默认启用的其他后台服务,性能会更好且更稳定。

如果检查 htop,我们会发现利用率几乎达到了 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 默认会使用所有可用的 CPU 核心。如果在树莓派上后台运行其他任务,可能会与模型推理产生竞争,导致延迟波动。为了缓解这种情况,可以减少线程数量,从而降低峰值延迟,但会带来轻微的性能损失。

torch.set_num_threads(2)

对于 shufflenet_v2_x1_5 使用 2 threads 而不是 4 threads 将最佳情况下的延迟增加到 72 ms,从 60 ms 开始,但消除了 128 ms 的延迟峰值。

下一步

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

查看更多:

  • 量化 以了解如何量化和融合您的模型。

  • 迁移学习教程 介绍如何使用迁移学习来微调现有模型以适应您的数据集。

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源