☰
PyTorch 2.x 的 torch.load FutureWarning 与 weights_only 深度解析
2026/10/14 2:31:30 网站建设 项目流程

升级到 PyTorch 2.x 后,我第一次在自己电脑上跑旧训练脚本,终端突然冒出一行明晃晃的FutureWarning: You are using weights_only=False...。当时第一反应是:我没动过模型加载的代码啊,怎么平白无故多了一句警告?等到后面几次在不同机器上部署,发现同样的问题在同事那边也出现,才意识到这不是偶发,而是 PyTorch 版本升级带来的行为变化。

这篇文章我想把torch.load的FutureWarning和weights_only=False这个参数掰开揉碎讲清楚。包括这个警告从哪来、为什么会默认禁用完整反序列化、什么场景下必须显式传weights_only=False、以及如何正确迁移到新的加载方式。适合正在维护旧代码库的 PyTorch 用户,也适合刚上手、把torch.load当普通文件读取函数用的新手。

1. 一次升级之后,代码开始对我喊话:FutureWarning 的现场还原

1.1 一个看似无关紧要的警告是怎么冒出来的

假设你有一个 PyTorch 项目,代码里用常规方式保存和加载模型:

import torch model = torch.nn.Linear(10, 5) torch.save(model.state_dict(), "model.pth") state_dict = torch.load("model.pth") print(state_dict.keys())

这段代码在 PyTorch 1.x 时代完全安静,没有任何提示。但到了 PyTorch 2.6 及更高版本,运行时会多出一行警告:

FutureWarning: You are using `weights_only=False` for the current loads. Please be aware that the load is unsafe and the weights may be exploited. FutureWarning: You are using `weights_only=False` for the current loads. Re-enable `weights_only=True` to avoid the risk.

第一眼看去,警告措辞相当严重,像是在指责你“正在做危险的事情”。但实际上,在 PyTorch 2.6 之前,torch.load的默认行为就是weights_only=False,也就是说你什么都没改动,只是换了一个新版本的解释器,默认行为变了。官方之所以加这个警告,是希望在下一个大版本里把默认值从False正式翻转成True,提前给开发者一个过渡期。

1.2 为什么说这是“行为迁移”而不是“代码错误”

很多项目在升级 PyTorch 后,第一反应是去查代码哪里写错了。但实际上这不是 bug,而是 API 的默认语义发生了变化。PyTorch 官方采取的方式是“渐进式默认值迁移”:先在某个版本让旧行为继续可用,但打印警告;等开发者社区普遍调整后,再在后续版本切换默认值。

这种做法在大型框架里很常见。它考虑到了生态里仍然有大量旧代码、旧脚本。如果直接一刀切把默认值改成True,很多依赖旧行为的项目在升级后直接崩溃,连警告的机会都没有。所以FutureWarning的作用就是给你一个缓冲期:现在还能跑,但请尽快适配。

这里要提醒一句:不要用“忽略警告”的方式混过去。我在某个生产推理服务里见过同事在启动脚本里直接加了warnings.filterwarnings("ignore"),结果模型加载从后台静默变成了不安全模式。等到模型文件来源被污染,反序列化弹出一个__reduce__载荷时,后悔都来不及。警告虽然啰嗦,但它确实想保护你。

2. weights_only 的真正含义:从“加载字典”到“执行任意代码”

2.1 torch.load 的底层机制:PyTorch 有自己的 pickle 变种

要理解weights_only,得先知道torch.load到底在做什么。PyTorch 的持久化格式基于 pickle,但不是标准 pickle。模型保存时,张量数据会被序列化进一个 ZIP 容器,内部有一个data.pkl文件,里面记录了对象结构、类引用、张量元数据等,实际数值则单独存储。

加载时,torch.load会读取这个容器,然后通过 pickle 字节码一步步重建 Python 对象。标准 pickle 在反序列化过程中,可以引用任意 Python 类,调用任意可调用对象,甚至通过__reduce__协议指定一个函数和参数去执行。这给torch.load带来了巨大的灵活性,但也带来了一个致命问题:反序列化不可信文件等于执行了文件内嵌入的代码。

2.2 weights_only 参数的工作边界

weights_only在 PyTorch 2.x 中是一个布尔参数。为True时,反序列化器只允许加载一小部分基础类型和 PyTorch 专用类型,包括张量、基本容器(dict、list、tuple、set)以及torch.nn.Parameter等。为False时,则使用完整 pickle 加载,允许任意类实例和代码执行。

这里的理解要点:weights_only=True不是“只加载权重数字”,而是一个“白名单机制”。名单之外的一切类型都会被拒绝。PyTorch 在自己的Unpickler子类里实现了一套精确到类的白名单,比如允许torch.Tensor、torch.Size、torch.dtype这些核心类型。白名单之外的东西,直接抛UnpicklingError。

