Могу ли я использовать модель «torch.hub.load» для загрузки 2?

Question

Я работаю над обнаружением оружия и хочу использовать обнаружение объектов и оценку позы для столкновения, но я не могу загрузить обе модели одновременно

      

WeaponDetectionModel = torch.hub.load('ultralytics/yolov5', 'custom', path='Model_WeaponDetection//YoloV5Weapon.pt')

HumanPose = torch.hub.load('TexasInstruments/edgeai-yolov5', 'custom', path='Model_HumanPose//YoloV5Human.pt')

Как я могу использовать оба как один и тот же процесс?

Answer 1

PyTorch Hub имеет ограничение на загрузку более одной модели из разных репозиториев. Я предполагаю, что вы получаете ошибки при загрузке второй модели.

Это связано с тем, что при загрузке первой модели модули импортируются в файлhubconf, а при попытке загрузки второй модели некоторые модули все еще доступны в кеше модулей.

Итак, грязный обходной путь — просто очистить импортированные модули:

      import torch
import sys

WeaponDetectionModel = torch.hub.load('ultralytics/yolov5', 'custom', path='Model_WeaponDetection//YoloV5Weapon.pt')

sys.modules.pop('models')  # pop any other possibly cached models 

HumanPose = torch.hub.load('TexasInstruments/edgeai-yolov5', 'custom', path='Model_HumanPose//YoloV5Human.pt')

Независимо от того, что это за две модели, модуль «модели» всегда кэшируется. В зависимости от того, насколько похожи зависимости ваших моделей, вам может понадобиться больше модулей, чтобы это работало. например, вам также может потребоваться загрузить модуль «utils».