加载自定义训练YOLOv5模型时触发ValueError错误求助

问题描述 #

我尝试用以下代码加载自定义训练的YOLOv5模型：

# Model
model = torch.hub.load('/home/yolov5/runs/train/yolo_sign_det2/weights', 'best')  # or yolov5n - yolov5x6, custom

运行后出现如下错误：

ValueError                                Traceback (most recent call last)
<ipython-input-3-c832ab8c1eab> in <module>
      2 
      3 # Model
----> 4 model = torch.hub.load('/home/yolov5/runs/train/yolo_sign_det2/weights', 'best')  # or yolov5n - yolov5x6, custom
      5 
      6 # Images

~/.conda/envs/yolo/lib/python3.6/site-packages/torch/hub.py in load(repo_or_dir, model, source, force_reload, verbose, skip_validation, *args, **kwargs)
    395 
    396     if source == 'github':
--> 397         repo_or_dir = _get_cache_or_reload(repo_or_dir, force_reload, verbose, skip_validation)
    398 
    399     model = _load_local(repo_or_dir, model, *args, **kwargs)

~/.conda/envs/yolo/lib/python3.6/site-packages/torch/hub.py in _get_cache_or_reload(github, force_reload, verbose, skip_validation)
    163         os.makedirs(hub_dir)
    164     # Parse github repo information
--> 165     repo_owner, repo_name, branch = _parse_repo_info(github)
    166     # Github allows branch name with slash '/',
    167     # this causes confusion with path on both Linux and Windows.

~/.conda/envs/yolo/lib/python3.6/site-packages/torch/hub.py in _parse_repo_info(github)
    110     else:
    111         repo_info, branch = github, None
--> 112     repo_owner, repo_name = repo_info.split('/')
    113 
    114     if branch is None:

ValueError: too many values to unpack (expected 2)

错误原因及解决方法 #

错误原因 #

torch.hub.load() 默认会把传入的路径当作GitHub仓库地址解析，它尝试将路径按/分割成仓库所有者和仓库名，但你传入的是本地权重文件的目录路径，分割后得到的部分远多于2个，因此抛出ValueError。

正确加载方式 #

加载本地自定义训练的YOLOv5模型，有两种常用方法：

方法1：指定source='local'

明确告诉torch.hub.load()加载本地模型，第一个参数传入YOLOv5项目的根目录，第二个参数传入'custom'，并通过path参数指定权重文件的完整路径：

model = torch.hub.load('/home/yolov5', 'custom', path='/home/yolov5/runs/train/yolo_sign_det2/weights/best.pt', source='local')

方法2：直接使用YOLOv5内置加载函数

如果你的代码在YOLOv5项目目录下，可直接导入模型类并加载权重：

from models.common import DetectMultiBackend

model = DetectMultiBackend('/home/yolov5/runs/train/yolo_sign_det2/weights/best.pt')

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