5️⃣Safetensors

Huggingface Safetensors

Safetensors는 텐서를 안전하게 저장하는(ML에서 자주 사용하는 pikle과 반대되는) 새로운 간단한 형식이며, 장점은 Zerocopy로 빠릅니다.

saftensors

실제 Huggingface에서 모델을 push 하면 자동적으로 safetensors로 변환되기 때문에 대다수 Model이 safetensors로 이루어져 있다.

%pip install safetensors

Requirement already satisfied: safetensors in /home/kubwa/anaconda3/envs/pytorch/lib/python3.11/site-packages (0.4.2)
Note: you may need to restart the kernel to use updated packages.

Load tensors

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)

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]

Save tensors

import torch
from safetensors.torch import save_file

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