当前位置:   article > 正文

开源项目 `safetensors` 使用教程

如何加载model.safetensors

开源项目 safetensors 使用教程

safetensorsSimple, safe way to store and distribute tensors项目地址:https://gitcode.com/gh_mirrors/sa/safetensors

项目介绍

safetensors 是一个新的、简单的格式,用于安全地存储张量(与 pickle 相反),同时保持快速(零拷贝)。这个项目由 Hugging Face 开发,旨在提供一个高效且安全的张量存储解决方案。safetensors 的主要特点包括:

  • 安全性:避免使用 pickle,减少安全风险。
  • 速度:实现零拷贝加载,提高性能。
  • 兼容性:支持多种深度学习框架,如 PyTorch、TensorFlow 等。

项目快速启动

安装

首先,你需要安装 safetensors。你可以通过 pipconda 进行安装:

pip install safetensors

或者使用 conda

conda install -c huggingface safetensors

加载张量

以下是一个简单的示例,展示如何加载 safetensors 格式的张量:

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

保存张量

以下是一个示例,展示如何将张量保存为 safetensors 格式:

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

应用案例和最佳实践

案例一:在多 GPU 环境下加载部分张量

在多 GPU 环境下,你可能只需要加载部分张量。以下是一个示例:

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

最佳实践

  • 安全性:始终使用 safetensors 格式来存储模型权重,避免使用 pickle
  • 性能:利用 safetensors 的零拷贝特性,提高模型加载速度。
  • 兼容性:确保你的项目支持多种深度学习框架。

典型生态项目

safetensors 已经被多个领先的 AI 企业和项目采用,以下是一些典型的生态项目:

  • Hugging Face Transformers:用于自然语言处理的库,广泛使用 safetensors 格式。
  • Stable Diffusion WebUI:用于图像生成的项目,采用 safetensors 格式存储模型权重。
  • ColossalAI:一个大规模深度学习库,支持 safetensors 格式。

这些项目展示了 safetensors 在不同领域的应用和其强大的生态支持。


通过以上内容,你应该对 safetensors 项目有了全面的了解,并能够快速上手使用。希望这篇教程对你有所帮助!

safetensorsSimple, safe way to store and distribute tensors项目地址:https://gitcode.com/gh_mirrors/sa/safetensors

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/笔触狂放9/article/detail/999533
推荐阅读
相关标签
  

闽ICP备14008679号