如何組合兩個YOLOv5模型？

Question

我有兩個模型，第一個是用於對圖像進行分類和裁剪到類。
從圖像中裁剪類后，我將其發送到第二個模型，該模型對數字進行分類。
它們都是 Yolo v5 型號。
但問題是我不能直接從 GPU 發送第二個。
首先我需要裁剪，我會得到一個 NumPy 數組。 收到 NumPy 數組后，我將其發送到第二個數組。
我想停止浪費時間將 NumPy 轉換為張量，反之亦然

model = torch.hub.load('.', 'custom', path=img_cls_path, source='local', force_reload=True)
model_ocr = torch.hub.load('.', 'custom', path=ocr_path, source='local', force_reload=True)
cap = cv2.VideoCapture(some_video_path)

while(cap.isOpened()):
    ret, frame = cap.read()
    results = model(frame)
    crops = results.crop(save=False)    
    for crop in crops:
        if 'number' in crop['label']:
            ocr_result = model_ocr(crop['im'])
            ocr_crop = ocr_result.crop(save=False)

如何組合兩個模型？

Answer 1

您將需要修改模型架構的源代碼，以防止第一個模型的輸出被寫入 numpy 類型，而是輸出 pytorch 張量。 除此之外，沒有辦法阻止 GPU->CPU->GPU 傳輸。