开源项目 safetensors
使用教程
项目介绍
safetensors
是一个新的、简单的格式,用于安全地存储张量(与 pickle
相反),同时保持快速(零拷贝)。这个项目由 Hugging Face 开发,旨在提供一个高效且安全的张量存储解决方案。safetensors
的主要特点包括:
- 安全性:避免使用
pickle
,减少安全风险。 - 速度:实现零拷贝加载,提高性能。
- 兼容性:支持多种深度学习框架,如 PyTorch、TensorFlow 等。
项目快速启动
安装
首先,你需要安装 safetensors
。你可以通过 pip
或 conda
进行安装:
pip install safetensors
或者使用 conda
:
conda install -c huggingface safetensors
加载张量
以下是一个简单的示例,展示如何加载 safetensors
格式的张量:
from safetensors import safe_open
tensors = {}
with safe_open("model.safetensors", framework="pt", device=0) as f:
for k in f.keys():
tensors[k] = f.get_tensor(k)
保存张量
以下是一个示例,展示如何将张量保存为 safetensors
格式:
import torch
from safetensors.torch import save_file
tensors = {
"embedding": torch.zeros((2, 2)),
"attention": torch.zeros((2, 3))
}
save_file(tensors, "model.safetensors")
应用案例和最佳实践
案例一:在多 GPU 环境下加载部分张量
在多 GPU 环境下,你可能只需要加载部分张量。以下是一个示例:
from safetensors import safe_open
tensors = {}
with safe_open("model.safetensors", framework="pt", device=0) as f:
tensor_slice = f.get_slice("embedding")
vocab_size, hidden_dim = tensor_slice.get_shape()
tensor = tensor_slice[:vocab_size, :hidden_dim]
最佳实践
- 安全性:始终使用
safetensors
格式来存储模型权重,避免使用pickle
。 - 性能:利用
safetensors
的零拷贝特性,提高模型加载速度。 - 兼容性:确保你的项目支持多种深度学习框架。
典型生态项目
safetensors
已经被多个领先的 AI 企业和项目采用,以下是一些典型的生态项目:
- Hugging Face Transformers:用于自然语言处理的库,广泛使用
safetensors
格式。 - Stable Diffusion WebUI:用于图像生成的项目,采用
safetensors
格式存储模型权重。 - ColossalAI:一个大规模深度学习库,支持
safetensors
格式。
这些项目展示了 safetensors
在不同领域的应用和其强大的生态支持。
通过以上内容,你应该对 safetensors
项目有了全面的了解,并能够快速上手使用。希望这篇教程对你有所帮助!