Safetensors
您正在查看的是需要从源码安装。如果你想使用常规的 pip install，请查看最新的稳定版本 (
v0.5.0-rc.0)。
加入 Hugging Face 社区
并获得增强的文档体验
Torch API
safetensors.torch.load_file
< 源代码 > ( filename: typing.Union[str, os.PathLike] device: typing.Union[str, int] = 'cpu' ) → Dict[str, torch.Tensor]
参数
- filename (
str, or os.PathLike) — 包含张量的文件名 - device (
Union[str, int], optional, defaults to cpu) — 加载后张量需要位于的设备。可用选项是所有常规的 torch 设备位置。
返回
Dict[str, torch.Tensor]
一个字典，其中键为名称，值为 torch.Tensor
将 safetensors 文件加载为 torch 格式。
示例
from safetensors.torch import load_file
file_path = "./my_folder/bert.safetensors"
loaded = load_file(file_path)
safetensors.torch.load
< 源代码 > ( data: bytes ) → Dict[str, torch.Tensor]
参数
- data (
bytes) — safetensors 文件的内容
返回
Dict[str, torch.Tensor]
一个字典，其中键为名称，值为 CPU 上的 torch.Tensor
从纯字节数据加载 safetensors 文件为 torch 格式。
示例
from safetensors.torch import load
file_path = "./my_folder/bert.safetensors"
with open(file_path, "rb") as f:
data = f.read()
loaded = load(data) safetensors.torch.save_file
< 源代码 > ( tensors: typing.Dict[str, torch.Tensor] filename: typing.Union[str, os.PathLike] metadata: typing.Optional[typing.Dict[str, str]] = None ) → None
参数
- tensors (
Dict[str, torch.Tensor]) — 输入的张量。张量需要是连续且密集的。 - filename (
str, or os.PathLike)) — 我们要保存到的文件名。 - metadata (
Dict[str, str], optional, defaults to None) — 您可能想要保存在文件头中的可选纯文本元数据。例如，它可以用于指定更多关于底层张量的信息。这纯粹是信息性的，不影响张量加载。
将张量字典以 safetensors 格式保存为原始字节。
示例
from safetensors.torch import save_file
import torch
tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
save_file(tensors, "model.safetensors") safetensors.torch.save
< 源代码 > ( tensors: typing.Dict[str, torch.Tensor] metadata: typing.Optional[typing.Dict[str, str]] = None ) → bytes
参数
- tensors (
Dict[str, torch.Tensor]) — 输入的张量。张量需要是连续且密集的。 - metadata (
Dict[str, str], optional, defaults to None) — 您可能想要保存在文件头中的可选纯文本元数据。例如，它可以用于指定更多关于底层张量的信息。这纯粹是信息性的，不影响张量加载。
代表该格式的原始字节
将张量字典以 safetensors 格式保存为原始字节。
示例
from safetensors.torch import save
import torch
tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
byte_data = save(tensors) safetensors.torch.load_model
< 源代码 > ( model: Module filename: typing.Union[str, os.PathLike] strict: bool = True device: typing.Union[str, int] = 'cpu' ) → `(missing, unexpected)
参数
- model (
torch.nn.Module) — 要加载到的模型。 - filename (
str, or os.PathLike) — 要从中加载文件的文件名位置。 - strict (
bool, optional, defaults to True) — 当存在缺失键或意外键时是否失败。当为 false 时，函数仅返回缺失和意外的名称。 - device (
Union[str, int], optional, defaults to cpu) — 加载后张量需要位于的设备。可用选项是所有常规的 torch 设备位置。
返回
`(missing, unexpected)
(List[str], List[str]) missing 是模型中在加载过程中未被修改的名称。unexpected` 是文件中存在但在加载过程中未被使用的名称。
将给定的文件名加载到 torch 模型中。此方法专门用于避免 safetensors 中不允许的张量共享问题。关于张量共享的更多信息
safetensors.torch.save_model
< 源代码 > ( model: Module filename: str metadata: typing.Optional[typing.Dict[str, str]] = None force_contiguous: bool = True )
参数
- model (
torch.nn.Module) — 要保存到磁盘的模型。 - filename (
str) — 保存文件的位置 - metadata (
Dict[str, str], optional) — 与文件一起保存的额外信息。将为每个被丢弃的张量添加一些元数据。这些信息不足以恢复整个共享结构，但可能有助于理解。 - force_contiguous (
boolean, optional, defaults to True) — 强制将 state_dict 保存为连续张量。这对模型的正确性没有影响，但如果张量的布局是为此特定原因选择的，则可能会改变性能。
将给定的 torch 模型保存到指定的文件名。此方法专门用于避免 safetensors 中不允许的张量共享问题。关于张量共享的更多信息
< > 在 GitHub 上更新
© . This site is unofficial and not affiliated with Hugging Face, Inc.