PyTorch 2.6问题解决：torch.load函数weights_only参数详解
1. 问题背景与现象
在PyTorch 2.6及以上版本中，许多用户在加载模型权重时遇到了类似如下的错误：
RuntimeError: unsupported pickle protocol: 5或者更具体的错误信息：
RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False这些错误通常发生在使用torch.load()函数加载预训练模型时，特别是当模型包含自定义类或复杂结构时（如YOLO系列模型）。问题的根源在于PyTorch 2.6引入的安全加载策略变更。
2. weights_only参数的作用机制
2.1 安全加载的背景
PyTorch 2.6对torch.load()函数进行了安全增强，新增了weights_only参数，默认设置为True。这个改变是为了防止潜在的恶意代码执行风险，因为模型文件本质上是通过Python的pickle机制序列化的，而pickle可以执行任意代码。
当weights_only=True时，torch.load()会：
- 仅允许加载包含张量、数字、字符串、列表、字典等基本类型的对象
- 禁止加载包含自定义类、函数或任意可执行代码的对象
- 对加载内容进行严格检查，确保不会执行任何潜在危险的代码
2.2 参数的具体影响
下表展示了weights_only参数不同设置下的行为差异：
| 参数值 | 允许加载的内容 | 安全性 | 适用场景 |
|---|---|---|---|
True | 基本数据类型（张量、数字、字符串等） | 高 | 加载纯权重文件、官方预训练模型 |
False | 任意可pickle对象（包括自定义类） | 低 | 加载包含自定义结构的完整模型 |
3. 问题解决方案
3.1 修改加载代码
对于大多数遇到问题的场景，最简单的解决方案是在调用torch.load()时显式设置weights_only=False：
model = torch.load('model.pth', weights_only=False)如果模型需要加载到特定设备上，可以同时指定map_location：
model = torch.load('model.pth', map_location='cuda:0', weights_only=False)3.2 YOLO系列模型的特殊处理
对于Ultralytics YOLO等框架，通常需要修改框架内部的加载代码。如参考博文中提到的，找到strip_optimizer函数中的加载代码：
# 修改前 x = torch.load(f, map_location=torch.device("cpu")) # 修改后 x = torch.load(f, map_location=torch.device("cpu"), weights_only=False)3.3 其他兼容性考虑
如果模型需要在不同PyTorch版本间共享，可以考虑以下策略：
- 导出纯权重：使用
torch.save(model.state_dict(), ...)只保存权重 - 版本说明：在模型文件中包含PyTorch版本要求
- 格式转换：将模型转换为ONNX等跨框架格式
4. 安全使用建议
虽然设置weights_only=False可以解决问题，但需要注意安全风险：
- 仅加载可信来源：确保模型文件来自可信赖的官方渠道
- 沙箱环境测试：首次加载未知模型时，在隔离环境中测试
- 最小权限原则：使用普通用户权限而非root权限运行加载代码
- 签名验证：对重要模型文件进行数字签名验证
对于生产环境，推荐的工作流程是：
try: # 先尝试安全加载 model = torch.load('model.pth', weights_only=True) except RuntimeError: # 安全加载失败时再降级 if is_trusted_source('model.pth'): model = torch.load('model.pth', weights_only=False) else: raise ValueError("Untrusted model file")5. 总结
PyTorch 2.6引入的weights_only参数是一项重要的安全改进，但在处理复杂模型结构时可能带来兼容性问题。通过理解其工作机制和合理设置参数，可以平衡安全性和功能性需求。关键要点包括：
weights_only=True是PyTorch 2.6+的默认设置，提供更高安全性- 加载包含自定义类的模型时需要设置
weights_only=False - 对于YOLO等框架，可能需要修改框架内部的加载代码
- 始终注意模型文件来源的可信度，平衡安全与功能
在实际应用中，建议根据具体场景选择最适合的加载策略，并在安全性和兼容性之间取得平衡。
获取更多AI镜像
想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。