你可以把它类比成一道安检门:weights_only=True相当于只放行有登记证的固定人员,其他任何人想进去都会被拦下;weights_only=False相当于大门敞开,谁都可以进来,而且进来之后还能在楼里跑一段自带的程序。

2.3 一个容易误解的地方:state_dict 与完整模型对象

很多人以为torch.save(model.state_dict(), ...)保存的只是一个普通 dict,所以weights_only=True加载它绝对没问题。这个判断基本正确,但要注意边的细节。state_dict里的键值对通常就是字符串、torch.Tensor、torch.nn.Parameter这些,全部在白名单内。

但如果保存的是完整模型对象:

torch.save(model, "model_full.pth")

这会把模型类引用也写进去。比如你在某个模块里定义了class MyNet(nn.Module),然后用torch.save(model, ...)保存。加载时如果weights_only=True,就会报错,因为MyNet不在白名单里。此时要么显式传weights_only=False,要么改变保存策略,只保存state_dict。

3. 这个警告背后是安全问题,不是 API 在闹脾气

3.1 pickle 反序列化漏洞:从一个“演示炸弹”说起

模型文件在深度学习工作流里经常通过网盘、Hugging Face、内部文件服务器分发。很多人觉得“模型文件就是二进制大数字,没什么可担心的”,这种想法低估了 pickle 的威力。

标准的 pickle__reduce__协议允许序列化数据指定一个可调用对象,并在反序列化时执行。举个例子:

import pickle import os class Exploit: def __reduce__(self): return (os.system, ("echo pwned",)) payload = pickle.dumps(Exploit())

如果把这个 payload 写进一个后缀为.pth的文件,再用torch.load(..., weights_only=False)加载,os.system("echo pwned")就会在加载方机器上执行。这只是演示,真实攻击可以把os.system换成下载木马、反弹 Shell、删除文件等任意命令。

PyTorch 也曾收到过相关的漏洞报告,社区对torch.load默认行为进行变更的源头正是这类攻击面。FutureWarning本质上是在告诉你:以前我在默认情况下放开了这个口子,未来我会收紧它,现在请你先自查一遍。

3.2 全局变量注入攻击:比执行命令更隐蔽的一种方式

还有一种攻击方式不直接弹 Shell,而是通过操作builtins或sys.modules里已有对象的属性来注入“后门”。比如攻击者可以在 pickle 流里写入一个torch.Tensor,但同时动态修改某个全局类的forward方法,把模型推理替换成恶意逻辑。这样即使模型加载成功,精度检测也完全正常,但推理结果已经被篡改。

这种攻击不需要触发os.system之类明显的危险操作,所以更难被发现。它尤其适合投毒开源模型:用户加载了模型,训练时一切正常,部署后模型的输出在某些特定输入下会走向攻击者预设的方向。

weights_only=True由于把可加载对象限制在白名单内,能够封堵这类注入路径。即便攻击者把恶意 pickle 字节码塞进文件,反序列化框架也不会执行白名单之外的__reduce__逻辑。

3.3 为什么weights_only=False仍然被保留

既然不安全,为什么不直接删除weights_only=False这个选项?原因很现实:PyTorch 生态里存在大量合法场景,需要在模型文件中混合保存自定义类、配置对象、tokenizer 等。比如 Hugging Face Transformers 过去有不少 checkpoint 直接在文件里保存了PreTrainedTokenizer对象,这类自定义类不可能全塞进 PyTorch 的内置白名单。

所以官方的思路不是禁止weights_only=False,而是把它从默认值改成显式选择。你要用可以,但必须自己确认文件来源可信。这个态度很明确:默认安全,风险自负。

4. 什么时候必须显式传 weights_only=False:场景清单与代码示例

4.1 场景一:保存的对象里包含自定义类实例

最常见的情况是训练脚本里用torch.save保存了一个大字典,里面既有模型权重,又有优化器状态、epoch 时间戳,还有某个自定义类:

import torch class ExperimentConfig: def __init__(self, lr, batch_size): self.lr = lr self.batch_size = batch_size config = ExperimentConfig(lr=1e-3, batch_size=32) model = torch.nn.Linear(10, 5) torch.save( { "model": model.state_dict(), "optimizer": torch.optim.Adam(model.parameters(), lr=1e-3).state_dict(), "config": config, "epoch": 10, }, "checkpoint.pth", )

随后在 PyTorch 2.6 里执行:

checkpoint = torch.load("checkpoint.pth")

你会看到一个明确的报错:

UnpicklingError: Weights only load failed: Unsupported class 'ExperimentConfig'

