#python #pytorch #yolo #yolov5
Вопрос:
У меня есть две модели, первая из которых предназначена для классификации изображений и обрезки по классам.
После обрезки классов с изображения я отправляю его во вторую модель, которая классифицирует цифры.
Обе они модели 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) Как объединить две модели?
Ответ №1:
Вам нужно будет изменить исходный код для архитектуры модели, чтобы выходные данные первой модели не записывались как тип numpy, а вместо этого выводили тензор pytorch. За исключением этого, нет никакого способа предотвратить передачу GPU->CPU->>GPU.
Комментарии:
1. Я попробовал, но второй не хочет получать тензор pytorch, только массив numpy
2. Затем перепишите начало
forwardфункции для сети, чтобы она могла обрабатывать любой тип ввода. Должно быть очень просто.3. Можете ли вы добавить пример кода или дать ссылку с примером, пожалуйста?
4. Не совсем, не зная точной модели Yolov5 из torch.hub, которую вы используете. Эти модели, как правило, предназначены для использования в готовом виде и их не так просто модифицировать, так как поверх базового CNN имеется несколько уровней абстракции. Я предлагаю вам найти строку кода, которая ожидает массив типа numpy и выдает ошибку при передаче массива pytorch. Вам нужно изменить эту строку.