在我的 Python 应用程序中，我使用 Detectron2 在图像上运行预测并检测图像中所有人的关键点。
我想在实时流式传输到我的应用程序中的帧上运行预测（使用 aiortc），但是我发现因为现在在新线程上运行（主线程被服务器占用），所以预测时间要长得多。
在线程上运行预测需要花费 1.5 到 4 秒不等，这太久了。
当在主线程上运行预测（没有视频流部分）时，我可以得到少于一秒的预测时间。
我的问题是为什么会发生这种情况，以及我该如何解决它？为什么从新线程使用 GPU 性能会下降得如此严重？
备注：
代码在 Google Colab 上测试，使用 Tesla P100 GPU，并通过从视频文件中读取帧来模拟视频流。
我使用这个问题中的代码计算运行帧预测所需的时间。
我尝试切换到 multiprocessing，但无法使用 cuda（我尝试过 import multiprocessing 和 import torch.multiprocessing，并使用 set_stratup_method('spawn')），在调用进程上的 start 时会卡住。
示例代码：
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
import threading
from typing import List
import numpy as np
import timeit
import cv2
# Prepare the configuration file
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set threshold for this model
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.DEVICE = "cuda"
predictor = DefaultPredictor(cfg)
def get_frames(video: cv2.VideoCapture):
frames = list()
while True:
has_frame, frame = video.read()
if not has_frame:
break
frames.append(frame)
return frames
class CodeTimer:
# Source: https://dev59.com/tGYq5IYBdhLWcg3weQgy#52749808
def __init__(self, name=None):
self.name = " '" + name + "'" if name else ''
def __enter__(self):
self.start = timeit.default_timer()
def __exit__(self, exc_type, exc_value, traceback):
self.took = (timeit.default_timer() - self.start) * 1000.0
print('Code block' + self.name + ' took: ' + str(self.took) + ' ms')
video = cv2.VideoCapture('DemoVideo.mp4')
num_frames = round(video.get(cv2.CAP_PROP_FRAME_COUNT))
frames_buffer = list()
predictions = list()
def send_frames():
# This function emulates the stream, so here we "get" a frame and add it to our buffer
for frame in get_frames(video):
frames_buffer.append(frame)
# Simulate delays between frames
time.sleep(random.uniform(0.3, 2.1))
def predict_frames():
predicted_frames = 0 # The number of frames predicted so far
while predicted_frames < num_frames: # Stop after we predicted all frames
buffer_length = len(frames_buffer)
if buffer_length <= predicted_frames:
continue # Wait until we get a new frame
# Read all the frames from the point we stopped
for frame in frames_buffer[predicted_frames:]:
# Measure the prediction time
with CodeTimer('In stream prediction'):
predictions.append(predictor(frame))
predicted_frames += 1
t1 = threading.Thread(target=send_frames)
t1.start()
t2 = threading.Thread(target=predict_frames)
t2.start()
t1.join()
t2.join()
- 你是对的，我刚刚注意到我过于简化了我的示例代码，它应该在另一个线程上发送帧。
- 我尝试过了，似乎确实已经初始化了，但仍然运行缓慢。
- 在
- SagiZivcfg中，我指定了预测器在cuda上运行，并且DefaultPredictor将帧移动到 GPU 上。AsyncPredictor在演示包中，所以导入不是那么容易。我正在检查当我安装 detectron 时是否有那个包。 - SagiZivYour session crashed after using all available RAM。它似乎一直在加载和加载东西而没有打印任何内容。我还尝试在10帧后停止预测而不是整个视频，但得到了相同的结果... - SagiZiv