背景:
刚刚,写了《【yolo】yolo推理报错,以及解决方案》,马上训练就遇到类似的报错。
我对我标注的图像进行了300轮的训练,但是训练完300轮后,报错了。。。
报错信息
300 epochs completed in 0.085 hours.
Traceback (most recent call last):File "<frozen runpy>", line 198, in _run_module_as_mainFile "<frozen runpy>", line 88, in _run_codeFile "F:\envs\yolo8\Scripts\yolo.exe\__main__.py", line 7, in <module>File "F:\envs\yolo8\Lib\site-packages\ultralytics\cfg\__init__.py", line 567, in entrypointgetattr(model, mode)(**overrides) # default args from model^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^File "F:\envs\yolo8\Lib\site-packages\ultralytics\engine\model.py", line 390, in trainself.trainer.train()File "F:\envs\yolo8\Lib\site-packages\ultralytics\engine\trainer.py", line 208, in trainself._do_train(world_size)File "F:\envs\yolo8\Lib\site-packages\ultralytics\engine\trainer.py", line 468, in _do_trainself.final_eval()File "F:\envs\yolo8\Lib\site-packages\ultralytics\engine\trainer.py", line 615, in final_evalstrip_optimizer(f) # strip optimizers^^^^^^^^^^^^^^^^^^File "F:\envs\yolo8\Lib\site-packages\ultralytics\utils\torch_utils.py", line 484, in strip_optimizerx = torch.load(f, map_location=torch.device("cpu"))^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^File "F:\envs\yolo8\Lib\site-packages\torch\serialization.py", line 1470, in loadraise 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的torch.load函数无法正确加载权重文件。具体原因是PyTorch 2.6及以上版本默认启用了weights_only=True,限制了加载的全局对象类型,而权重文件中包含了一个不被允许的全局对象(ultralytics.nn.tasks.DetectionModel)。
这个原因和推理报错类似。
从错误信息来看,问题出在训练完成后加载模型权重时,PyTorch的torch.load
函数无法正确加载权重文件。具体原因是PyTorch 2.6及以上版本默认启用了weights_only=True
,限制了加载的全局对象类型,而权重文件中包含了一个不被允许的全局对象(ultralytics.nn.tasks.DetectionModel
)。
-
PyTorch的
weights_only
机制
PyTorch 2.6及以上版本默认启用了weights_only=True
,这是一种安全机制,防止加载不受信任的权重文件时执行任意代码。如果你的权重文件中包含不被允许的全局对象(如自定义类或函数),就会触发此错误。 -
权重文件问题
你的权重文件中可能包含了一些自定义的类或函数(如ultralytics.nn.tasks.DetectionModel
),这些对象不被weights_only=True
允许加载。
解决方法,详细步骤
以下是具体的操作步骤:
- 找到加载权重的代码
在你的错误日志中,加载权重的代码位于ultralytics/utils/torch_utils.py
的第484行:x = torch.load(f, map_location=torch.device("cpu"))
如果你是新建的虚拟环境,直接安装的pytorch的包,那么源代码会在虚拟环境下:
F:\envs\yolo8\Lib\site-packages\ultralytics/utils/torch_utils.py
- 修改加载代码
将上述代码修改为:x = torch.load(f, map_location=torch.device("cpu"), weights_only=False)
- 重新运行训练
保存修改后的代码,并重新运行训练脚本。
修改后立竿见影,不再报错:
注意事项
- 安全性:将
weights_only
设置为False
或使用add_safe_globals
会降低安全性,请确保权重文件来自可信来源。 - 兼容性:如果降级PyTorch版本,请确保其他依赖库与新版本兼容。
- 备份:在修改代码或降级PyTorch之前,建议备份当前环境和代码。