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

最新推荐文章于 2025-03-26 20:55:49 发布

Python学研大本营

最新推荐文章于 2025-03-26 20:55:49 发布

阅读量3.2k

点赞数 5

文章标签：安全深度学习人工智能

本文链接：https://blog.csdn.net/weixin_39915649/article/details/134456765

版权

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

微信搜索关注《Python学研大本营》，加入读者群，分享更多精彩

什么是Safetensors？

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

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

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

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