safetensors 개요 — 설치와 사용법

safetensors 개요 — 설치와 사용법

safetensors는 텐서를 pickle 대신 안전하게, 그리고 zero-copy로 빠르게 저장하는 포맷이에요. "Safetensors is really fast 🚀"라는 말처럼 로딩 속도가 큰 강점이에요.

출처: Safetensors 문서

설치

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")

어디서 쓰나요

transformers, mlx, candle, stable-diffusion-webui, llama.cpp(convert.py), diffusers, oobabooga/text-generation-webui, CivitAI 등 광범위한 프로젝트에서 사용해요.

더 알아보기

• Safetensors 문서
• Safetensors GitHub