我可以使用“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”模块。