【yolo】yolo推理报错,以及解决方案
背景
我安装好了yolo8的环境后出现的报错,安装命令如下:
1.创建python环境
conda create -n yolo python==3.8
查看现有环境
conda env list
激活环境
conda activate yolo
2.安装库
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install ultralytics==8.1 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnx -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnxsim -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple
(最新的yolo8改用默认GPU版的onnxruntime,要安装一下下面的库,否则ONNX转换会有警告)
pip install onnxruntime-gpu -i https://pypi.tuna.tsinghua.edu.cn/simple
报错内容如下:
命令:
推理检测
yolo predict model=yolov8n.pt source=4.png
输入命令推理时报错:
(yolo8) F:\temp\yolotest>yolo predict model=yolov8n.pt source=2.bmp
Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "F:\envs\yolo8\Scripts\yolo.exe\__main__.py", line 7, in <module>
File "F:\envs\yolo8\Lib\site-packages\ultralytics\cfg\__init__.py", line 540, in entrypoint
model = YOLO(model, task=task)
^^^^^^^^^^^^^^^^^^^^^^
File "F:\envs\yolo8\Lib\site-packages\ultralytics\engine\model.py", line 95, in __init__
self._load(model, task)
File "F:\envs\yolo8\Lib\site-packages\ultralytics\engine\model.py", line 161, in _load
self.model, self.ckpt = attempt_load_one_weight(weights)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "F:\envs\yolo8\Lib\site-packages\ultralytics\nn\tasks.py", line 701, in attempt_load_one_weight
ckpt, weight = torch_safe_load(weight) # load ckpt
^^^^^^^^^^^^^^^^^^^^^^^
File "F:\envs\yolo8\Lib\site-packages\ultralytics\nn\tasks.py", line 634, in torch_safe_load
return torch.load(file, map_location="cpu"), file # load
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "F:\envs\yolo8\Lib\site-packages\torch\serialization.py", line 1470, in load
raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
_pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
(1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
(2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.
WeightsUnpickler error: Unsupported global: GLOBAL ultralytics.nn.tasks.DetectionModel was not an allowed global by default. Please use `torch.serialization.add_safe_globals([DetectionModel])` or the `torch.serialization.safe_globals([DetectionModel])` context manager to allowlist this global if you trust this class/function.
Check the documentation of torch.load to learn more about types accepted by default with weights_only https://pytorch.org/docs/stable/generated/torch.load.html.
解决方案
这个错误是由于 PyTorch 2.6 及以上版本中 torch.load
的默认行为发生了变化,weights_only
参数默认设置为 True
,导致加载 YOLO 模型权重时出现问题。以下是解决该问题的方法:
修改 torch.load
参数
在 YOLO 的代码中,找到 torch.load
的调用位置(通常是 ultralytics/nn/tasks.py
文件),将 weights_only
参数设置为 False
。具体步骤如下:
-
打开文件:
# 找到文件路径 F:\envs\yolo8\Lib\site-packages\ultralytics\nn\tasks.py
-
找到
torch.load
的调用位置(大约在第 634 行):return torch.load(file, map_location="cpu"), file # load
-
修改为:
return torch.load(file, map_location="cpu", weights_only=False), file # load
改完后,立竿见影,推理成功:
这说明环境搭建成功!