PyTorch 2.6+ 下 YOLOv8 模型加载失败?手把手教你解决 weights_only 报错
最近在升级到 PyTorch 2.6 或更高版本后,不少 YOLOv8 用户遇到了模型加载失败的问题。错误信息通常包含_pickle.UnpicklingError: Weights only load failed,这确实让人头疼。作为一位长期使用 YOLOv8 进行目标检测的开发者,我也踩过这个坑。今天就来详细分析这个问题,并分享几种可靠的解决方案。
1. 问题根源解析
这个问题的核心在于 PyTorch 2.6 引入的一项安全变更。在之前的版本中,torch.load()函数的weights_only参数默认值为False,这意味着它可以加载包含任意 Python 对象的模型文件。但从 PyTorch 2.6 开始,这个默认值改为了True,只允许加载纯权重数据,禁止执行模型文件中可能包含的任何代码。
YOLOv8 的模型文件(如 yolov8s.pt)包含了自定义的DetectionModel类的序列化信息。当weights_only=True时,PyTorch 的安全机制会阻止加载这些自定义类,从而抛出错误。
关键错误信息通常包含以下内容:
_pickle.UnpicklingError: Weights only load failed. This file can still be loaded... Unsupported global: GLOBAL ultralytics.nn.tasks.DetectionModel was not an allowed global by default2. 解决方案一:临时关闭安全检查
最直接的解决方法是在加载模型时显式设置weights_only=False。这适用于你完全信任模型来源的情况。
操作步骤:
- 找到 YOLOv8 代码中加载模型的位置(通常在
yolov8/nn/tasks.py文件中的torch_safe_load函数) - 修改加载代码:
# 原始代码 ckpt = torch.load(file, map_location="cpu") # 修改为 ckpt = torch.load(file, map_location="cpu", weights_only=False)注意:这种方法虽然简单,但会降低安全性。只应在你完全信任模型文件来源时使用。
3. 解决方案二:安全地允许特定类
PyTorch 提供了一种更安全的方式来处理这种情况 - 通过add_safe_globals显式允许特定的自定义类。
具体实现:
from torch import serialization from ultralytics.nn.tasks import DetectionModel # 在加载模型前添加以下代码 serialization.add_safe_globals([DetectionModel]) # 然后正常加载模型 model = torch.load("yolov8s.pt", weights_only=True)这种方法的好处是既保持了weights_only的安全检查,又允许加载你明确信任的自定义类。
4. 解决方案三:使用上下文管理器
如果你不想永久修改全局设置,可以使用上下文管理器的方式临时允许特定类:
from torch import serialization from ultralytics.nn.tasks import DetectionModel with serialization.safe_globals([DetectionModel]): model = torch.load("yolov8s.pt", weights_only=True)这种方式更加灵活,作用范围仅限于上下文块内部,不会影响其他代码。
5. 解决方案四:降级 PyTorch 版本
如果上述方法都不适合你的场景,可以考虑暂时降级 PyTorch 版本:
pip install torch==2.5.0不过这不是长期解决方案,建议尽快适配新版本的 PyTorch。
6. 最佳实践建议
根据不同的使用场景,我推荐以下策略:
| 场景 | 推荐方案 | 优点 | 缺点 |
|---|---|---|---|
| 开发环境 | 方案二或三 | 保持安全性 | 需要少量代码修改 |
| 生产环境 | 方案二 | 平衡安全与功能 | 需要代码审核 |
| 临时使用 | 方案一 | 简单直接 | 安全性较低 |
| 长期兼容 | 等待 Ultralytics 官方更新 | 最规范 | 可能需要等待 |
在实际项目中,我通常会选择方案二,因为它提供了良好的安全性和兼容性平衡。记得在代码中添加适当的注释,说明为什么需要这样做。
7. 深入理解 weights_only 机制
为了更好地解决这类问题,了解 PyTorch 的weights_only机制很有帮助。当weights_only=True时,PyTorch 只允许加载以下类型的数据:
- 基本 Python 类型(int, float, str, list, dict 等)
- torch.Tensor
- numpy.ndarray
- 特定允许的 torch 类
任何自定义类都需要显式添加到安全列表才能加载。这种机制有效防止了恶意模型文件执行任意代码的风险。
常见需要允许的 YOLOv8 类:
ultralytics.nn.tasks.DetectionModelultralytics.nn.modules.Convultralytics.nn.modules.Bottleneck
8. 未来兼容性考虑
随着 PyTorch 安全机制的不断加强,建议关注以下几点:
- 定期更新 Ultralytics 库,官方可能会发布兼容性修复
- 在 CI/CD 流程中加入 PyTorch 版本兼容性测试
- 考虑将自定义模型结构分离到单独的模块中,便于管理安全列表
我在最近的一个项目中就遇到了类似问题,通过建立允许类清单的方式,系统化地解决了兼容性问题。具体做法是创建一个safe_classes.py文件,集中管理所有需要允许的类,然后在应用启动时统一注册。