如何组合两个 YOLOv5 模型？答案

【问题标题】：How to combine two YOLOv5 models?如何组合两个 YOLOv5 模型？
【发布时间】：2021-12-01 02:17:43
【问题描述】：

我有两个模型，第一个用于对图像进行分类和裁剪。
从图像中裁剪类后，我将其发送到第二个模型，该模型对数字进行分类。
它们都是 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)

如何组合两个模型？

【问题讨论】：

标签： python pytorch yolo yolov5

【解决方案1】：

您需要修改模型架构的源代码，以防止第一个模型的输出被写为 numpy 类型，而是输出 pytorch 张量。除此之外，没有办法阻止 GPU->CPU->GPU 传输。

【讨论】：

我试过了，但是第二个不想得到pytorch张量，只有numpy数组
然后为网络重写 forward 函数的开头，以便它可以处理任一类型的输入。应该很简单。
能否添加示例代码或提供示例链接？
不是真的，在不知道您正在使用的 torch.hub 中的 Yolov5 的确切型号的情况下。这些模型往往是为现成的使用而设计的，并且不容易修改，因为在基础 CNN 之上有几个抽象层。我建议您找到需要类型为 numpy 数组并在传递 pytorch 数组时引发错误的代码行。你需要修改这一行。