赞
踩
safetensors
使用教程safetensors
是一个新的、简单的格式,用于安全地存储张量(与 pickle
相反),同时保持快速(零拷贝)。这个项目由 Hugging Face 开发,旨在提供一个高效且安全的张量存储解决方案。safetensors
的主要特点包括:
pickle
,减少安全风险。首先,你需要安装 safetensors
。你可以通过 pip
或 conda
进行安装:
pip install safetensors
或者使用 conda
:
conda install -c huggingface safetensors
以下是一个简单的示例,展示如何加载 safetensors
格式的张量:
- 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)
以下是一个示例,展示如何将张量保存为 safetensors
格式:
- import torch
- from safetensors.torch import save_file
-
- tensors = {
- "embedding": torch.zeros((2, 2)),
- "attention": torch.zeros((2, 3))
- }
- save_file(tensors, "model.safetensors")
在多 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[:vocab_size, :hidden_dim]
safetensors
格式来存储模型权重,避免使用 pickle
。safetensors
的零拷贝特性,提高模型加载速度。safetensors
已经被多个领先的 AI 企业和项目采用,以下是一些典型的生态项目:
safetensors
格式。safetensors
格式存储模型权重。safetensors
格式。这些项目展示了 safetensors
在不同领域的应用和其强大的生态支持。
通过以上内容,你应该对 safetensors
项目有了全面的了解,并能够快速上手使用。希望这篇教程对你有所帮助!
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。