Torch API

safetensors.torch.load_file

< source >

( 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], 可选, 默认为 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

< source >

( 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

< source >

( 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], 可选, 默认为 None) — 您可能想要保存在标头中的可选纯文本元数据。例如，它可以用于更详细地说明底层张量。这纯粹是信息性的，不会影响张量加载。

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

< source >

( tensors: typing.Dict[str, torch.Tensor] metadata: typing.Optional[typing.Dict[str, str]] = None ) → bytes

参数

tensors (Dict[str, torch.Tensor]) — 传入的张量。张量需要是连续且密集的。
metadata (Dict[str, str], 可选, 默认为 None) — 您可能想要保存在标头中的可选纯文本元数据。例如，它可以用于更详细地说明底层张量。这纯粹是信息性的，不会影响张量加载。

bytes

表示格式的原始字节

以 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

< source >

( 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, 可选, 默认为 True) — 是否在缺少键或有意外键时失败。当为 false 时，该函数仅返回缺失和意外的名称。
device (Union[str, int], 可选, 默认为 cpu) — 张量加载后需要放置的设备。可用选项是所有常规 torch 设备位置。

`(missing, unexpected)

(List[str], List[str]) missing是模型中在加载期间未修改的名称 unexpected` 是文件上存在的名称，但在加载期间未使用。

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

safetensors.torch.save_model

< source >

( 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], 可选) — 与文件一起保存的额外信息。将为每个删除的张量添加一些元数据。此信息不足以恢复整个共享结构，但可能有助于理解内容
force_contiguous (boolean, 可选, 默认为 True) — 强制将 state_dict 保存为连续张量。这对模型的正确性没有影响，但如果张量的布局是为此目的专门选择的，则可能会改变性能。

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

< > 在 GitHub 上更新

Safetensors

Torch API

safetensors.torch.load_file

safetensors.torch.load

safetensors.torch.save_file

safetensors.torch.save

safetensors.torch.load_model

safetensors.torch.save_model