Safetensors

Safetensors est un nouveau format simple pour stocker des tenseurs de manière sécurisée (contrairement à pickle) tout en restant rapide (zéro copie). Safetensors est vraiment rapide 🚀.

Installation

avec pip :

pip install safetensors

avec conda :

conda install -c conda-forge safetensors

Utilisation

Charger les tenseurs

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)

Charger uniquement une partie des tenseurs (intéressant lors de l'exécution sur plusieurs 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]

Enregistrer les tenseurs

import torch
from safetensors.torch import save_file

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

Format

Supposons que vous ayez un fichier safetensors nommé model.safetensors, alors model.safetensors aura le format interne suivant :

Projets en vedette

Safetensors est largement utilisé par de grandes entreprises de pointe dans le domaine de l'IA, telles que Hugging Face, EleutherAI et StabilityAI. Voici une liste non exhaustive des projets qui utilisent safetensors :

Update on GitHub