safetensors Torch API
safetensors Torch API
safetensors의 PyTorch 바인딩은 파일 로드·저장, 그리고 torch 모델 단위 로드·저장 함수를 제공해요.
safetensors.torch.load_file
safetensors 파일을 torch 포맷으로 로드해요.
from safetensors.torch import load_file
file_path = "./my_folder/bert.safetensors"
loaded = load_file(file_path)
파라미터:
filename: 텐서가 들어 있는 파일 경로device(기본"cpu"): 로드 후 텐서가 위치할 디바이스backend(기본"mmap"): 텐서 바이트를 서빙하는 스토리지 백엔드."mmap"(기본)과"pread"지원
반환: Dict[str, torch.Tensor]
safetensors.torch.load
순수 바이트에서 torch 포맷으로 로드해요.
from safetensors.torch import load
file_path = "./my_folder/bert.safetensors"
with open(file_path, "rb") as f:
data = f.read()
loaded = load(data)
safetensors.torch.save_file
텐서 딕셔너리를 filename에 safetensors 포맷으로 저장해요.
from safetensors.torch import save_file
import torch
tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
save_file(tensors, "model.safetensors")
safetensors.torch.save
텐서 딕셔너리를 raw 바이트로 저장해요.
from safetensors.torch import save
import torch
tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
byte_data = save(tensors)
safetensors.torch.load_model / save_model
torch 모델(torch.nn.Module) 단위로 로드/저장해요. 텐서 공유(tensor sharing) 문제를 피하기 위해 존재해요 — safetensors는 텐서 공유를 허용하지 않거든요.
from safetensors.torch import save_model, load_model
# save_model(model, filename, force_contiguous=True)
# load_model(model, filename)
force_contiguous (기본 True): state_dict를 연속(contiguous) 텐서로 강제 저장. 모델 정확성엔 영향이 없지만 성능 레이아웃에 영향을 줄 수 있어요.