什么是Safetensors？

【Safetensors】：https://huggingface.co/docs/safetensors/index

Hugging Face開發(fā)了一種名為Safetensors的新序列化格式，旨在簡化和精簡大型復(fù)雜張量的存儲和加載。張量是深度學習中使用的主要數(shù)據(jù)結(jié)構(gòu)，其大小會給效率帶來挑戰(zhàn)。

Safetensors結(jié)合使用高效的序列化和壓縮算法來減少大型張量的大小，使其比pickle等其他序列化格式更快、更高效。這意味著，與傳統(tǒng)PyTorch序列化格式pytorch_model.bin和model.safetensors相比，Safetensors在CPU上的速度快76.6倍，在GPU上的速度快2倍。請查看速度比較（https://huggingface.co/docs/safetensors/speed）。

使用Safetensors的好處

易用性

Safetensors具有簡單直觀的API，可以在Python中序列化和反序列化張量。這意味著開發(fā)人員可以專注于搭建深度學習模型，而不必在序列化和反序列化上花費時間。

跨平臺兼容性

可以用Python進行序列化，并方便地使用各種編程語言和平臺（如C++、Java和JavaScript）加載生成的文件。這樣就可以實現(xiàn)在不同的編程環(huán)境中無縫共享模型。

速度

Safetensors針對速度進行了優(yōu)化，可以高效處理大型張量的序列化和反序列化。因此，它是使用大型語言模型的應(yīng)用程序的絕佳選擇。

大小優(yōu)化

它混合使用了有效的序列化和壓縮算法，以減小大型張量的大小，與其他序列化格式（如pickle）相比，性能更快、更高效。

安全

為了防止序列化張量在存儲或傳輸過程中出現(xiàn)損壞，Safetensors使用了校驗和機制。這保證了額外的安全性，確保存儲在Safetensors中的所有數(shù)據(jù)都準確可靠。此外，它還能防止DOS攻擊。

懶性加載

在使用多個節(jié)點或GPU的分布式環(huán)境中工作時，只在每個模型上加載部分張量是很有幫助的。BLOOM利用這種格式在8個 GPU上加載模型僅需45秒，而普通PyTorch加權(quán)則需要10分鐘。

開始使用Safetensors

在本節(jié)中，我們將介紹safetensors API，以及如何保存和加載張量文件。

可以使用pip管理器安裝safetensors：

pip install safetensors

本文將使用Torch共享張量中的示例來搭建一個簡單的神經(jīng)網(wǎng)絡(luò)，并使用PyTorch的safetensors.torch API保存模型。

from torch import nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = nn.Linear(100, 100)
        self.b = self.a

    def forward(self, x):
        return self.b(self.a(x))


model = Model()
print(model.state_dict())

正如所看到的，已經(jīng)成功創(chuàng)建了模型。

OrderedDict([('a.weight', tensor([[-0.0913, 0.0470, -0.0209, ..., -0.0540, -0.0575, -0.0679], [ 0.0268, 0.0765, 0.0952, ..., -0.0616, 0.0146, -0.0343], [ 0.0216, 0.0444, -0.0347, ..., -0.0546, 0.0036, -0.0454], ...,

現(xiàn)在，我們將通過提供model對象和文件名來保存模型。然后，我們將把保存的文件加載到使用nn.Module創(chuàng)建的model對象中。

from safetensors.torch import load_model, save_model

save_model(model, "model.safetensors")

load_model(model, "model.safetensors")
print(model.state_dict())

OrderedDict([('a.weight', tensor([[-0.0913, 0.0470, -0.0209, ..., -0.0540, -0.0575, -0.0679], [ 0.0268, 0.0765, 0.0952, ..., -0.0616, 0.0146, -0.0343], [ 0.0216, 0.0444, -0.0347, ..., -0.0546, 0.0036, -0.0454], ...,

在第二個示例中，我們將嘗試保存使用torch.zeros創(chuàng)建的張量。為此，我們將使用save_file函數(shù)。

import torch
from safetensors.torch import save_file, load_file

tensors = {
   "weight1": torch.zeros((1024, 1024)),
   "weight2": torch.zeros((1024, 1024))
}
save_file(tensors, "new_model.safetensors")

為了加載張量，我們將使用load_file函數(shù)。

load_file("new_model.safetensors")

{'weight1': tensor([[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]]),
 'weight2': tensor([[0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         ...,
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.],
         [0., 0., 0.,  ..., 0., 0., 0.]])}

Safetensors API適用于Pytorch、Tensorflow、PaddlePaddle、Flax和Numpy。可以通過閱讀Safetensors文檔來了解它。

圖片來自Torch API

結(jié)論

簡而言之，Safetensors是一種存儲深度學習應(yīng)用中使用的大型張量的新方法。與其他技術(shù)相比，它具有更快、更高效和用戶友好的特點。此外，它還能確保數(shù)據(jù)的保密性和安全性，同時支持各種編程語言和平臺。通過使用Safetensors，機器學習工程師可以優(yōu)化時間，專注于開發(fā)更優(yōu)秀的模型。

強烈推薦在項目中使用Safetensors。許多頂級AI公司，如Hugging Face、EleutherAI和StabilityAI，都在他們的項目中使用了Safetensors。

參考資料

文檔：Safetensors（https://huggingface.co/docs/safetensors/index）

博客：https://medium.com/@zergtant/what-is-safetensors-and-how-to-convert-ckpt-model-to-safetensors-13d36eb94d57

GitHub：https://github.com/huggingface/safetensors