Safetensors

Safetensors 是一种新的简单格式，用于安全地存储张量（与 pickle 相对），并且速度依然很快（零拷贝）。Safetensors 确实非常快速 🚀。

安装

使用 pip：

pip install safetensors

使用 conda：

conda install -c conda-forge 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)

仅加载部分张量（在多 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[:, :hidden_dim]

保存张量

import torch
from safetensors.torch import save_file

tensors = {
    "embedding": torch.zeros((2, 2)),
    "attention": torch.zeros((2, 3))
}
save_file(tensors, "model.safetensors")

格式

假设你有一个名为 model.safetensors 的 safetensors 文件，那么 model.safetensors 的内部格式如下：

精选项目

Safetensors 已被众多领先的 AI 企业广泛采用，例如 Hugging Face、EleutherAI 和 StabilityAI。以下是一个使用 safetensors 的项目的非详尽列表：

Update on GitHub