只加载权重更安全,也能避免未来 PyTorch 默认行为变更带来的坑。
最近在训练 YOLO/Ultralytics 的同学可能都遇到过这样一行黄字提示:
C:\7zProject\SkinDiseaseTrain_v8\ultralytics\utils\torch_utils.py:521: FutureWarning: You are using torch.load with weights_only=False (the current default value), which uses the default pickle module implicitly.
It is possible to construct malicious pickle data which will execute arbitrary code during unpickling…
很多人第一反应是:“是不是模型坏了?”其实不是 bug,而是 PyTorch 在提醒你:模型权重文件如果来自不可信来源,真的可能“有毒”。
1.警告出现的原因
在 Ultralytics 内部,权重加载是这样写的:
x = torch.load(f, map_location=torch.device("cpu"))
因为没有传 weights_only=True,所以触发了 PyTorch 的 FutureWarning。
根本原因在于:
(1)weights_only=False(默认值)会调用 Python 的 pickle 去反序列化。
(2)pickle 能执行任意对象,如果权重文件被篡改,加载时可能执行恶意代码。
2.什么时候会出问题
目前这只是一个 警告,程序还能跑。但在未来版本中,PyTorch 将把默认值改成:
这会带来两种情况:
(1)推理/迁移学习:只需要权重张量 (state_dict),完全不受影响,反而更安全。
(2)继续训练:如果 checkpoint 里还有优化器、调度器、自定义对象,那么未来就会加载失败,除非你显式写 weights_only=False。
3.解决方案
显式开启 weights_only=True(立刻可用)
ckpt = torch.load(path, map_location="cpu", weights_only=True)
评论(0)