在类似Hugging Face的AI模型中心浏览时,我们开始越来越多地看到带有扩展名.safetensors的文件,而不是过去常见的.bin或.pth文件。
在这篇文章中,我们将详细探讨.safetensors究竟是什么,为什么会出现以及相较于传统方法的强大优势。
1. .safetensors是什么?
.safetensors是Hugging Face开发的一种新的张量(Tensor)存储格式。
深度学习模型由数十亿个参数(权重)组成,而这些庞大的数字集需要存储和加载。为了克服基于Python的pickle模块作为事实标准所面临的重大缺陷(安全性和速度问题),这种格式应运而生。
简单来说,“更安全、更快速的模型存储文件”。
2. 为什么出现:传统方式(Pickle)的问题
现有的PyTorch模型(.bin, .pth)内部使用Python的pickle模块来序列化数据。但pickle存在致命问题。
安全漏洞(任意代码执行)
pickle不只是保存数据,它实际上是保存Python对象本身。在这个过程中,可以包含Python代码,恶意黑客可能会在模型文件内部植入破坏系统或窃取个人信息的代码。一旦用户无意中执行load(),该恶意代码就会被执行。
问题情景示例:
用户加载了从互联网上下载的model.bin -> 内部隐藏的代码被执行 -> 用户的SSH密钥或密码被发送到黑客服务器。
为了解决这些安全威胁,.safetensors应运而生。
3. .safetensors的核心特点
3.1. 安全性
.safetensors正如其名,具有安全的特点。该格式仅保存纯粹的张量数据和元数据(JSON格式)。因为没有可执行代码的空间,所以即使是从不可信来源下载的文件也可以放心加载。
3.2. 零拷贝(Zero-Copy)及速度
在加载大型LLM(大型语言模型)或Stable Diffusion模型时速度极快。
-
传统方式:文件复制到CPU内存 -> 反序列化(Unpickling) -> 再转换为张量形式 -> 移动到GPU。(产生不必要的复制过程)
-
safetensors:利用内存映射(Memory Mapping, mmap)技术。操作系统将文件直接映射到内存地址,因此可以毫不费力地从磁盘即时使用数据。这被称为零拷贝。
3.3. 懒加载(Lazy Loading)
它允许在不将整个模型加载到内存中的情况下,仅选择必要的部分进行快速读取。
例如,如果想要查看100GB模型文件中特定层的权重,传统方式需要读取全部100GB,而.safetensors则可以精确读取需要的部分。这在分布式学习环境或推理优化时非常有利。
3.4. 框架兼容性
它不依赖于特定的深度学习框架(如PyTorch等)。
-
PyTorch
-
TensorFlow
-
JAX
-
PaddlePaddle
设计使得它可以在以上各种框架中轻松读写。
4. 文件结构
.safetensors文件具有非常简单的结构。
-
头部(Header):位于文件的前部,遵循JSON格式。存储了每个张量的名称、数据类型(dtype)、形状(shape)及数据存储位置(offset)等信息。
-
数据(Data):紧随头部的二进制数据块。包含了纯粹的张量值。
由于这种结构,甚至可以仅通过读取头部来了解模型的结构,而无需读取整个文件。
4-1. 实用技巧:在终端中检查头部和元数据
要确认下载的.safetensors文件是量化模型,还是包含哪些层,加载数GB的模型是非常低效的。
safetensors库提供了一个功能,可以在不读取整个文件的情况下快速扫描头部信息。在终端输入以下Python一行命令即可快速查看。
当然,你也可以访问hugging face下载的页面,点击右上角的文件信息来查看。但这里,我将介绍如何直接在我的终端中进行检查。
准备工作
首先,需要安装库。
pip install safetensors
命令(终端输入)
请在model.safetensors的位置填入实际文件路径。
python -c "from safetensors import safe_open; \
with safe_open('model.safetensors', framework='pt', device='cpu') as f: \
print('--- Metadata ---'); \
print(f.metadata()); \
print('\n--- Tensor Keys (Layers) ---'); \
print(list(f.keys())[:5])" # 因为太长只输出前5个
输出结果解释
执行此命令后,可以获得两个重要信息。
-
Metadata: 模型创建者嵌入的信息。如果这里写有
format: gptq或quantization: int4等信息,即使文件名没有写,也可以确认是量化模型(然而,如果创建者没有填写元数据,则可能显示为None)。 -
Keys: 构成模型的层的名称。通过这些信息可以了解模型的结构。
5. 比较总结:.bin (Pickle) vs .safetensors
| 特征 | .bin / .pth (基于Pickle) | .safetensors |
|---|---|---|
| 安全性 | 风险(可能执行恶意代码) | 安全(仅保存数据) |
| 加载速度 | 较慢(CPU负载) | 非常快(零拷贝) |
| 内存效率 | 需要全部加载 | 按需加载(懒加载) |
| 兼容性 | 依赖Python/PyTorch | 框架独立 |
6. 使用示例(Python)
这是使用safetensors库保存和加载张量的简单示例。
import torch
from safetensors.torch import save_file, load_file
# 1. 创建并保存张量
tensors = {
"embedding": torch.zeros((1024, 512)),
"attention": torch.rand((512, 512))
}
# 以字典形式将张量保存到文件
save_file(tensors, "model.safetensors")
# 2. 加载文件
loaded = load_file("model.safetensors")
print(loaded["embedding"].shape)
# 输出: torch.Size([1024, 512])
7. 结论
.safetensors不仅仅是文件扩展名的变化,而是为了AI模型的安全性与效率的重要进化。包括Hugging Face在内的主要社区已经将其作为基本格式。
今后在下载或分发模型时,建议尽量使用.safetensors而非.bin。这将使您免于安全威胁,同时大幅缩短模型加载时间。

目前没有评论。