Safetensors 文档

Torch API

您正在查看 主分支 版本,需要从源代码安装。如果您想要使用常规的 pip 安装,请查看最新的稳定版本 (v0.3.2)。
Hugging Face's logo
加入 Hugging Face 社区

并获取增强文档体验的访问权限

开始使用

Torch API

safetensors.torch.load_file

< >

( filename: Union device: Union = 'cpu' ) Dict[str, torch.Tensor]

参数

  • filename (str, 或 os.PathLike) — 包含张量的文件的名称
  • device (Union[str, int], 可选,默认为 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]

包含名称作为键,值作为 torch.Tensor(在 cpu 上)的字典

从纯字节中将 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: Dict filename: Union metadata: Optional = None ) None

参数

  • 张量 (tensors) (Dict[str, torch.Tensor]) — 输入张量。张量需要是连续且密集的。
  • 文件名 (filename) (stros.PathLike) — 保存到的文件名。
  • 元数据 (metadata) (Dict[str, str]可选,默认为 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

< >

( 张量: Dict 元数据: Optional = None ) 字节

参数

  • 张量 (tensors) (Dict[str, torch.Tensor]) — 输入张量。张量需要是连续且密集的。
  • 元数据 (metadata) (Dict[str, str]可选,默认为 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

< >

( 模型: Module 文件名: Union 严格: bool = True 设备: Union = 'cpu' ) `(缺失,意外)`

参数

  • 模型 (model) (torch.nn.Module) — 要加载到的模型。
  • 文件名 (filename) (stros.PathLike) — 加载文件的路径。
  • device (Union[str, int]可选,默认为 cpu) — 加载后张量需要位于的设备。可用选项为所有常规的 torch 设备位置。

返回值

`(missing, unexpected)

(List[str], List[str]) missing表示模型中在加载过程中未修改的名称unexpected` 表示文件中存在的名称,但在加载过程中未使用。

将给定文件名加载到 PyTorch 模型中。此方法专门用于避免 safetensors 中不允许的张量共享问题。 有关张量共享的更多信息

safetensors.torch.save_model

< >

( model: Module filename: str metadata: Optional = None force_contiguous: bool = True )

参数

  • model (torch.nn.Module) — 要保存到磁盘的模型。
  • filename (str) — 保存文件的目标文件名位置。
  • metadata (Dict[str, str]可选) — 以及文件一起保存的其他信息。每个删除的张量都会添加一些元数据。这些信息不足以恢复整个共享结构,但可能有助于理解。
  • force_contiguous (布尔值可选,默认为 True) — 强制将 state_dict 保存为连续张量。这不会影响模型的正确性,但如果出于特定原因选择了张量的布局,则可能会改变性能。

将给定的 PyTorch 模型保存到指定的文件名中。此方法专门用于避免 safetensors 中不允许的张量共享问题。 有关张量共享的更多信息

< > 在 GitHub 上更新