这就是weights_only=True生效的结果。此时你必须明确告诉 PyTorch“这个文件我信任,放开限制”:

checkpoint = torch.load("checkpoint.pth", weights_only=False)

没有这个参数,代码在跑训练时连 checkpoint 都恢复不了,前功尽弃。所以对于旧项目,排查到这种场景时建议先加参数,让项目恢复运行,再考虑逐步迁移。

4.2 场景二:加载旧版 PyTorch 保存的非标准文件

PyTorch 版本之间的序列化格式基本保持兼容,但某些库(比如早期的第三方分布式框架)会在文件里保存额外的自定义对象。这些对象在保存时合法,在旧版加载时也合法,到了 PyTorch 2.6 的白名单机制下就会卡住。

遇到这种问题,错误信息里通常会带Unsupported class和对应的模块路径。我的建议是先用weights_only=False跑通,确认文件内容里到底是什么,再决定要不要彻底清理这些非必要对象。

4.3 一个值得尝试的最小复现脚本

如果你想在自己的环境里完整复现一遍这个“Warning → 报错 → 显式解除”的链路,可以用下面这个小实验,在 Jupyter 或脚本里直接跑:

import torch import warnings # 定义一个不在白名单里的自定义类 class CustomBox: def __init__(self, value): self.value = value box = CustomBox(42) # 用 torch.save 保存 torch.save({"box": box, "data": torch.ones(3)}, "custom_box.pth") # 触发 FutureWarning / 报错 with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") try: loaded = torch.load("custom_box.pth", weights_only=True) except Exception as e: print("Load failed:", type(e).__name__, str(e)) else: print(loaded) # 使用 weights_only=False 加载,正常 loaded = torch.load("custom_box.pth", weights_only=False) print(loaded["box"].value)

这个脚本可以验证两点:第一,在较新的 PyTorch 2.x 里,如果不传weights_only,会按默认值走并显示FutureWarning;第二,当文件里有自定义类时,weights_only=True会直接报错,而不是“尝试加载后忽略”。理解这些行为,后续排查才有方向。

5. 全面消除 FutureWarning 的实操方案对比

5.1 方案一:明确传入 weights_only=True

如果你的文件里只有权重、优化器状态、基础类型,那直接在加载处显式传入weights_only=True,警告立即消失:

state_dict = torch.load("model.pth", weights_only=True)

这相当于告诉 PyTorch:我已经确认这个文件只需要白名单类型,你无需为兼容旧行为而警告我。

需要注意的是,优化器状态里有些字段可能并不在白名单内,比如个别优化器的state中保存了torch.Tensor,正常是允许的。但如果是第三方优化器,字段里混入了自定义类,就得看实际情况了。所以不要机械地以为“优化器状态一定安全”,在一个项目里批量改参数之前,最好先把每个 checkpoint 类型跑一遍。

5.2 方案二:明确传入 weights_only=False

对于确实包含自定义类的 checkpoint,显式写:

checkpoint = torch.load("checkpoint.pth", weights_only=False)

这是目前最直接的“消除警告”方法。但前提是你对文件来源有绝对信心,比如文件是本地训练产出的,没有经过第三方环节。如果你的模型是从某个开源平台下载的,文件内容又不完全可知,那么weights_only=False就存在被投毒的风险。

5.3 方案三:自定义 Unpickler 白名单

官方还提供了一个更精细的接口:你可以自定义Unpickler,允许加载你确实需要的类,而不是一刀切地放开全部限制。这样既能通过白名单机制安全加载自定义对象,又不会退回到完全无防护状态。

具体思路是继承torch.serialization.Unpickler或利用pickle模块的find_class钩子。下面是一个基于 PyTorch 官方_weights_only_unpickler扩展思路的示例:

import pickle import torch from torch.serialization import add_safe_globals # 把你项目里真正需要的自定义类加入安全全局白名单 add_safe_globals([ExperimentConfig])

在 PyTorch 2.x 的torch.serialization模块里,add_safe_globals是一条便捷路径,可以把自定义类加入允许列表,然后仍然用weights_only=True加载。如果版本支持,这个方案比weights_only=False更优雅,因为它只开放了你指定的类,其余依旧严格受限。

如果当前版本没有add_safe_globals,也可以自己继承pickle.Unpickler,重写find_class。下面是一个兼容性较强的写法:

import pickle class SafeUnpickler(pickle.Unpickler): SAFE = { ("torch", "Tensor"), ("torch", "Size"), ("experiment", "ExperimentConfig"), } def find_class(self, module, name): if (module, name) not in self.SAFE: raise pickle.UnpicklingError(f"blocked {module}.{name}") return super().find_class(module, name) def load_safe(fileobj): return SafeUnpickler(fileobj).load()

