Can i use "torch.hub.load" load 2 model

Question

i working on weapon detection and want to use Object Detection and Pose Estimation for collision but i cant load both model same time

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')

How can i use both as same process ?

Answer 1

PyTorch Hub has the limitation of loading more than one model from different repositories. I am guessing that you get errors while loading the second model.

This is because of the fact that when you load the first model, modules are imported into the hubconf file, and when you try to load the second model, some modules are still available in the module cache.

So, a dirty workaround is to just clear the imported modules:

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')

No matter what the two models are, the 'models' module is always cached, depending on how similar the dependencies of your models are, there might be more modules that you need to pop, to make this work. e.g. you might need to pop the 'utils' module as well.