python 我可以使用“torch.hub.load”load2型号吗

tjjdgumg  于 2023-04-10  发布在  Python
关注(0)|答案(1)|浏览(133)

我正在研究武器探测,希望使用目标探测和碰撞姿态估计,但我不能同时加载两个模型

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

如何将两者作为同一个过程使用?

bxgwgixi

bxgwgixi1#

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

无论这两个模型是什么,“models”模块总是被缓存,这取决于你的模型的依赖关系有多相似,你可能需要弹出更多的模块,才能使其工作。例如,你可能需要弹出“utils”模块。

相关问题