Safetensors，高效安全易用的深度学习新工具

大家好，本文将介绍一种为深度学习应用提供速度、效率、跨平台兼容性、用户友好性和安全性的新工具。

Safetensors简介

Hugging Face开发了一种名为Safetensors的新序列化格式，旨在简化和精简大型复杂张量的存储和加载。张量是深度学习中使用的主要数据结构，其大小会给效率带来挑战。

Safetensors结合使用高效的序列化和压缩算法来减少大型张量的大小，使其比pickle等其他序列化格式更快、更高效。这意味着，与传统PyTorch序列化格式pytorch_model.bin和model.safetensors相比，Safetensors在CPU上的速度快76.6倍，在GPU上的速度快2倍。

 使用Safetensors的好处

Safetensors具有简单直观的API，可以在Python中序列化和反序列化张量。这意味着开发人员可以专注于搭建深度学习模型，而不必在序列化和反序列化上花费时间。

可以用Python进行序列化，并方便地使用各种编程语言和平台（如C++、Java和JavaScript）加载生成的文件，这样就可以实现在不同的编程环境中无缝共享模型。

Safetensors针对速度进行了优化，可以高效处理大型张量的序列化和反序列化，因此它是使用大型语言模型的应用程序的绝佳选择。

它混合使用了有效的序列化和压缩算法，以减小大型张量的大小，与其他序列化格式（如pickle）相比，性能更快、更高效。

为了防止序列化张量在存储或传输过程中出现损坏，Safetensors使用了校验和机制。这保证了额外的安全性，确保存储在Safetensors中的所有数据都准确可靠。此外，它还能防止DOS攻击。

在使用多个节点或GPU的分布式环境中工作时，只在每个模型上加载部分张量是很有帮助的。BLOOM利用这种格式在8个 GPU上加载模型仅需45秒，而普通PyTorch加权则需要10分钟。

使用Safetensors

在本节中我们将介绍safetensors API，以及如何保存和加载张量文件。可以使用pip管理器安装safetensors：

pip install safetensors

本文将使用Torch共享张量中的示例来搭建一个简单的神经网络，并使用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())

正如所看到的，已经成功创建了模型。

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

现在我们将通过提供model对象和文件名来保存模型，然后把保存的文件加载到使用nn.Module创建的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创建的张量，为此将使用save_file函数。

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函数。

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文档来了解它。

简而言之，Safetensors是一种存储深度学习应用中使用的大型张量的新方法。与其他技术相比，它具有更快、更高效和用户友好的特点，此外它还能确保数据的保密性和安全性，同时支持各种编程语言和平台。通过使用Safetensors，机器学习工程师可以优化时间，专注于开发更优秀的模型。

强烈推荐在项目中使用Safetensors，许多顶级AI公司，如Hugging Face、EleutherAI和StabilityAI，都在他们的项目中使用了Safetensors。