Safetensors

Safetensors 是一種新的簡潔格式，用於安全地儲存張量（與 pickle 不同），而且速度仍然很快（零拷貝）。Safetensors 真的很快 🚀。

安裝

使用 pip

pip install safetensors

使用 conda

conda install -c huggingface 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 檔案，那麼該檔案將具有以下內部格式

精選專案

Safetensors 被廣泛應用於領先的 AI 企業，例如 Hugging Face、EleutherAI 和 StabilityAI。以下是一些使用 Safetensors 的專案列表（非詳盡）

在 GitHub 上更新