【AI大模型】Transformers大模型库（十六）：safetensors存储类型

一、引言

这里的Transformers指的是huggingface开发的大模型库，为huggingface上数以万计的预训练大模型提供预测、训练等服务。

🤗 Transformers 提供了数以千计的预训练模型，支持 100 多种语言的文本分类、信息抽取、问答、摘要、翻译、文本生成。它的宗旨是让最先进的 NLP 技术人人易用。 🤗 Transformers 提供了便于快速下载和使用的API，让你可以把预训练模型用在给定文本、在你的数据集上微调然后通过 model hub 与社区共享。同时，每个定义的 Python 模块均完全独立，方便修改和快速研究实验。 🤗 Transformers 支持三个最热门的深度学习库： Jax, PyTorch 以及 TensorFlow — 并与之无缝整合。你可以直接使用一个框架训练你的模型然后用另一个加载和推理。

本文重点介绍safetensors库用法。

二、safetensors库

2.1 概述

safetensors是一个库，旨在安全地存储和加载机器学习模型的权重，特别是针对PyTorch模型。它通过加密和验证模型数据来增强安全性，防止数据篡改。

2.2 使用方法

2.2.1 安装safetensors

首先，确保你安装了safetensors库。可以通过pip安装：

   pip install safetensors

2.2.2 保存模型权重

使用safetensors保存模型权重，而不是直接使用PyTorch的.save()方法。

   import torch
   from safetensors.torch import save_file

   # 假设model是你的模型实例
   model_state_dict = model.state_dict()
   # 保存模型到safetensors格式
   save_file(model_state_dict, "model.safetensors")

对应的pytorch保存模型的方法

# 保存模型状态字典
torch.save(model.state_dict(), 'model.pth')

# 加载模型状态字典
model = YourModelClass()  # 初始化模型实例
model.load_state_dict(torch.load('model.pth'))  # 加载权重
model.eval()  # 如果是预训练模型，通常设置为评估模式

2.2.3 加载模型权重

加载时，同样使用safetensors的专用函数。

   from safetensors.torch import load_file

   # 加载模型权重
   loaded_state_dict = load_file("model.safetensors")
   # 加载到模型中
   model.load_state_dict(loaded_state_dict)

使用safetensors时，模型的加载和保存方式与直接使用PyTorch的.pt或.pth文件不同，它提供了额外的安全特性，特别是在模型的分发和共享方面

三、总结

本篇内容展示了如何使用safetensors库，主要功能旨在安全地存储和加载机器学习模型的权重，特别是针对PyTorch模型。它通过加密和验证模型数据来增强安全性，防止数据篡改。