这种做法的好处是“最小权限”:只有你列出的类型允许加载。坏处是维护成本偏高,每次 checkpoint 里新增一个自定义类,都要记得更新白名单。

5.4 四个方案的横向对比

方案代码改动量安全性适用场景
显式weights_only=True极小高纯权重、基础类型文件
显式weights_only=False极小低,依赖信任含自定义类的旧文件
add_safe_globals([类])小较高自定义类少且可控
自定义SafeUnpickler中等高需要精细控制加载类

对大多数内部训练脚本来说,方案一和方案二能解决 90% 的问题。方案三适合那些需要在加载端兼顾安全与自定义类的项目。

6. 更稳的长期解法:从 checkpoint 格式层入手

6.1 不要把所有东西都塞进 torch.save 的字典里

FutureWarning暴露的只是表层问题,底层问题是“torch.save + torch.load”这个组合承担了太多本不该它承担的任务。很多人图省事,把训练配置、数据预处理参数、tokenizer、自定义类实例、日志元信息统统塞进同一个.pth文件。这种“万能字典”在单机个人项目里确实方便,但一旦要分发、部署、跨版本加载,麻烦就来了。

更好的做法是拆开存储:

  • 权重保存为state_dict或直接用safetensors格式。
  • 配置文件保存为 JSON 或 YAML。
  • tokenizer、词表等文本类数据保存为独立文件。
  • 只有真正需要组合速度的巨型 checkpoint 才考虑封装。

拆开以后,模型加载路径可以保持weights_only=True或直接走safetensors,完全避开了 pickle 风险。

6.2 safetensors 是 torch.load 之外的重要替代

huggingface/safetensors是一个专门为张量存储设计的格式,它不基于 pickle,序列化和反序列化都不涉及任意代码执行。它的 API 也很简单:

from safetensors.torch import save_file, load_file save_file(model.state_dict(), "model.safetensors") tensors = load_file("model.safetensors")

load_file返回一个普通的 dict,可以直接加载进模型:

model.load_state_dict(tensors)

使用 safetensors 最大的优势就是“不需要担心反序列化执行代码”,因为格式本身决定了它只能存储张量数据。对于开源模型分发、跨框架部署,这是目前生态里更受推荐的方向。

但要注意,safetensors 并不支持任意自定义对象。如果你确实需要保存包含优化器状态、自定义类实例的完整训练 checkpoint,safetensors 就不能完全替代torch.load。折中方案是:训练过程中把“纯张量类”数据存成 safetensors,把“非张量元信息”存成 JSON,两者分开管理。

6.3 给旧项目的迁移路线建议

如果手头有一堆旧项目,全部推倒重构不现实。我建议按下面的优先级逐步迁移:

  1. 先给所有torch.load调用补上显式参数,用weights_only=True或weights_only=False,把FutureWarning清零。这是成本最低的操作。
  2. 检查哪些文件用weights_only=True能正常加载。这些文件说明内容全是白名单类型,后续可以优先转 safetensors。
  3. 那些必须用weights_only=False才能加载的文件,排查里面的自定义类到底是什么。能用add_safe_globals白名单替代的就用白名单,减少全量开放的窗口。
  4. 新项目一开始就设计好 checkpoint 结构,默认用 safetensors 存放张量,用 JSON 存放配置。

这套流程不需要一次性完成。可以先保证代码在 PyTorch 2.6+ 下跑起来,再慢慢优化。

6.4 团队协作中的一个常见坑:警告被吞掉

如果你的项目里有人图省事在conftest.py或启动了全局忽略警告:

warnings.filterwarnings("ignore", category=FutureWarning)

那么FutureWarning会消失,但如果未来 PyTorch 把默认值翻转成weights_only=True,旧代码就会在没有任何预兆的情况下直接报UnpicklingError。这比看到警告要难排查得多,因为报错点往往不在torch.load,而在随后的model.load_state_dict()或者某个属性访问处。建议大家把警告当信号,而不是噪音。如果团队里确实需要压制警告,也建议保留日志记录,让行为变化有迹可循。

根据我的实际经验,在升级 PyTorch 2.x 之后,第一批需要处理的就是所有直接出现在模型推理和训练代码里的torch.load调用。批量排查时可以先用一个简单脚本把所有加载入口的调用方式打出来:

grep -rn "torch.load" --include="*.py" .

然后逐个检查:加载对象是否来自可信来源、文件内容是否包含自定义类、能否切换到 safetensors。这样一遍梳理下来,不仅是消警告,也是对整个项目的模型持久化方案做了一次体检。经过这一轮调整,后续再碰到torch.load相关的告警或报错,心里会非常清楚问题出在哪里,不会再被模糊的提示绕晕。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询