在Python中如何利用PyTorch Safetensors格式高效存储权重？-Python教程-PHP中文网

不能直接用torch.save存safetensors，因其默认使用pickle协议，而safetensors是零序列化、内存映射友好的二进制格式，不兼容pickle；正确方式是用safetensors.torch.save_file保存cpu上的state_dict，并确保张量连续、键为合法字符串。

在python中如何利用pytorch safetensors格式高效存储权重？

为什么不能直接用 `torch.save` 存 Safetensors？

因为 torch.save 默认用 Python 的 pickle 协议，而 Safetensors 是零序列化、内存映射友好的二进制格式，不兼容 pickle。直接保存会得到一个普通 .pt 文件，不是 .safetensors，也无法被 Hugging Face 生态（如 transformers、diffusers）安全加载。

如何用 `safetensors` 库正确导出模型权重？

需先将模型状态字典转为纯 numpy 或 torch.Tensor（保持 CPU 内存），再传给 safetensors.torch.save_file。关键点：

state_dict() 必须在 CPU 上：GPU 张量无法直接写入 Safetensors（会报 NotImplementedError: Cannot store tensors with device cuda）
键名必须是字符串，且不能含非法字符（如 /、\0）；常见做法是保留原始键名，但确保无嵌套结构（Safetensors 不支持嵌套 dict）
推荐显式指定 metadata（如 {"format": "pt"}），方便下游识别来源

示例：

from safetensors.torch import save_file
import torch
<p>state = model.cpu().state_dict()</p><h1>确保所有 tensor 是 contiguous 且在 CPU 上</h1><p>state = {k: v.contiguous() for k, v in state.items()}
save_file(state, "model.safetensors", metadata={"format": "pt"})
</p>

加载时为何出现 `KeyError` 或权重形状不匹配？

常见于模型结构与保存时的 state_dict 键不一致，比如：

python全能编程助手

SkillSub Pro - Python 题解与代码注释双功能技能功能概述SkillSub Pro - Python 题解与代码注释双功能技能是一项面向实际任务的技能，主要用于SkillSub Pro 是一个 Python 题解生成与代码注释的双功能合体技能，专为学生、算法学习者和开发者设计；✅ 一个技能，两种用途：；核心要点📝 题解模式：输入题目/题号，自动生成完整 Python 题解（含详细注释、解题思路、复杂度分析）；💬 注释模式：输入 Python 代码，自动添加详细中。它将相关步骤、

下载

模型用了 nn.DataParallel 或 FSDP，保存时带 module. 前缀，但加载时模型没包装——需统一用 state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()} 清洗
使用了 torch.compile 后的模型，state_dict 键可能含编译器插入的临时前缀（如 _orig_mod.），需提前剥离
加载时未调用 model.load_state_dict(..., strict=False)，遇到新增/缺失层直接报错

安全加载建议：

from safetensors.torch import load_file
<p>state = load_file("model.safetensors")</p><h1>检查 key 是否对齐</h1><p>missing, unexpected = model.load_state_dict(state, strict=False)
if missing: print("Missing:", missing)
if unexpected: print("Unexpected:", unexpected)
</p>

和 `torch.save` 相比，Safetensors 在哪些场景真正提效？

不是所有情况都更快——它优势集中在 I/O 密集、多进程或安全敏感场景：

冷启动加载快：Safetensors 支持 mmap，load_file 默认只读 header，按需加载 tensor，比 torch.load(..., map_location="cpu") 全量解压快 2–5×（尤其 >1GB 模型）
多进程安全：没有 pickle 反序列化，不会触发任意代码执行，适合不可信权重源（如社区模型 Hub）
跨框架互通：同一 .safetensors 文件可被 Rust（tokenizers）、JavaScript（@xenova/transformers）直接读取，无需转换
但训练中频繁保存（如每步 checkpoint）反而更慢：Safetensors 是单次写入格式，不支持追加或增量更新

真正省时间的地方，是部署时首次加载大模型、或 CI 中校验权重完整性——而不是训练循环里反复 dump。