还在用torch.save()？快来看看safetensors吧！

目录

• 1. 前言
• 2. safetensors.torch
• • 2.1 读取
  • • 2.1.1 load_file
    • 2.1.2 load
    • 2.1.3 load_model
  • 2.2 写入
  • • 2.2.1 save_file
    • 2.2.2 save
    • 2.2.3 save_model
• Ref

1. 前言

Safetensors 是一个由 HuggingFace 开发的库（采用 Rust 编写），用于安全地存储和加载张量（tensors）。与其他序列化格式（如 pickle）相比，Safetensors 提供了更快的加载速度和更高的安全性。这个库特别适用于深度学习模型的序列化，尤其是在需要高效加载大型模型权重时。

安装：

pip install safetensors

2. safetensors.torch

这一节主要讲解safetensors中有关PyTorch的API。

2.1 读取

2.1.1 load_file

函数签名如下：

def load_file(filename: Union[str, os.PathLike], device="cpu") -> Dict[str, torch.Tensor]:

用于直接读取 .safetensors 文件并将其转化成Python字典的格式：

from safetensors.torch import load_file

file_path = "./my_folder/bert.safetensors"
loaded = load_file(file_path)

2.1.2 load

函数签名如下：

def load(data: bytes) -> Dict[str, torch.Tensor]:

以二进制格式读取 .safetensors 文件并将其转化成Python字典：

from safetensors.torch import load

file_path = "./my_folder/bert.safetensors"
with open(file_path, "rb") as f:
    data = f.read()

loaded = load(data)

2.1.3 load_model

函数签名如下：

def load_model(model: torch.nn.Module, filename: str, strict=True) -> Tuple[List[str], List[str]]:

如果使用 load_file 来加载模型参数，我们通常会执行：

model.load_state_dict(load_file("model.safetensors"))

如果使用 load_model 就可以大大简洁：

load_model(model, "model.safetensors")

2.2 写入

2.2.1 save_file

函数签名如下：

def save_file(
    tensors: Dict[str, torch.Tensor],
    filename: Union[str, os.PathLike],
    metadata: Optional[Dict[str, str]] = None,
):

用于存储给定的张量（字典格式）：

from safetensors.torch import save_file
import torch

tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
save_file(tensors, "model.safetensors")

2.2.2 save

函数签名如下：

def save(tensors: Dict[str, torch.Tensor], metadata: Optional[Dict[str, str]] = None) -> bytes:

用于将给定的张量（字典格式）存储为二进制格式。

注意，该函数并不会直接生成文件，只会返回一个二进制结果：

from safetensors.torch import save
import torch

tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
byte_data = save(tensors)

2.2.3 save_model

函数签名如下：

def save_model(
    model: torch.nn.Module, filename: str, metadata: Optional[Dict[str, str]] = None, force_contiguous: bool = True
):

如果使用 save_file 来存储模型参数，我们通常会执行：

save_file(model.state_dict(), "model.safetensors")

如果使用 save_model 就可以大大简洁：

save_model(model, "model.safetensors")

---

Ref

[1] https://huggingface.co/docs/safetensors/v0.3.2/index
[2] https://blog.eleuther.ai/safetensors-security-audit/
[3] https://medium.com/@mandalsouvik/safetensors-a-simple-and-safe-way-to-store-and-distribute-tensors-d9ba1931ba04

文章知识点与官方知识档案匹配，可进一步学习相关知识

Python入门技能树人工智能深度学习422732 人正在系统学习中

本文转自 https://raelum.blog.csdn.net/article/details/136636819，如有侵权，请联系删除。