YOLOv8 训练报错:PyTorch 2.6+ 模型加载兼容性问题解决

发布于:2025-08-13 ⋅ 阅读:(19) ⋅ 点赞:(0)

问题

训练时报错,核心错误信息为_pickle.UnpicklingError: Weights only load failed,涉及torch.load函数的weights_only参数及DetectionModel类的加载限制

yolo8) yfzx@yfzx-4090:~/project/go2/yolov8_train$ python main.py
Ultralytics YOLOv8.3.176 🚀 Python-3.10.18 torch-2.8.0+cu128 CUDA:0 (NVIDIA GeForce RTX 4090, 24092MiB)
trainer: task=detect, mode=train, model=yolov8s.pt, data=coco128.yaml, epochs=300, time=None, patience=100, batch=16, imgsz=640, save=True, save_period=-1, cache=False, device=None, workers=8, project=None, name=train8, exist_ok=False, pretrained=True, optimizer=auto, verbose=True, seed=0, deterministic=True, single_cls=False, rect=False, cos_lr=False, close_mosaic=10, resume=False, amp=True, fraction=1.0, profile=False, freeze=None, multi_scale=False, overlap_mask=True, mask_ratio=4, dropout=0.0, val=True, split=val, save_json=False, save_hybrid=False, conf=None, iou=0.7, max_det=300, half=False, dnn=False, plots=True, source=None, vid_stride=1, stream_buffer=False, visualize=False, augment=False, agnostic_nms=False, classes=None, retina_masks=False, embed=None, save_dir=runs/detect/train8, show=False, save_frames=False, save_txt=False, save_conf=False, save_crop=False, show_labels=True, show_conf=True, show_boxes=True, line_width=None, format=torchscript, keras=False, optimize=False, int8=False, dynamic=False, simplify=True, opset=None, workspace=4, nms=False, lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=7.5, cls=0.5, dfl=1.5, pose=12.0, kobj=1.0, label_smoothing=0.0, nbs=64, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, bgr=0.0, mosaic=1.0, mixup=0.0, copy_paste=0.0, copy_paste_mode=flip, auto_augment=randaugment, erasing=0.4, crop_fraction=1.0, cfg=None, tracker=botsort.yaml
Traceback (most recent call last):
  File "/home/yfzx/project/go2/yolov8_train/main.py", line 7, in <module>
    trainer.train()
  File "/home/yfzx/project/go2/yolov8_train/yolov8/engine/trainer.py", line 198, in train
    self._do_train(world_size)
  File "/home/yfzx/project/go2/yolov8_train/yolov8/engine/trainer.py", line 312, in _do_train
    self._setup_train(world_size)
  File "/home/yfzx/project/go2/yolov8_train/yolov8/engine/trainer.py", line 226, in _setup_train
    ckpt = self.setup_model()
  File "/home/yfzx/project/go2/yolov8_train/yolov8/engine/trainer.py", line 530, in setup_model
    weights, ckpt = attempt_load_one_weight(model)
  File "/home/yfzx/project/go2/yolov8_train/yolov8/nn/tasks.py", line 806, in attempt_load_one_weight
    ckpt, weight = torch_safe_load(weight)  # load ckpt
  File "/home/yfzx/project/go2/yolov8_train/yolov8/nn/tasks.py", line 732, in torch_safe_load
    ckpt = torch.load(file, map_location="cpu")
  File "/home/yfzx/conda/anaconda/envs/yolo8/lib/python3.10/site-packages/torch/serialization.py", line 1529, 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([ultralytics.nn.tasks.DetectionModel])` or the `torch.serialization.safe_globals([ultralytics.nn.tasks.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参数默认值从False改为True。
参数设为True时,为安全起见仅允许加载权重数据,禁止执行模型文件中可能包含的自定义类代码。
YOLOv8 的模型文件(如yolov8s.pt)包含自定义的DetectionModel类序列化信息,因此被拦截。

解决方法

临时允许加载完整模型
修改 YOLOv8 代码中加载模型的部分,显式设置weights_only=False。
找到报错中提到的torch_safe_load函数(路径:yolov8/nn/tasks.py第 732 行),将:

ckpt = torch.load(file, map_location="cpu")

修改为

ckpt = torch.load(file, map_location="cpu", weights_only=False)  # 关闭安全检查

再次运行即可
在这里插入图片描述


网站公告

今日签到

点亮在社区的每一天
去签到