序列化
huggingface_hub 包含一些帮助程序，可以帮助机器学习库以标准化方式序列化模型权重。这部分库仍在开发中，将在未来的版本中得到改进。目标是统一权重在 Hub 上的序列化方式，既可以消除库之间的代码重复，也可以促进 Hub 上的约定。
保存 torch 状态字典
serialization 模块的主要帮助程序将 torch nn.Module 作为输入并将其保存到磁盘。它处理保存共享张量（参见 safetensors 说明）的逻辑，以及使用 split_torch_state_dict_into_shards() 将状态字典拆分为分片的逻辑。目前，仅支持 torch 框架。
如果要保存状态字典（例如层名称和相关张量之间的映射）而不是 nn.Module，可以使用 save_torch_state_dict()，它提供了相同的功能。例如，如果要在保存状态字典之前对其应用自定义逻辑，这将非常有用。
huggingface_hub.save_torch_model
< 源代码 >( model: torch.nn.Module save_directory: Union filename_pattern: Optional = None force_contiguous: bool = True max_shard_size: Union = '5GB' metadata: Optional = None safe_serialization: bool = True )
参数
- model (
torch.nn.Module) — 要保存到磁盘的模型。 - save_directory (
str或Path) — 模型将保存到的目录。 - filename_pattern (
str, 可选) — 用于生成保存模型的文件名的模式。模式必须是一个可以使用filename_pattern.format(suffix=...)格式化的字符串，并且必须包含关键字suffix。默认为"model{suffix}.safetensors"或pytorch_model{suffix}.bin，具体取决于safe_serialization参数。 - force_contiguous (
boolean, 可选) — 强制将 state_dict 保存为连续张量。这对模型的正确性没有影响，但如果张量的布局是专门为此目的选择的，则可能会改变性能。默认为True。 - max_shard_size (
int或str, 可选) — 每个分片的最大大小（以字节为单位）。默认为 5GB。 - metadata (
Dict[str, str], 可选) — 要与模型一起保存的额外信息。对于每个丢弃的张量，都会添加一些元数据。这些信息不足以恢复整个共享结构，但可能有助于理解情况。 - safe_serialization (
bool，可选) — 是否以 safetensors 格式保存，这是默认行为。如果为False，则分片将保存为 pickle 格式。出于安全原因，建议使用安全序列化。以 pickle 格式保存已被弃用，并将在未来的版本中删除。
将给定的 PyTorch 模型保存到磁盘，处理分片和共享张量问题。
另请参阅 save_torch_state_dict() 以更灵活地保存状态字典。
有关张量共享的更多信息，请查看 本指南。
模型状态字典被拆分为多个分片，以便每个分片都小于给定的大小。分片将使用给定的 filename_pattern 保存到 save_directory 中。如果模型太大而无法放入单个分片中，则会在 save_directory 中保存一个索引文件，以指示每个张量的保存位置。此辅助函数在底层使用 split_torch_state_dict_into_shards()。如果 safe_serialization 为 True，则分片将保存为 safetensors 格式（默认）。否则，分片将保存为 pickle 格式。
在保存模型之前，会清除 save_directory 中所有先前的分片文件。
如果模型的某个张量大于 max_shard_size，它最终将位于其自己的分片中，该分片的大小将大于 max_shard_size。
示例
>>> from huggingface_hub import save_torch_model
>>> model = ... # A PyTorch model
# Save state dict to "path/to/folder". The model will be split into shards of 5GB each and saved as safetensors.
>>> save_torch_model(model, "path/to/folder")
# Load model back
>>> from huggingface_hub import load_torch_model # TODO
>>> load_torch_model(model, "path/to/folder")
>>>huggingface_hub.save_torch_state_dict
< 源代码 >( state_dict: Dict save_directory: Union filename_pattern: Optional = None force_contiguous: bool = True max_shard_size: Union = '5GB' metadata: Optional = None safe_serialization: bool = True )
参数
- state_dict (
Dict[str, torch.Tensor]) — 要保存的状态字典。 - save_directory (
str或Path) — 模型将保存到的目录。 - filename_pattern (
str，可选) — 用于生成模型将保存到的文件名的模式。模式必须是一个可以使用filename_pattern.format(suffix=...)格式化的字符串，并且必须包含关键字suffix。默认为"model{suffix}.safetensors"或pytorch_model{suffix}.bin，具体取决于safe_serialization参数。 - force_contiguous (
boolean，可选) — 强制将 state_dict 保存为连续张量。这对模型的正确性没有影响，但如果张量的布局是专门为此原因选择的，则它可能会改变性能。默认为True。 - max_shard_size (
int或str，可选) — 每个分片的最大大小，以字节为单位。默认为 5GB。 - metadata (
Dict[str, str]，可选) — 随模型一起保存的额外信息。每个删除的张量都会添加一些元数据。这些信息不足以恢复整个共享结构，但可能有助于理解一些事情。 - safe_serialization (
bool，可选) — 是否保存为 safetensors，这是默认行为。如果为False，则分片将保存为 pickle。出于安全原因，建议使用安全序列化。保存为 pickle 已弃用，将在未来的版本中删除。
将模型状态字典保存到磁盘，处理分片和共享张量问题。
另请参阅 save_torch_model() 以直接保存 PyTorch 模型。
有关张量共享的更多信息，请查看 本指南。
模型状态字典被拆分为多个分片，以便每个分片都小于给定的大小。分片将使用给定的 filename_pattern 保存到 save_directory 中。如果模型太大而无法放入单个分片中，则会在 save_directory 中保存一个索引文件，以指示每个张量的保存位置。此辅助函数在底层使用 split_torch_state_dict_into_shards()。如果 safe_serialization 为 True，则分片将保存为 safetensors 格式（默认）。否则，分片将保存为 pickle 格式。
在保存模型之前，会清除 save_directory 中所有先前的分片文件。
如果模型的某个张量大于 max_shard_size，它最终将位于其自己的分片中，该分片的大小将大于 max_shard_size。
示例
>>> from huggingface_hub import save_torch_state_dict
>>> model = ... # A PyTorch model
# Save state dict to "path/to/folder". The model will be split into shards of 5GB each and saved as safetensors.
>>> state_dict = model_to_save.state_dict()
>>> save_torch_state_dict(state_dict, "path/to/folder")将状态字典拆分为分片
serialization 模块还包含低级助手，用于将状态字典拆分为多个分片，同时在此过程中创建适当的索引。这些助手可用于 torch 和 tensorflow 张量，并且旨在轻松扩展到任何其他机器学习框架。
split_tf_state_dict_into_shards
huggingface_hub.split_tf_state_dict_into_shards
< 源代码 >( state_dict: Dict filename_pattern: str = 'tf_model{suffix}.h5' max_shard_size: Union = '5GB' ) → StateDictSplit
参数
- state_dict (
Dict[str, Tensor]) — 要保存的状态字典。 - filename_pattern (
str, 可选) — 用于生成保存模型的文件名的模式。模式必须是一个可以使用filename_pattern.format(suffix=...)格式化的字符串，并且必须包含关键字suffix。默认为"tf_model{suffix}.h5"。 - max_shard_size (
int或str, 可选) — 每个分片的最大大小（以字节为单位）。默认为 5GB。
返回值
StateDictSplit
一个包含分片和检索它们索引的 StateDictSplit 对象。
将模型状态字典拆分为多个分片，以便每个分片都小于给定的大小。
分片是通过按其键的顺序迭代 state_dict 来确定的。没有进行优化以使每个分片尽可能接近传递的最大大小。例如，如果限制为 10GB，并且我们有大小为 [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] 的张量，它们将被分片为 [6GB], [6+2GB], [6+2+2GB]，而不是 [6+2+2GB], [6+2GB], [6GB]。
如果模型的某个张量大于 max_shard_size，它最终将位于其自己的分片中，该分片的大小将大于 max_shard_size。
split_torch_state_dict_into_shards
huggingface_hub.split_torch_state_dict_into_shards
< 源代码 >( state_dict: Dict filename_pattern: str = 'model{suffix}.safetensors' max_shard_size: Union = '5GB' ) → StateDictSplit
参数
- state_dict (
Dict[str, torch.Tensor]) — 要保存的状态字典。 - filename_pattern (
str，可选) — 用于生成模型保存文件名的模式。模式必须是一个字符串，可以使用filename_pattern.format(suffix=...)格式化，并且必须包含关键字suffix。默认为"model{suffix}.safetensors"。 - max_shard_size (
int或str，可选) — 每个分片的最大大小，以字节为单位。默认为 5GB。
返回值
StateDictSplit
一个包含分片和检索它们索引的 StateDictSplit 对象。
将模型状态字典拆分为多个分片，以便每个分片都小于给定的大小。
分片是通过按其键的顺序迭代 state_dict 来确定的。没有进行优化以使每个分片尽可能接近传递的最大大小。例如，如果限制为 10GB，并且我们有大小为 [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] 的张量，它们将被分片为 [6GB], [6+2GB], [6+2+2GB]，而不是 [6+2+2GB], [6+2GB], [6GB]。
要将模型状态字典保存到磁盘，请参阅 save_torch_state_dict()。此辅助函数在底层使用了 split_torch_state_dict_into_shards。
如果模型的某个张量大于 max_shard_size，它最终将位于其自己的分片中，该分片的大小将大于 max_shard_size。
示例
>>> import json
>>> import os
>>> from safetensors.torch import save_file as safe_save_file
>>> from huggingface_hub import split_torch_state_dict_into_shards
>>> def save_state_dict(state_dict: Dict[str, torch.Tensor], save_directory: str):
... state_dict_split = split_torch_state_dict_into_shards(state_dict)
... for filename, tensors in state_dict_split.filename_to_tensors.items():
... shard = {tensor: state_dict[tensor] for tensor in tensors}
... safe_save_file(
... shard,
... os.path.join(save_directory, filename),
... metadata={"format": "pt"},
... )
... if state_dict_split.is_sharded:
... index = {
... "metadata": state_dict_split.metadata,
... "weight_map": state_dict_split.tensor_to_filename,
... }
... with open(os.path.join(save_directory, "model.safetensors.index.json"), "w") as f:
... f.write(json.dumps(index, indent=2))split_state_dict_into_shards_factory
这是每个框架特定辅助函数派生的底层工厂。在实践中，除非您需要将其适配到尚不支持的框架，否则您不应该直接使用此工厂。如果是这种情况，请通过在 huggingface_hub 存储库上 新建一个问题 告知我们。
huggingface_hub.split_state_dict_into_shards_factory
< 源代码 >( state_dict: Dict get_storage_size: Callable filename_pattern: str get_storage_id: Callable = <function <lambda> at 0x7fbffa9d93f0> max_shard_size: Union = '5GB' ) → StateDictSplit
参数
- state_dict (
Dict[str, Tensor]) — 要保存的状态字典。 - get_storage_size (
Callable[[Tensor], int]) — 一个函数，返回保存到磁盘上的张量的大小（以字节为单位）。 - get_storage_id (
Callable[[Tensor], Optional[Any]], 可选) — 返回张量存储唯一标识符的函数。多个不同的张量可以共享同一个底层存储。此标识符在其生命周期内保证对该张量的存储是唯一的且恒定的。生命周期不重叠的两个张量存储可能具有相同的 ID。 - filename_pattern (
str, 可选) — 用于生成模型保存文件名的模式。模式必须是一个可以用filename_pattern.format(suffix=...)格式化的字符串，并且必须包含关键字suffix - max_shard_size (
int或str, 可选) — 每个分片的最大大小，以字节为单位。默认为 5GB。
返回值
StateDictSplit
一个包含分片和检索它们索引的 StateDictSplit 对象。
将模型状态字典拆分为多个分片，以便每个分片都小于给定的大小。
分片是通过按其键的顺序迭代 state_dict 来确定的。没有进行优化以使每个分片尽可能接近传递的最大大小。例如，如果限制为 10GB，并且我们有大小为 [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] 的张量，它们将被分片为 [6GB], [6+2GB], [6+2+2GB]，而不是 [6+2+2GB], [6+2GB], [6GB]。
如果模型的某个张量大于 max_shard_size，它最终将位于其自己的分片中，该分片的大小将大于 max_shard_size。
辅助函数
get_torch_storage_id
返回张量存储的唯一标识符。
多个不同的张量可以共享同一个底层存储。此标识符在其生命周期内保证对该张量的存储是唯一的且恒定的。生命周期不重叠的两个张量存储可能具有相同的 ID。 对于元张量，我们返回 None，因为我们无法判断它们是否共享同一个存储。