【Python | 深度学习】safetensors 包的介绍和使用案例（含源代码）

safetensors 是一种用于安全存储张量（与 pickle 相反）的新型简单格式，并且仍然很快（零拷贝）。

safetensors 真的很快。

一、安装

1.1 pip 安装

pip install safetensors

• 1

1.2 conda 安装

conda install -c huggingface safetensors

• 1

二、加载张量

from safetensors import safe_open tensors = {} with safe_open("model.safetensors", framework="pt", device=0) as f: for k in f.keys(): tensors[k] = f.get_tensor(k)

• 1
• 2
• 3
• 4
• 5
• 6

仅加载部分张量（在多个GPU上运行时很有趣）：

from safetensors import safe_open tensors = {} with safe_open("model.safetensors", framework="pt", device=0) as f: tensor_slice = f.get_slice("embedding") vocab_size, hidden_dim = tensor_slice.get_shape() tensor = tensor_slice[:, :hidden_dim]

• 1
• 2
• 3
• 4
• 5
• 6
• 7

三、保存张量

import torch from safetensors.torch import save_file tensors = { "embedding": torch.zeros((2, 2)), "attention": torch.zeros((2, 3)) } save_file(tensors, "model.safetensors")

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8

四、速度比较

4.1 下载 gpt2 的文件

safetensors 真的很快。让我们通过加载 gpt2 权重将其进行比较。要运行 GPU 基准测试，请确保您的机器具有 GPU，或者您已选择是否使用的是 Google Colab。

在开始之前，请确保已安装所有必要的库：

pip install safetensors huggingface_hub torch

• 1

让我们从导入所有将使用的包开始：

import os import datetime from huggingface_hub import hf_hub_download from safetensors.torch import load_file import torch

• 1
• 2
• 3
• 4
• 5

Download safetensors & torch weights for gpt2：

sf_filename = hf_hub_download("gpt2", filename="model.safetensors") pt_filename = hf_hub_download("gpt2", filename="pytorch_model.bin")

• 1
• 2

4.2 CPU 基准测试

start_st = datetime.datetime.now() weights = load_file(sf_filename, device="cpu") load_time_st = datetime.datetime.now() - start_st print(f"Loaded safetensors {load_time_st}")

• 1
• 2
• 3
• 4

输出结果为：

Loaded safetensors 0:00:00.026842

• 1

start_pt = datetime.datetime.now() weights = torch.load(pt_filename, map_location="cpu") load_time_pt = datetime.datetime.now() - start_pt print(f"Loaded pytorch {load_time_pt}")

• 1
• 2
• 3
• 4

输出结果为：

Loaded pytorch 0:00:00.182266

• 1

print(f"on CPU, safetensors is faster than pytorch by: {load_time_pt/load_time_st:.1f} X")

• 1

输出结果为：

on CPU, safetensors is faster than pytorch by: 6.8 X

• 1

这种加速是由于该库通过直接映射文件来避免不必要的副本。实际上可以在 torch 上完成。 当前显示的加速比已打开：

• 操作系统： Windows
• 处理器： 英特尔® 至强® CPU @ 2.00GHz

4.3 GPU 基准测试

os.environ["SAFETENSORS_FAST_GPU"] = "1" torch.zeros((2, 2)).cuda() start_st = datetime.datetime.now() weights = load_file(sf_filename, device="cuda:0") load_time_st = datetime.datetime.now() - start_st print(f"Loaded safetensors {load_time_st}") start_pt = datetime.datetime.now() weights = torch.load(pt_filename, map_location="cuda:0") load_time_pt = datetime.datetime.now() - start_pt print(f"Loaded pytorch {load_time_pt}") print(f"on GPU, safetensors is faster than pytorch by: {load_time_pt/load_time_st:.1f} X")

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10
• 11

输出结果为：

Loaded safetensors 0:00:00.497415 Loaded pytorch 0:00:00.250602 on GPU, safetensors is faster than pytorch by: 0.5 X

• 1
• 2
• 3

加速有效是因为此库能够跳过不必要的 CPU 分配。不幸的是，据我们所知，它无法在纯 pytorch 中复制。该库的工作原理是内存映射文件，使用 pytorch 创建空张量，并直接调用以直接在 GPU 上移动张量。

显卡：GTX 3060