Python 官方文档:入门教程 => 点击学习
safetensors 是一种用于安全存储张量(与 pickle 相反)的新型简单格式,并且仍然很快(零拷贝)。 safetensors 真的很快。 一、安装 1.1 pip 安装 pip insta
safetensors 是一种用于安全存储张量(与 pickle 相反)的新型简单格式,并且仍然很快(零拷贝)。
safetensors 真的很快。
pip install safetensors
conda install -c huggingface safetensors
from safetensors import safe_opentensors = {}with safe_open("model.safetensors", framework="pt", device=0) as f: for k in f.keys(): tensors[k] = f.get_tensor(k)
仅加载部分张量(在多个GPU上运行时很有趣):
from safetensors import safe_opentensors = {}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]
import torchfrom safetensors.torch import save_filetensors = { "embedding": torch.zeros((2, 2)), "attention": torch.zeros((2, 3))}save_file(tensors, "model.safetensors")
safetensors 真的很快。让我们通过加载 gpt2 权重将其进行比较。要运行 GPU 基准测试,请确保您的机器具有 GPU,或者您已选择是否使用的是 Google Colab。
在开始之前,请确保已安装所有必要的库:
pip install safetensors huggingface_hub torch
让我们从导入所有将使用的包开始:
import osimport datetimefrom huggingface_hub import hf_hub_downloadfrom safetensors.torch import load_fileimport torch
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")
start_st = datetime.datetime.now()weights = load_file(sf_filename, device="cpu")load_time_st = datetime.datetime.now() - start_stprint(f"Loaded safetensors {load_time_st}")
输出结果为:
Loaded safetensors 0:00:00.026842
start_pt = datetime.datetime.now()weights = torch.load(pt_filename, map_location="cpu")load_time_pt = datetime.datetime.now() - start_ptprint(f"Loaded pytorch {load_time_pt}")
输出结果为:
Loaded pytorch 0:00:00.182266
print(f"on CPU, safetensors is faster than pytorch by: {load_time_pt/load_time_st:.1f} X")
输出结果为:
on CPU, safetensors is faster than pytorch by: 6.8 X
这种加速是由于该库通过直接映射文件来避免不必要的副本。实际上可以在 torch 上完成。 当前显示的加速比已打开:
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_stprint(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_ptprint(f"Loaded pytorch {load_time_pt}")print(f"on GPU, safetensors is faster than pytorch by: {load_time_pt/load_time_st:.1f} X")
输出结果为:
Loaded safetensors 0:00:00.497415Loaded pytorch 0:00:00.250602on GPU, safetensors is faster than pytorch by: 0.5 X
加速有效是因为此库能够跳过不必要的 CPU 分配。不幸的是,据我们所知,它无法在纯 pytorch 中复制。该库的工作原理是内存映射文件,使用 pytorch 创建空张量,并直接调用以直接在 GPU 上移动张量。
显卡:GTX 3060
来源地址:https://blog.csdn.net/wzk4869/article/details/130668642
--结束END--
本文标题: 【Python | 深度学习】safetensors 包的介绍和使用案例(含源代码)
本文链接: https://www.lsjlt.com/news/429630.html(转载时请注明来源链接)
有问题或投稿请发送至: 邮箱/279061341@qq.com QQ/279061341
下载Word文档到电脑,方便收藏和打印~
2024-03-01
2024-03-01
2024-03-01
2024-02-29
2024-02-29
2024-02-29
2024-02-29
2024-02-29
2024-02-29
2024-02-29
回答
回答
回答
回答
回答
回答
回答
回答
回答
回答
0