什么是Safetensors？

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

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

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

使用Safetensors的好處

易用性

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

跨平臺(tái)兼容性

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

速度

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

大小優(yōu)化

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

安全

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

懶性加載

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

開(kāi)始使用Safetensors

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

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

pip install safetensors

本文將使用Torch共享張量中的示例來(lái)搭建一個(gè)簡(jiǎn)單的神經(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)在，我們將通過(guò)提供model對(duì)象和文件名來(lái)保存模型。然后，我們將把保存的文件加載到使用nn.Module創(chuàng)建的model對(duì)象中。

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], ...,

在第二個(gè)示例中，我們將嘗試保存使用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。可以通過(guò)閱讀Safetensors文檔來(lái)了解它。

圖片來(lái)自Torch API

結(jié)論

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

強(qiáng)烈推薦在項(xiàng)目中使用Safetensors。許多頂級(jí)AI公司，如Hugging Face、EleutherAI和StabilityAI，都在他們的項(xiàng)目中使用了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