如何在同一进程中同时加载两个torch.hub模型实现武器检测与姿态估计？

问题核心 #

你遇到的问题源于两个模型分别基于不同仓库的YOLOv5实现，torch.hub默认的缓存机制和命名空间可能导致冲突，无法同时加载。以下是几种可行的解决方法：

方法1：指定独立缓存目录并强制重载

通过给两个模型分配不同的缓存目录，避免仓库文件冲突，同时强制重新加载确保环境隔离：

import torch

# 加载武器检测模型，指定专属缓存目录
WeaponDetectionModel = torch.hub.load(
    'ultralytics/yolov5', 
    'custom', 
    path='Model_WeaponDetection//YoloV5Weapon.pt',
    force_reload=True,
    cache_dir='./hub_cache/ultralytics_yolov5'
)

# 加载姿态估计模型，使用不同的缓存目录
HumanPose = torch.hub.load(
    'TexasInstruments/edgeai-yolov5', 
    'custom', 
    path='Model_HumanPose//YoloV5Human.pt',
    force_reload=True,
    cache_dir='./hub_cache/edgeai_yolov5'
)

cache_dir参数将两个仓库的缓存文件分开存储，force_reload=True会跳过本地缓存，重新下载并构建对应仓库的模型环境，避免版本冲突。

方法2：本地克隆仓库后加载

如果hub远程加载仍有冲突，可手动克隆两个仓库到本地，从本地路径加载模型，实现完全隔离：

1. 手动克隆仓库到本地指定目录：
   • 克隆ultralytics/yolov5到./local_models/ultralytics_yolov5
   • 克隆TexasInstruments/edgeai-yolov5到./local_models/edgeai_yolov5
2. 从本地加载模型：

import torch

# 从本地目录加载武器检测模型
WeaponDetectionModel = torch.hub.load(
    './local_models/ultralytics_yolov5', 
    'custom', 
    path='Model_WeaponDetection//YoloV5Weapon.pt',
    source='local'
)

# 从本地目录加载姿态估计模型
HumanPose = torch.hub.load(
    './local_models/edgeai_yolov5', 
    'custom', 
    path='Model_HumanPose//YoloV5Human.pt',
    source='local'
)

source='local'参数让torch.hub直接读取本地仓库代码，彻底避免远程仓库的依赖冲突。

方法3：排查设备资源问题

如果是GPU显存不足导致加载失败，可先切换到CPU测试：

import torch

# 强制使用CPU加载模型，排除显存问题
WeaponDetectionModel = torch.hub.load('ultralytics/yolov5', 'custom', path='Model_WeaponDetection//YoloV5Weapon.pt', device='cpu')
HumanPose = torch.hub.load('TexasInstruments/edgeai-yolov5', 'custom', path='Model_HumanPose//YoloV5Human.pt', device='cpu')

若CPU能正常加载，再考虑优化显存使用：比如启用半精度推理（model.half()）、使用更小的模型版本，或针对多GPU环境做模型并行处理。

内容的提问来源于stack exchange，提问作者Kaemaros