Safetensors 文档
Torch API
并获得增强的文档体验
开始使用
Torch API
safetensors.torch.load_file
< 源代码 >( filename: typing.Union[str, os.PathLike] device: typing.Union[str, int] = 'cpu' ) → Dict[str, torch.Tensor]
将 safetensors 文件加载为 torch 格式。
safetensors.torch.load
< 源代码 >( data: bytes ) → Dict[str, torch.Tensor]
从纯字节数据加载 safetensors 文件为 torch 格式。
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
将张量字典以 safetensors 格式保存为原始字节。
safetensors.torch.save
< 源代码 >( tensors: typing.Dict[str, torch.Tensor] metadata: typing.Optional[typing.Dict[str, str]] = None ) → bytes
将张量字典以 safetensors 格式保存为原始字节。
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
, oros.PathLike
) — 要从中加载文件的文件名位置。 - strict (
bool
, optional, defaults to True) — 当存在缺失键或意外键时是否失败。当为 false 时,函数仅返回缺失和意外的名称。 - device (
Union[str, int]
, optional, defaults tocpu
) — 加载后张量需要位于的设备。可用选项是所有常规的 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 )
将给定的 torch 模型保存到指定的文件名。此方法专门用于避免 safetensors
中不允许的张量共享问题。关于张量共享的更多信息