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

问题

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

bash 复制代码
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 行),将:

python 复制代码
ckpt = torch.load(file, map_location="cpu")

修改为

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

再次运行即可

相关推荐
龙卷风04051 小时前
深入理解Spring AI Alibaba多Agent系统:图结构驱动的智能协作
人工智能·后端
mqiqe1 小时前
【Spring AI MCP】四、MCP 服务端
java·人工智能·spring
好奇龙猫1 小时前
【AI学习-lora-定义-comfyUI相关-相关学习-了解概念(1)】
人工智能·学习
Jay20021111 小时前
【机器学习】10 正则化 - 减小过拟合
人工智能·机器学习
sxwuyanzu1 小时前
企业知识库的隐形危机:从“文档堆“到“知识系统“的进化之路
人工智能
5***79002 小时前
人工智能在环保监测中的数据分析
人工智能
算家计算2 小时前
芯片战打响!谷歌TPU挑战英伟达:AI算力战争背后的行业变局
人工智能·nvidia·芯片
技术支持者python,php2 小时前
训练模型,物体识别(opencv)
人工智能·opencv·计算机视觉
爱笑的眼睛112 小时前
深入理解MongoDB PyMongo API:从基础到高级实战
java·人工智能·python·ai
软件开发技术深度爱好者3 小时前
基于多个大模型自己建造一个AI智能助手
人工智能