开源项目 `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 项目有了全面的了解，并能够快速上手使用。希望这篇教程对你有所帮助！