在 torch 2.6 里 torch.load 加载自己保存的 checkpoint,直接抛 UnpicklingError,原文 Weights only load failed,同一文件在 2.3 上能读,服务启动就崩。

一开始怀疑是「旧文件是用 pickle 存的整个训练状态」,查了一圈,方向不对。

排查方向其实有限,先把可能性排一排。2.6 把 weights_only 的默认值改成了 True。也可能是checkpoint 里存了 argparse.Namespace 等自定义对象。

另外,torchvision 的旧模型定义不在安全白名单里。

第一件事是确认 torch 版本与默认值

python -c "import torch;print(torch.__version__, torch.load.__defaults__)"

显式传 weights_only 读一次做对照

python -c "import torch;sd=torch.load('ckpt.pt',map_location='cpu',weights_only=False);print(type(sd))"

接着把模型权重单独存成 state_dict

python -c "import torch;sd=torch.load('ckpt.pt',map_location='cpu',weights_only=False);torch.save(sd['model'],'model_only.pt')"

末了在新格式上验证可以安全加载

python -c "import torch;sd=torch.load('model_only.pt',map_location='cpu');print(len(sd))"

当时还开着另一个排查窗口,两边日志对照着看才定位到。

同一份数据我导出了两份,一份留底,一份拿来试。

恢复的判定其实很直接。权重文件能被默认参数加载,服务启动不再抛 UnpicklingError,推理输出与旧版本一致。

默认值收紧是安全改进,短期可以在加载处显式关闭,长期还是把权重和训练状态分开存。拖得越久,排查成本越高,因为现场状态会被后续操作覆盖。

作者 9823189