背景:
剛剛,寫了《【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之前,建議備份當前環境和代碼。