PyTorch模型保存:ckpt与pth差异、断点续训与避坑
深度模型的保存表面上看就是一行 torch.save 的事但真正在生产里跑过一轮的人都知道这里面藏着能让你加班到凌晨的坑。我自己就遇到过接手同事的项目看到目录里躺着 model.ckpt、best.pth、checkpoint-epoch12.pth 三个文件一时搞不清楚哪个是权重、哪个是断点、哪个能直接拿去推理最后靠比对文件大小和用 torch.load 一个个打印 key 才把关系理清楚。这篇文章就围绕 ckpt 和 pth 这两个高频后缀把深度模型保存这件事从文件长什么样讲到工程上该怎么定策略中间会把 state_dict、断点续训、map_location、参数不匹配排查这些实操细节全部铺开。不管你是刚跑完第一个 MNIST 的新手还是已经在带团队做模型交付的老手都能在这里找到能直接抄走的代码模板和避坑经验。1. 先弄明白ckpt 和 pth 到底是不是两种格式1.1 后缀只是习惯真正决定内容的是你存进去的对象很多人第一反应是ckpt 和 pth 是两种不同的文件格式这个理解在 90% 的情况下是错的。PyTorch 的torch.save底层走的是 pickle 序列化加 zip 打包文件名后缀写什么它根本不管你写成.model、.bin、.weights都能正常读回来。官方文档的示例里习惯用.pt和.pth而 TensorFlow 系的 checkpoint 机制习惯用.ckptPyTorch Lightning 为了和 TF 的习惯对齐也把自己的断点文件默认命名为.ckpt。于是就有了看到 ckpt 以为是 TF看到 pth 以为是 PyTorch这种经验主义判断实际上并不可靠。真正决定一个文件里装了什么是你传给torch.save的那个对象。它可能是整个nn.Module实例可能是一个state_dict字典也就是参数名到张量的映射也可能是一个包含了模型权重、优化器状态、当前 epoch、学习率调度器状态的复合字典。这三种东西虽然都能叫.pth但加载方式和适用场景完全不同。所以我判断一个权重文件从来不看后缀而是先跑一段探测代码把它里面的 key 打出来。import torch ckpt torch.load(model.pth, map_locationcpu, weights_onlyTrue) print(type(ckpt)) if isinstance(ckpt, dict): print(list(ckpt.keys())[:20])如果打印出来是一堆layer1.0.conv1.weight这种名字那它就是纯 state_dict如果打印出来是epoch、optimizer、lr_scheduler、model这种顶层 key那就是断点文件如果打印出来直接是一个 Module 对象那就是整包保存。三种情况对应的加载写法完全不一样这也是后面所有坑的源头。1.2 state_dict 和整个模型对象的本质区别理解 ckpt 和 pth 的差异绕不开 state_dict 这个概念。nn.Module的state_dict()返回的是一个有序字典里面只有可学习的参数和注册过的 buffer比如 BatchNorm 的 running_mean、running_var不包含网络结构本身。这就意味着光有 state_dict 你没有模型定义代码是跑不起来的因为 PyTorch 不知道这些张量该往哪个层里塞。而torch.save(model, path)走的是另一条路它用 pickle 把整个 Module 对象连同它的类定义引用一起序列化。听起来很方便问题是 pickle 保存的是类的引用路径比如mymodels.resnet.CustomResNet。你把文件拷到另一台机器如果那个模块路径不存在、或者类定义改了、或者文件夹结构变了加载直接报ModuleNotFoundError或者AttributeError。这种耦合在单人实验环境里还好一旦进入多人协作或者模型交付环节就是灾难。我在团队里推的规则很明确只保存 state_dict永远不保存整个模型对象。理由有三条。第一可移植换机器换目录都不影响第二文件小pickle 整个对象会把一些冗余的 Python 属性也带上第三安全反序列化一个完整的类对象比反序列化一个纯张量字典的风险高得多。代价是你必须维护模型定义代码但这个代价在工程上完全可以接受因为模型结构本来就应该进版本管理。1.3 不同生态下 ckpt 的含义差异虽然 PyTorch 系也能随便叫 ckpt但在 TensorFlow/Keras 的世界里ckpt 是一个有明确协议的东西。Keras 的 ModelCheckpoint 回调默认生成的是一组文件.ckpt-5.index、.ckpt-5.data-00000-of-00001还可能带一个checkpoint文本文件记录最新的是哪一步。这一组文件必须放在一起缺一个都读不了。看到这种多文件结构基本可以确定是 TensorFlow 系的产物和 PyTorch 的单文件 ckpt 不是一回事。还有一种情况是 Hugging Face Transformers 保存出来的权重通常是pytorch_model.bin或model.safetensors配一份config.json。有些人为了统一命名手动把它改名成.ckpt这就更让人迷惑了。所以养成习惯拿到陌生权重文件先看它是单文件还是多文件组再看它有没有配套的 config最后用探测代码确认内容结构三步走下来就不会认错。来源生态典型文件名是否单文件内部结构PyTorch 手写训练脚本model.pth / best.ckpt单文件通常为 state_dict 或复合字典PyTorch Lightningepoch12-step900.ckpt单文件复合字典含 hyper_parametersTensorFlow / Kerasmodel.ckpt-5.index .data多文件组图变量需配套 metaHugging Facepytorch_model.bin / .safetensors单文件加 configstate_dictkey 命名有前缀TorchScript 导出traced_model.pt单文件可执行图非字典这张表建议收藏遇到陌生文件先对号入座能省掉大量试错时间。2. PyTorch 三种保存姿势的取舍逻辑2.1 整包保存方便是真方便坑也是真坑torch.save(model, full_model.pth)这种写法在教程里出现频率极高因为它加载时只要一行model torch.load(full_model.pth)就完事不需要你在加载侧再写一遍网络定义。对于做快速验证、写 demo、跑课设作业的场景它确实省事。但它的三个硬伤决定了它上不了生产。第一是环境耦合我踩过一次特别典型的坑实验机上用 PyTorch 1.12 保存的整包模型换到只有 1.8 的服务器上加载报了一长串AttributeError原因是某些层在序列化时记录了旧版本的内部属性。第二是代码耦合模型类只要改了构造函数签名、加了一个新参数旧文件就可能加载失败因为 pickle 在重建对象时会重新调用__init__。第三是 pickle 反序列化的安全风险pickle 文件在加载时是可以执行任意代码的来源不明的权重文件绝对不能直接 load这一点后面第 4 章会展开讲。还有一种情况是保存了整个对象但只想取权重这时候你会写torch.save(model.state_dict(), ...)注意这已经切换成第二种姿势了。我见过不少人在同一个项目里两种混用结果加载时一半报错一半正常排查起来非常痛苦。所以定一条死规矩一个项目只用一种保存姿势团队内统一。2.2 state_dict 才是工程上的标准姿势标准做法是把模型和权重分开对待结构定义写在代码里权重单独存文件。保存端就一行torch.save(model.state_dict(), best.pth)加载端需要两步先实例化同结构的模型再灌权重model build_model(num_classes10) # 结构必须和训练时完全一致 state torch.load(best.pth, map_locationcpu, weights_onlyTrue) model.load_state_dict(state) model.eval()这里load_state_dict默认是严格模式strictTrue要求两边的 key 集合完全一致多一个少一个都会抛错。这个默认值其实是好事它帮你把结构不匹配的问题在第一时间暴露出来而不是悄悄加载了一半参数然后推理结果莫名其妙。如果你确实需要放宽比如只加载骨干网络、分类头重新初始化那就显式写strictFalse然后务必打印缺失和多余的 key 列表确认一遍别糊里糊涂就往下跑。map_location这个参数是我认为最被低估的一个。它的作用是告诉 PyTorch 把张量映射到哪个设备上。GPU 上保存的权重直接在没有 GPU 的机器上加载不加map_location会直接报 CUDA 不可用加了map_locationcpu就能正常读进来。常见的还有map_locationcuda:1指定卡号或者map_location{cuda:0: cuda:1}做设备重映射。多卡场景下这个参数用得特别多后面细说。2.3 断点续训要保存的东西远不止权重如果你只保存 state_dict训练中断之后重新开始优化器的动量、自适应学习率的历史累积、学习率调度器走到第几步、AMP 的梯度缩放因子这些全部丢失。表面上看模型能继续训实际上收敛轨迹已经变了尤其是在训练后期优化器状态的重要性不比权重低。所以断点文件的正确打开方式是存一个复合字典。我常用的模板长这样def save_checkpoint(path, model, optimizer, scheduler, scaler, epoch, best_metric, cfg): torch.save({ epoch: epoch, model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict() if scheduler else None, scaler: scaler.state_dict() if scaler else None, best_metric: best_metric, config: cfg, }, path)注意这里用了model.state_dict()而不是model外层字典的 key 我叫model所以完整的结构是字典里有 model、optimizer 等子字典。这种嵌套结构在加载时容易写错很多人会写成state_dict[state_dict]或者state_dict[model_state_dict]其实取决于你保存时怎么命名的。这也是为什么我在团队里要求保存的 key 名固定为model、optimizer、scheduler、epoch、best_metric五个不允许自由发挥谁改谁负责同步所有下游加载代码。顺带说一句config这个字段。把训练配置学习率、batch size、数据增强参数、类别数一起存进 ckpt好处是一年后你翻出这个文件还能复现实验坏处是如果配置里塞了不可序列化的对象比如数据集实例、lambda 函数就会保存失败。我的做法是只存基础类型和列表字典复杂对象转成字符串描述。3. 一套可以直接复现的保存与加载实操流程3.1 保存端区分 best 和 last 两个文件训练脚本里我通常维护两个文件last.pth每个 epoch 覆盖写best.pth只在指标刷新时写。这样做的原因很实际断点续训要用 last模型交付要用 best两个用途分开互不干扰。如果只存一个文件你会在继续训练和保留最好模型之间反复纠结。目录结构建议这样组织runs/exp_20240612/ ├── config.yaml ├── last.pth ├── best.pth └── log.txt每个实验一个时间戳目录两个权重文件加一份配置。这个结构看起来朴素但它解决了一个大问题三个月后你回来看结果不用去猜best.pth是哪次实验、什么参数。我见过太多人把十几个实验的 best.pth 全堆在一个目录里最后靠文件修改时间排序那画面太心酸了。保存 best 的判断逻辑本身也有讲究。分类任务常用准确率或 F1我一般用验证集 F1因为它在类别不均衡时比准确率更能反映真实水平。检测类任务可能用 mAP分割类用 mIoU。注意阈值方向要写对if metric best_metric和if metric best_metric差了十万八千里我在早期代码里犯过一次反号错误结果保存下来的 best.pth 是整个训练过程中最差的那个白白浪费了一晚上算力。3.2 加载端三种场景的写法对照加载场景可以归成三类写法各有侧重。纯推理只关心权重其他一律不管model build_model(num_classescfg[num_classes]) state torch.load(best.pth, map_locationcpu, weights_onlyTrue) model.load_state_dict(state[model] if model in state else state) model.eval() with torch.no_grad(): out model(x)注意这里我加了一个model in state的判断用来兼容纯 state_dict和复合字典两种情况。这种容错写法在接手别人项目时特别有用能少写两行调试代码。生产环境我还是建议解析清楚再写死别把容错逻辑留在线上。继续训练需要完整恢复优化器状态ckpt torch.load(last.pth, map_locationcpu, weights_onlyFalse) model.load_state_dict(ckpt[model]) optimizer.load_state_dict(ckpt[optimizer]) scheduler.load_state_dict(ckpt[scheduler]) start_epoch ckpt[epoch] 1这里weights_onlyFalse是因为优化器状态里包含了 param_groups 这类结构纯张量模式读不了。这一点非常关键PyTorch 2.6 起torch.load的weights_only默认值变成了 True如果你还在用旧代码加载断点文件升级后会突然报错提示反序列化被拒绝。解决办法就是显式传weights_onlyFalse同时确认文件来源可信。迁移学习加载骨干、换掉分类头state torch.load(pretrained.pth, map_locationcpu, weights_onlyTrue) missing, unexpected model.load_state_dict(state, strictFalse) print(missing:, missing) print(unexpected:, unexpected)strictFalse会返回两个列表missing 是你模型里有但权重文件里没有的 keyunexpected 反过来。迁移学习场景下你期望看到的 missing 应该正好是新的分类头unexpected 应该正好是旧的分类头。如果 missing 里出现了骨干网络的层名说明结构对不上得回头检查。3.3 断点续训的完整实现与恢复点选择把前面几块拼起来一个能用的训练循环大概是这个样子start_epoch 0 best_metric 0.0 resume_path runs/exp/last.pth if os.path.exists(resume_path) and args.resume: ckpt torch.load(resume_path, map_locationcpu, weights_onlyFalse) model.load_state_dict(ckpt[model]) optimizer.load_state_dict(ckpt[optimizer]) if ckpt.get(scheduler): scheduler.load_state_dict(ckpt[scheduler]) if ckpt.get(scaler): scaler.load_state_dict(ckpt[scaler]) start_epoch ckpt[epoch] 1 best_metric ckpt.get(best_metric, 0.0) print(fresumed from epoch {start_epoch}) for epoch in range(start_epoch, total_epochs): train_one_epoch(...) metric evaluate(...) save_checkpoint(runs/exp/last.pth, model, optimizer, scheduler, scaler, epoch, max(best_metric, metric), cfg) if metric best_metric: best_metric metric torch.save({model: model.state_dict(), epoch: epoch, metric: metric}, runs/exp/best.pth)这段代码里有几个细节值得强调。第一start_epoch ckpt[epoch] 1的加一是必须的我见过漏掉加一的写法导致恢复后同一个 epoch 训了两遍数据采样和调度器步数都会错位。第二best_metric也要跟着恢复否则恢复后第一次验证不管多差都会被判定为新最优直接把真正的好模型覆盖掉。第三如果用了自动混合精度scaler的状态必须一起存因为它的缩放因子是动态累积的重置后前几十步的梯度会不稳。还有一个容易忽略的点是数据加载器的随机种子。断点续训的理想状态是从断点处严格继续但 DataLoader 的 shuffle 顺序、各种数据增强的随机数如果不做种子管理恢复后的数据流和中断前是不一样的。严格复现需要保存随机数生成器状态通常在研究场景才这么干工程上大家接受少量偏差。但如果你在做对复现性要求很高的对比实验那就得把torch.get_rng_state()、numpy.random.get_state()也塞进 ckpt。3.4 从 pth 到 pt、TorchScript、ONNX 的转换思路经常有人问pth 是不是从 pt 导出的这个问题其实建立在pt 和 pth 是两种格式的误会上。在 PyTorch 的语境里.pt和.pth从来都是同一种东西的两种叫法官方示例里两者混用不存在导出关系。真正存在转换关系的是另外几组state_dict 到 TorchScript、state_dict 到 ONNX。TorchScript 导出是为了脱离 Python 环境部署生成的文件通常叫traced.pt它是一个可以独立执行的图不依赖你的模型定义代码model.eval() example torch.randn(1, 3, 224, 224) traced torch.jit.trace(model, example) traced.save(traced.pt)这里model.eval()是必须的因为 trace 会把当时的计算路径固定下来如果模型在训练模式下Dropout 和 BatchNorm 的行为会被错误地烘焙进图里推理结果全乱。这个坑非常隐蔽模型不会报错只是结果不对。ONNX 导出用于跨框架部署torch.onnx.export( model, example, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version12, )dynamic_axes这行很关键不写的话导出的模型 batch 维度是写死的部署时只能一张一张推理吞吐量直接砍到十分之一。opset_version的选择取决于你的推理引擎支持到哪个版本选高了跑不起来选低了某些算子不支持我一般从 11 或 12 起步试。4. 踩坑实录ckpt 与 pth 相关的常见问题排查4.1 加载报错速查表下面这张表是我这些年攒下来的报错对照基本覆盖了九成以上的加载问题。报错信息关键词大概率原因处理方式Missing key(s) in state_dict结构不一致或键名前缀不同打印两边 key 对比检查是否多了module.前缀Unexpected key(s) in state_dict权重文件比模型多参数确认是否加载了错误文件或用 strictFalsesize mismatch for ...某层维度对不上常见于类别数改动检查输出层维度必要时裁掉重训Attempting to deserialize object on a CUDA device无 GPU 环境加载 GPU 权重加 map_locationcpuWeights only load failed新版默认 weights_onlyTrue断点文件显式传 weights_onlyFalseModuleNotFoundError整包保存且类路径失效改用 state_dict 保存重训或重建结构RuntimeError: Error(s) in loading state_dict复合字典没取对层级打印顶层 keys确认取 ckpt[model]这张表我用得很频繁基本看到报错第一行就能定位方向。4.2 key 不匹配的定位方法打印对比是唯一正解面对 size mismatch 或者 missing key最有效的手段就是把两边的 key 都打出来做对比别靠猜。model_keys set(model.state_dict().keys()) ckpt_keys set(ckpt[model].keys()) print(only in model:, sorted(model_keys - ckpt_keys)[:10]) print(only in ckpt:, sorted(ckpt_keys - model_keys)[:10])如果发现 ckpt 里的 key 全都带module.前缀说明权重是用nn.DataParallel或DistributedDataParallel训练时保存的包了一层。这时候有两个办法一是加载时手动去前缀new_state {k.replace(module., , 1): v for k, v in ckpt[model].items()} model.load_state_dict(new_state)二是用load_state_dict的strictFalse配合key重映射。我更推荐第一种因为改完的字典可以继续当普通 state_dict 用逻辑干净。反过来如果模型本身被 DDP 包了而权重没前缀加载前也要相应处理。这里有个经验保存的时候统一存model.module.state_dict()也就是剥掉 DDP 外壳这样存出来的文件在任何场景下都能用不必每次加载都做前缀清洗。这个习惯我从第一次遇到 DDP 权重问题之后就一直保持。4.3 文件体积、安全与版本兼容模型文件的安全问题值得单独说。pickle 格式在反序列化时可以触发任意代码执行这已经不是理论风险了。所以第一条纪律是只加载你自己训练或者可信来源的权重文件。网上随手下载的权重文件先看来源能用 safetensors 格式的就用 safetensors它只存张量不支持代码执行天然安全。PyTorch 也在这方面做了收紧weights_onlyTrue就是只允许加载张量、基础类型和少量安全对象。从 2.6 开始这个参数默认变成 True是好事但会打破一批老代码。我的处理方式是推理加载用默认的 True断点续训显式写 False并且把这条规则写进项目的 README。体积方面一个 50M 参数的模型float32 权重差不多 200MB如果断点文件里还带了优化器状态Adam 会存两份动量体积直接变成三倍。所以断点文件动辄几个 G 是正常的磁盘规划要提前算好。我一般会保留最近三个 last.pth 加全部 best.pth老的用脚本自动清理不然跑一个月的实验能把盘塞满。4.4 多卡、EMA 与混合精度下的特殊处理多卡训练时的保存策略有两个主流选择。一是只在 rank 0 上保存配合torch.distributed.barrier()保证其他进程不写文件这样避免多个进程同时写同一个文件导致损坏。二是每个 rank 存自己的分片用于后续并行加载大模型这是近年来大模型训练的常见做法但需要配套的加载逻辑。if dist.get_rank() 0: torch.save({model: model.module.state_dict()}, ckpt.pth) dist.barrier()barrier()不能省否则 rank 0 还在写文件的时候其他进程可能已经进下一轮训练损毁文件的风险是真实存在的。EMA指数移动平均权重的情况也常见。做深度伪造检测模型时比如基于 Xception 骨干的检测网络很多人会同时维护原始权重和 EMA 权重两套参数。保存时要把两套都存下来因为验证阶段通常用 EMA 权重评估而继续训练要用原始权重更新梯度。命名上建议用model和ema_model两个 key 明确区分别都叫state_dict。混合精度的坑前面提过一次这里再补一个scaler的状态如果没保存续训后缩放因子会从默认值重新开始。在前几千步里梯度缩放可能偏小导致下溢loss 会突然抖动。判断方法很简单续训后前 100 步的 loss 曲线如果出现明显跳变基本就是这个原因。5. 真实项目里怎么定保存策略5.1 从三个实际模型类型的保存需求说起不同模型对保存策略的需求差别挺大我用三个具体场景来说明。第一个是深度伪造检测模型典型结构是 Xception 骨干加一个二分类头。这类任务的特点是数据不平衡严重指标波动大所以 best 的判定不能只看单次验证结果我一般会在验证集上跑多次或用滑动平均。另外这类模型经常要做跨数据集测试权重文件的可移植性要求高所以必须是纯 state_dict 加一份 config坚决不能整包保存。第二个是深度平衡模型Deep Equilibrium Model它的特点是前向过程是求解一个不动点本身不带很多中间层参数权重的存储结构反而比较轻。但它有个特殊点训练时需要保存求解器的迭代次数和收敛阈值等状态否则续训时数值行为会变。所以这类模型的 ckpt 里除了常规字段还要额外记录求解器配置。第三个是深度循环模型比如各种 RNN、状态空间模型。它们的隐藏状态在网络内部流转一般不作为参数保存但如果你的实现里把某些状态做成了 buffer比如某些变体中的初始状态那就会进 state_dict加载时必须保证 buffer 也对得上。我遇到过 buffer 维度不一致导致的 size mismatch排查了半天才发现是序列长度配置变了。5.2 命名规范与版本管理的一些约定混乱的文件名是排查成本的最大来源。我现在的命名规范是这样的{实验名}_{数据集}_{骨干网}_{指标值}_{epoch}.pth比如dfdetect_ffpp_xception_f1-0.923_ep18.pth。指标值放文件名里好处是一眼能看出好坏不用加载。缺点是每次刷新 best 都要重命名所以实践上 best.pth 用固定名同时软链接指向带指标名的文件兼顾便利和可读性。版本管理方面权重文件不要进 git用 git-lfs 也尽量避免因为大文件会拖慢仓库。我的做法是权重存独立的存储路径git 里只保留一个weights_manifest.json记录每个实验的权重路径、指标、训练命令。这个小文件几十行但它让整个项目可追溯。顺带说一句保存本地模型配置失败这类问题。它通常不是保存逻辑本身的错而是配置里塞了不可序列化的东西比如数据集对象、文件句柄、lambda 函数。写配置的时候坚持只用基础类型和容器类型这条原则基本就不会遇到。如果确实需要存复杂对象转成字符串描述或者只存它的构造参数。5.3 保存频率与性能开销的权衡每个 epoch 都保存断点文件的代价经常被低估。一个 3GB 的断点文件写盘在普通机械盘上要好几秒如果训练一个 epoch 只要 20 秒那保存开销就占了 20% 以上而且写盘 IO 会阻塞训练进程。优化方式有几个一是只在满足条件时保存比如每 N 个 epoch 或指标刷新时二是先写临时文件再原子重命名避免写一半崩溃导致文件损坏三是用后台线程异步写盘训练继续跑。tmp path .tmp torch.save(obj, tmp) os.replace(tmp, path) # 原子操作避免半截文件os.replace这一步看着不起眼但它能防止断电或者进程被杀导致 last.pth 变成损坏文件。我在一次集群任务被抢占之后就加上了这个习惯代价几乎为零收益是文件永远可用。6. 一些容易被问到的细节问题6.1 关于 weights_only 的取舍前面反复提到这个参数值得单独梳理一下判断标准。纯推理加载 state_dict用weights_onlyTrue安全且够用。加载含优化器状态的断点用weights_onlyFalse因为优化器状态字典里有非张量结构。加载整包模型对象必然要 False但这条路径我建议直接放弃。还有一个细节如果你用的是较老版本 PyTorchweights_only参数可能根本不存在那说明版本在 2.0 之前加载行为一直是宽松模式。升级版本时要留意这个行为变化最好在升级前先把项目里的torch.load调用点全部列出来逐个确认该传什么。6.2 CPU 与 GPU 权重互转的实际影响从 GPU 保存的权重加载到 CPU除了要map_locationcpu还有一点要注意张量的 dtype 和布局不变所以内存占用和 GPU 显存占用是一样的量级。一个大模型加载到 CPU 内存可能直接把内存吃满这时候可以考虑用torch.load(..., mmapTrue)做内存映射按需读取能显著降低峰值内存。反向的 CPU 到 GPU 加载不需要特殊处理load 完之后调用model.to(cuda)就行。但要注意map_locationcuda和先 load 到 CPU 再 to cuda 的区别前者在加载过程中就分配显存如果显存不够会直接失败且可能留下碎片后者更可控。我一般倾向于先加载到 CPU 再搬出问题好排查。6.3 权重文件损坏的识别与预防权重文件损坏在长时间训练里不算罕见尤其是写到一半进程被杀、或者是网络存储抖动。识别方法很直接加载时如果报UnpicklingError或者zipfile.BadZipFile基本就是文件坏了。预防手段就是前面说的原子写加上临时文件。如果只是损坏了尾部有时候还能抢救用torch.load加mmapTrue可能读到部分内容但这个方法不保证成功只适合应急。真正靠谱的还是保留多个历史断点别只留一个文件。我的习惯是 last 文件保留最近三个轮转这样即使最新一个坏了损失也就一个 epoch。6.4 跨框架权重迁移的现实难度有时候需要把一个 PyTorch 的 pth 迁到别的框架或者反过来。这件事的难度取决于层的对应关系。卷积、全连接、BatchNorm 这类标准层基本能一一对应但涉及自定义算子、特殊的 padding 方式、不同的默认初始化就会出现权重对得上但结果对不上的情况。我的建议是做数值对齐验证构造一个固定输入在两边分别跑前向逐层对比输出。哪一层开始出现明显偏差问题就在那一层附近。这个流程我在做跨框架部署时走过好几次虽然麻烦但比盲目试错高效得多。最后分享一个我自己用下来最省事的习惯。不管项目大小我都会在训练脚本旁边放一个inspect_ckpt.py功能就一个传进任意权重文件打印它的顶层结构、每个子字典的 key 数量和前几个 key 名、以及所有张量的 dtype 和 shape 概览。这个小脚本我用了好几年它不能解决任何问题但它能在三十秒内让你搞清楚手上这个文件到底是什么。踩过太多次以为它是纯权重、结果是复合字典的坑之后我现在的原则是任何陌生的 ckpt 或 pth先 inspect再动手。多花这半分钟往往能省下半夜的排查时间。
上一篇/下一篇内容由系统自动关联
返回资讯列表 →