一、引言
这里的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模型。它通过加密和验证模型数据来增强安全性，防止数据篡改。
