樹莓派 4 上的即時推理 (30 fps!)¶
創建於: Feb 08, 2022 | 最後更新於: Jan 16, 2024 | 最後驗證於: Nov 05, 2024
作者: Tristan Rice
PyTorch 對樹莓派 4 提供原生支援。本教程將指導你如何在樹莓派 4 上設定以執行 PyTorch,並在 CPU 上即時(30 fps+)執行 MobileNet v2 分類模型。
所有測試均在樹莓派 4 Model B 4GB 版本上進行,但也應適用於 2GB 版本以及效能降低的 3B 版本。
樹莓派 4 設定¶
PyTorch 只提供 Arm 64 位 (aarch64) 的 pip 包,因此你需要為你的樹莓派安裝一個 64 位的作業系統版本。
你可以從以下地址下載最新的 arm64 樹莓派作業系統:https://downloads.raspberrypi.org/raspios_arm64/images/ 並透過 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 版本,所以你可以透過 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 位樹莓派作業系統上不可用,而且比 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 |
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:量化和 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
執行結果顯示幀率穩定在約 30 fps。
這是在樹莓派作業系統所有預設設定下的表現。如果你停用 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 預設會使用所有可用的核心。如果你的樹莓派上有其他後臺程式執行,可能會與模型推理產生競爭,導致延遲激增。為了緩解此問題,你可以減少執行緒數量,這會以輕微的效能損失為代價降低峰值延遲。
torch.set_num_threads(2)
對於 shufflenet_v2_x1_5,使用 2 threads 而非 4 threads 會將最佳情況延遲從 60 ms 提高到 72 ms,但消除了高達 128 ms 的延遲激增。
下一步¶
你可以建立自己的模型或微調現有模型。如果你基於 torchvision.models.quantized 中的模型進行微調,大部分的融合和量化工作已經為你完成,這樣你就可以直接在樹莓派上以良好效能進行部署。
檢視更多