Hugging Face's logo

Hub Python 库文档

序列化

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

序列化

huggingface_hub 提供了以标准化方式保存和加载 ML 模型权重的辅助工具。库的这一部分仍在开发中，将在未来的版本中得到改进。目标是协调 Hub 中权重保存和加载的方式，既要消除跨库的代码重复，又要建立一致的约定。

DDUF 文件格式

DDUF 是一种为扩散模型设计的文件格式。它允许将运行模型的所有信息保存在单个文件中。这项工作受到 GGUF 格式的启发。huggingface_hub 提供了保存和加载 DDUF 文件的辅助工具，确保文件格式得到遵守。

这是解析器的早期版本。API 和实现可能会在不久的将来发生变化。

解析器目前只进行非常少的验证。有关文件格式的更多详细信息，请查看 https://github.com/huggingface/huggingface.js/tree/main/packages/dduf。

如何编写 DDUF 文件？

以下是如何使用 export_folder_as_dduf() 导出包含扩散模型不同部分的文件夹：

# Export a folder as a DDUF file
>>> from huggingface_hub import export_folder_as_dduf
>>> export_folder_as_dduf("FLUX.1-dev.dduf", folder_path="path/to/FLUX.1-dev")

为了获得更大的灵活性，您可以使用 export_entries_as_dduf() 并传递要包含在最终 DDUF 文件中的文件列表

# Export specific files from the local disk.
>>> from huggingface_hub import export_entries_as_dduf
>>> export_entries_as_dduf(
...     dduf_path="stable-diffusion-v1-4-FP16.dduf",
...     entries=[ # List entries to add to the DDUF file (here, only FP16 weights)
...         ("model_index.json", "path/to/model_index.json"),
...         ("vae/config.json", "path/to/vae/config.json"),
...         ("vae/diffusion_pytorch_model.fp16.safetensors", "path/to/vae/diffusion_pytorch_model.fp16.safetensors"),
...         ("text_encoder/config.json", "path/to/text_encoder/config.json"),
...         ("text_encoder/model.fp16.safetensors", "path/to/text_encoder/model.fp16.safetensors"),
...         # ... add more entries here
...     ]
... )

entries 参数还支持传递路径或字节的可迭代对象。如果您有一个已加载的模型，并希望直接将其序列化为 DDUF 文件，而不是必须先将每个组件序列化到磁盘，然后再序列化为 DDUF 文件，这将非常有用。以下是如何将 StableDiffusionPipeline 序列化为 DDUF 的示例

# Export state_dicts one by one from a loaded pipeline 
>>> from diffusers import DiffusionPipeline
>>> from typing import Generator, Tuple
>>> import safetensors.torch
>>> from huggingface_hub import export_entries_as_dduf
>>> pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
... # ... do some work with the pipeline

>>> def as_entries(pipe: DiffusionPipeline) -> Generator[Tuple[str, bytes], None, None]:
...     # Build a generator that yields the entries to add to the DDUF file.
...     # The first element of the tuple is the filename in the DDUF archive (must use UNIX separator!). The second element is the content of the file.
...     # Entries will be evaluated lazily when the DDUF file is created (only 1 entry is loaded in memory at a time)
...     yield "vae/config.json", pipe.vae.to_json_string().encode()
...     yield "vae/diffusion_pytorch_model.safetensors", safetensors.torch.save(pipe.vae.state_dict())
...     yield "text_encoder/config.json", pipe.text_encoder.config.to_json_string().encode()
...     yield "text_encoder/model.safetensors", safetensors.torch.save(pipe.text_encoder.state_dict())
...     # ... add more entries here

>>> export_entries_as_dduf(dduf_path="stable-diffusion-v1-4.dduf", entries=as_entries(pipe))

注意： 实际上，diffusers 提供了一种直接将 pipeline 序列化为 DDUF 文件的方法。上面的代码片段仅用作示例。

如何读取 DDUF 文件？

>>> import json
>>> import safetensors.torch
>>> from huggingface_hub import read_dduf_file

# Read DDUF metadata
>>> dduf_entries = read_dduf_file("FLUX.1-dev.dduf")

# Returns a mapping filename <> DDUFEntry
>>> dduf_entries["model_index.json"]
DDUFEntry(filename='model_index.json', offset=66, length=587)

# Load model index as JSON
>>> json.loads(dduf_entries["model_index.json"].read_text())
{'_class_name': 'FluxPipeline', '_diffusers_version': '0.32.0.dev0', '_name_or_path': 'black-forest-labs/FLUX.1-dev', 'scheduler': ['diffusers', 'FlowMatchEulerDiscreteScheduler'], 'text_encoder': ['transformers', 'CLIPTextModel'], 'text_encoder_2': ['transformers', 'T5EncoderModel'], 'tokenizer': ['transformers', 'CLIPTokenizer'], 'tokenizer_2': ['transformers', 'T5TokenizerFast'], 'transformer': ['diffusers', 'FluxTransformer2DModel'], 'vae': ['diffusers', 'AutoencoderKL']}

# Load VAE weights using safetensors
>>> with dduf_entries["vae/diffusion_pytorch_model.safetensors"].as_mmap() as mm:
...     state_dict = safetensors.torch.load(mm)

辅助函数

huggingface_hub.export_entries_as_dduf

< >

( dduf_path: typing.Union[str, os.PathLike] entries: typing.Iterable[typing.Tuple[str, typing.Union[str, pathlib.Path, bytes]]] )

参数

  • dduf_path (str 或 os.PathLike) — 要写入的 DDUF 文件的路径。
  • entries (Iterable[Tuple[str, Union[str, Path, bytes]]]) — 要写入 DDUF 文件的条目的可迭代对象。每个条目都是一个元组，包含文件名和内容。文件名应该是 DDUF 存档中文件的路径。内容可以是字符串或 pathlib.Path，表示本地磁盘上文件的路径，也可以直接是字节内容。

引发

    • — DDUFExportError：如果在导出过程中发生任何错误（例如，条目名称无效，缺少 'model_index.json' 等）。

从条目的可迭代对象写入 DDUF 文件。

这是一个比 export_folder_as_dduf() 更低级别的辅助函数，在序列化数据时允许更大的灵活性。特别是，您不需要在将数据导出到 DDUF 文件之前将其保存在磁盘上。

示例

# Export specific files from the local disk.
>>> from huggingface_hub import export_entries_as_dduf
>>> export_entries_as_dduf(
...     dduf_path="stable-diffusion-v1-4-FP16.dduf",
...     entries=[ # List entries to add to the DDUF file (here, only FP16 weights)
...         ("model_index.json", "path/to/model_index.json"),
...         ("vae/config.json", "path/to/vae/config.json"),
...         ("vae/diffusion_pytorch_model.fp16.safetensors", "path/to/vae/diffusion_pytorch_model.fp16.safetensors"),
...         ("text_encoder/config.json", "path/to/text_encoder/config.json"),
...         ("text_encoder/model.fp16.safetensors", "path/to/text_encoder/model.fp16.safetensors"),
...         # ... add more entries here
...     ]
... )
# Export state_dicts one by one from a loaded pipeline
>>> from diffusers import DiffusionPipeline
>>> from typing import Generator, Tuple
>>> import safetensors.torch
>>> from huggingface_hub import export_entries_as_dduf
>>> pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
... # ... do some work with the pipeline

>>> def as_entries(pipe: DiffusionPipeline) -> Generator[Tuple[str, bytes], None, None]:
...     # Build an generator that yields the entries to add to the DDUF file.
...     # The first element of the tuple is the filename in the DDUF archive (must use UNIX separator!). The second element is the content of the file.
...     # Entries will be evaluated lazily when the DDUF file is created (only 1 entry is loaded in memory at a time)
...     yield "vae/config.json", pipe.vae.to_json_string().encode()
...     yield "vae/diffusion_pytorch_model.safetensors", safetensors.torch.save(pipe.vae.state_dict())
...     yield "text_encoder/config.json", pipe.text_encoder.config.to_json_string().encode()
...     yield "text_encoder/model.safetensors", safetensors.torch.save(pipe.text_encoder.state_dict())
...     # ... add more entries here

>>> export_entries_as_dduf(dduf_path="stable-diffusion-v1-4.dduf", entries=as_entries(pipe))

huggingface_hub.export_folder_as_dduf

< >

( dduf_path: typing.Union[str, os.PathLike] folder_path: typing.Union[str, os.PathLike] )

参数

  • dduf_path (str 或 os.PathLike) — 要写入的 DDUF 文件的路径。
  • folder_path (str 或 os.PathLike) — 包含扩散模型的文件夹的路径。

将文件夹导出为 DDUF 文件。

使用 export_entries_as_dduf() 作为底层函数。

示例

>>> from huggingface_hub import export_folder_as_dduf
>>> export_folder_as_dduf(dduf_path="FLUX.1-dev.dduf", folder_path="path/to/FLUX.1-dev")

huggingface_hub.read_dduf_file

< >

( dduf_path: typing.Union[os.PathLike, str] ) → Dict[str, DDUFEntry]

参数

  • dduf_path (str 或 os.PathLike) — 要读取的 DDUF 文件的路径。

返回值

Dict[str, DDUFEntry]

一个 DDUFEntry 字典，按文件名索引。

引发

    • — DDUFCorruptedFileError：如果 DDUF 文件已损坏（即不符合 DDUF 格式）。

读取 DDUF 文件并返回条目字典。

仅读取元数据，数据不会加载到内存中。

示例

>>> import json
>>> import safetensors.torch
>>> from huggingface_hub import read_dduf_file

# Read DDUF metadata
>>> dduf_entries = read_dduf_file("FLUX.1-dev.dduf")

# Returns a mapping filename <> DDUFEntry
>>> dduf_entries["model_index.json"]
DDUFEntry(filename='model_index.json', offset=66, length=587)

# Load model index as JSON
>>> json.loads(dduf_entries["model_index.json"].read_text())
{'_class_name': 'FluxPipeline', '_diffusers_version': '0.32.0.dev0', '_name_or_path': 'black-forest-labs/FLUX.1-dev', ...

# Load VAE weights using safetensors
>>> with dduf_entries["vae/diffusion_pytorch_model.safetensors"].as_mmap() as mm:
...     state_dict = safetensors.torch.load(mm)

class huggingface_hub.DDUFEntry

< >

( filename: str length: int offset: int dduf_path: Path )

参数

  • filename (str) — DDUF 存档中文件的名称。
  • offset (int) — 文件在 DDUF 存档中的偏移量。
  • length (int) — 文件在 DDUF 存档中的长度。
  • dduf_path (str) — DDUF 存档的路径（供内部使用）。

表示 DDUF 文件中文件条目的对象。

有关如何读取 DDUF 文件，请参阅 read_dduf_file()。

as_mmap

< >

( )

将文件作为内存映射文件打开。

用于直接从文件加载 safetensors。

示例

>>> import safetensors.torch
>>> with entry.as_mmap() as mm:
...     tensors = safetensors.torch.load(mm)

read_text

< >

( encoding: str = 'utf-8' )

将文件读取为文本。

适用于 ‘.txt’ 和 ‘.json’ 条目。

示例

>>> import json
>>> index = json.loads(entry.read_text())

错误

class huggingface_hub.errors.DDUFError

< >

( )

与 DDUF 格式相关的错误的基本异常。

class huggingface_hub.errors.DDUFCorruptedFileError

< >

( )

当 DDUF 文件损坏时引发的异常。

class huggingface_hub.errors.DDUFExportError

< >

( )

DDUF 导出期间发生错误的基本异常。

class huggingface_hub.errors.DDUFInvalidEntryNameError

< >

( )

当条目名称无效时引发的异常。

保存张量

serialization 模块的主要助手接受 torch nn.Module 作为输入，并将其保存到磁盘。它处理保存共享张量的逻辑（参见safetensors 说明）以及将状态字典拆分为分片的逻辑，在底层使用 split_torch_state_dict_into_shards()。目前，仅支持 torch 框架。

如果您想保存状态字典（例如，层名称和相关张量之间的映射）而不是 nn.Module，您可以使用 save_torch_state_dict()，它提供相同的功能。例如，如果您想在保存状态字典之前应用自定义逻辑，这将非常有用。

save_torch_model

huggingface_hub.save_torch_model

< >

( model: torch.nn.Module save_directory: typing.Union[str, pathlib.Path] filename_pattern: typing.Optional[str] = None force_contiguous: bool = True max_shard_size: typing.Union[int, str] = '5GB' metadata: typing.Optional[typing.Dict[str, str]] = None safe_serialization: bool = True is_main_process: bool = True shared_tensors_to_discard: typing.Optional[typing.List[str]] = None )

参数

  • 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，并且将在未来版本中删除。
  • is_main_process (bool, 可选) — 调用此函数的进程是否为主进程。在分布式训练（如 TPU）中并且需要从所有进程调用此函数时非常有用。在这种情况下，仅在主进程上设置 is_main_process=True 以避免竞争条件。默认为 True。
  • shared_tensors_to_discard (List[str], 可选) — 保存共享张量时要删除的张量名称列表。如果未提供且检测到共享张量，它将删除按字母顺序排列的第一个名称。

将给定的 torch 模型保存到磁盘，处理分片和共享张量问题。

另请参阅 save_torch_state_dict() 以更灵活地保存状态字典。

有关张量共享的更多信息，请查看本指南。

模型状态字典被拆分为分片，以便每个分片都小于给定大小。分片保存在 save_directory 中，并使用给定的 filename_pattern。如果模型太大而无法放入单个分片，则会在 save_directory 中保存一个索引文件，以指示每个张量保存在哪里。此助手在底层使用 split_torch_state_dict_into_shards()。如果 safe_serialization 为 True，则分片将保存为 safetensors（默认）。否则，分片将保存为 pickle。

在保存模型之前，save_directory 将被清除任何先前的分片文件。

如果模型的某个张量大于 max_shard_size，它将最终出现在自己的分片中，该分片的大小将大于 max_shard_size。

如果您的模型是 transformers.PreTrainedModel，则应将 model._tied_weights_keys 作为 shared_tensors_to_discard 传递，以正确处理共享张量的保存。这确保在保存期间丢弃正确的重复张量。

示例

>>> 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")
>>>

save_torch_state_dict

huggingface_hub.save_torch_state_dict

< >

( state_dict: typing.Dict[str, ForwardRef('torch.Tensor')] save_directory: typing.Union[str, pathlib.Path] filename_pattern: typing.Optional[str] = None force_contiguous: bool = True max_shard_size: typing.Union[int, str] = '5GB' metadata: typing.Optional[typing.Dict[str, str]] = None safe_serialization: bool = True is_main_process: bool = True shared_tensors_to_discard: typing.Optional[typing.List[str]] = None )

参数

  • 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, optional) — 强制将 state_dict 保存为连续张量。这对模型的正确性没有影响，但如果张量的布局是专门为此目的而选择的，则可能会改变性能。默认为 True。
  • max_shard_size (int 或 str, optional) — 每个分片的最大大小，以字节为单位。默认为 5GB。
  • metadata (Dict[str, str], optional) — 与模型一起保存的额外信息。将为每个丢弃的张量添加一些元数据。这些信息不足以恢复整个共享结构，但可能有助于理解事物。
  • safe_serialization (bool, optional) — 是否保存为 safetensors 格式，这是默认行为。如果为 False，则分片保存为 pickle 格式。出于安全原因，建议使用安全序列化。以 pickle 格式保存已被弃用，并将在未来版本中删除。
  • is_main_process (bool, optional) — 调用此函数的进程是否为主进程。在分布式训练（如 TPU）中，需要从所有进程调用此函数时很有用。在这种情况下，仅在主进程上设置 is_main_process=True 以避免竞争条件。默认为 True。
  • shared_tensors_to_discard (List[str], optional) — 保存共享张量时要丢弃的张量名称列表。如果未提供且检测到共享张量，则将按字母顺序丢弃第一个名称。

将模型状态字典保存到磁盘，处理分片和共享张量问题。

另请参阅 save_torch_model() 以直接保存 PyTorch 模型。

有关张量共享的更多信息，请查看本指南。

模型状态字典被拆分为分片，以便每个分片都小于给定大小。分片保存在 save_directory 中，并使用给定的 filename_pattern。如果模型太大而无法放入单个分片，则会在 save_directory 中保存一个索引文件，以指示每个张量保存在哪里。此助手在底层使用 split_torch_state_dict_into_shards()。如果 safe_serialization 为 True，则分片将保存为 safetensors（默认）。否则，分片将保存为 pickle。

在保存模型之前，save_directory 将被清除任何先前的分片文件。

如果模型的某个张量大于 max_shard_size，它将最终出现在自己的分片中，该分片的大小将大于 max_shard_size。

如果您的模型是 transformers.PreTrainedModel，则应将 model._tied_weights_keys 作为 shared_tensors_to_discard 传递，以正确处理共享张量的保存。这确保在保存期间丢弃正确的重复张量。

示例

>>> 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 张量，并且设计为易于扩展到任何其他 ML 框架。

split_tf_state_dict_into_shards

huggingface_hub.split_tf_state_dict_into_shards

< >

( state_dict: typing.Dict[str, ForwardRef('tf.Tensor')] filename_pattern: str = 'tf_model{suffix}.h5' max_shard_size: typing.Union[int, str] = '5GB' ) → StateDictSplit

参数

  • state_dict (Dict[str, Tensor]) — 要保存的状态字典。
  • filename_pattern (str, optional) — 生成用于保存模型的文件名的模式。模式必须是可以使用 filename_pattern.format(suffix=...) 格式化的字符串，并且必须包含关键字 suffix 默认为 "tf_model{suffix}.h5"。
  • max_shard_size (int 或 str, optional) — 每个分片的最大大小，以字节为单位。默认为 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: typing.Dict[str, ForwardRef('torch.Tensor')] filename_pattern: str = 'model{suffix}.safetensors' max_shard_size: typing.Union[int, str] = '5GB' ) → StateDictSplit

参数

  • state_dict (Dict[str, torch.Tensor]) — 要保存的状态字典。
  • filename_pattern (str, optional) — 生成用于保存模型的文件名的模式。模式必须是可以使用 filename_pattern.format(suffix=...) 格式化的字符串，并且必须包含关键字 suffix 默认为 "model{suffix}.safetensors"。
  • max_shard_size (int 或 str, optional) — 每个分片的最大大小，以字节为单位。默认为 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 仓库上打开新 issue 告知我们。

huggingface_hub.split_state_dict_into_shards_factory

< >

( state_dict: typing.Dict[str, ~TensorT] get_storage_size: typing.Callable[[~TensorT], int] filename_pattern: str get_storage_id: typing.Callable[[~TensorT], typing.Optional[typing.Any]] = <function <lambda> at 0x7f61cb4364d0> max_shard_size: typing.Union[int, str] = '5GB' ) → StateDictSplit

参数

  • state_dict (Dict[str, Tensor]) — 要保存的状态字典。
  • get_storage_size (Callable[[Tensor], int]) — 一个函数，用于返回张量在磁盘上保存时的大小（以字节为单位）。
  • get_storage_id (Callable[[Tensor], Optional[Any]], optional) — 一个函数，用于返回张量存储的唯一标识符。多个不同的张量可以共享相同的底层存储。此标识符保证在此张量的存储的生命周期内是唯一且恒定的。生命周期不重叠的两个张量存储可能具有相同的 ID。
  • filename_pattern (str, optional) — 生成用于保存模型的文件名的模式。模式必须是可以使用 filename_pattern.format(suffix=...) 格式化的字符串，并且必须包含关键字 suffix
  • max_shard_size (int 或 str, optional) — 每个分片的最大大小，以字节为单位。默认为 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。

加载张量

加载助手支持 safetensors 或 pickle 格式的单文件和分片检查点。load_torch_model() 接受一个 nn.Module 和一个检查点路径（可以是单个文件或目录）作为输入，并将权重加载到模型中。

load_torch_model

huggingface_hub.load_torch_model

< >

( model: torch.nn.Module checkpoint_path: typing.Union[str, os.PathLike] strict: bool = False safe: bool = True weights_only: bool = False map_location: typing.Union[str, ForwardRef('torch.device'), NoneType] = None mmap: bool = False filename_pattern: typing.Optional[str] = None ) → NamedTuple

参数

  • model (torch.nn.Module) — The model in which to load the checkpoint.
  • checkpoint_path (str 或 os.PathLike) — 指向检查点文件或包含检查点目录的路径。
  • strict (bool, 可选, 默认为 False) — 是否严格执行模型状态字典中的键与检查点中的键匹配。
  • safe (bool, 可选, 默认为 True) — 如果 safe 为 True，将加载 safetensors 文件。如果 safe 为 False，该函数将首先尝试加载 safetensors 文件（如果可用），否则将回退到加载 pickle 文件。filename_pattern 参数优先于 safe 参数。
  • weights_only (bool, 可选, 默认为 False) — 如果为 True，则仅加载模型权重，而不加载优化器状态和其他元数据。仅在 PyTorch >= 1.13 中受支持。
  • map_location (str 或 torch.device, 可选) — 一个 torch.device 对象、字符串或字典，用于指定如何重新映射存储位置。它指示所有张量应加载到的位置。
  • mmap (bool, 可选, 默认为 False) — 是否使用内存映射文件加载。对于 PyTorch >= 2.1.0 中基于 zipfile 的检查点，内存映射可以提高大型模型的加载性能。
  • filename_pattern (str, 可选) — 用于查找索引文件的模式。模式必须是一个字符串，可以使用 filename_pattern.format(suffix=...) 格式化，并且必须包含关键字 suffix。默认为 "model{suffix}.safetensors"。

返回值

NamedTuple

一个具名元组，包含 missing_keys 和 unexpected_keys 字段。

  • missing_keys 是一个字符串列表，包含缺失的键，即模型中存在但在检查点中不存在的键。
  • unexpected_keys 是一个字符串列表，包含未预期的键，即检查点中存在但在模型中不存在的键。

引发

FileNotFoundError 或 ImportError 或 ValueError

  • FileNotFoundError — 如果检查点文件或目录不存在。
  • ImportError — 如果尝试加载 .safetensors 文件或 PyTorch 检查点时，未安装 safetensors 或 torch。
  • ValueError — 如果检查点路径无效，或者无法确定检查点格式。

将检查点加载到模型中，处理分片和非分片检查点。

示例

>>> from huggingface_hub import load_torch_model
>>> model = ... # A PyTorch model
>>> load_torch_model(model, "path/to/checkpoint")

load_state_dict_from_file

huggingface_hub.load_state_dict_from_file

< >

( checkpoint_file: typing.Union[str, os.PathLike] map_location: typing.Union[str, ForwardRef('torch.device'), NoneType] = None weights_only: bool = False mmap: bool = False ) → Union[Dict[str, "torch.Tensor"], Any]

参数

  • checkpoint_file (str 或 os.PathLike) — 要加载的检查点文件的路径。可以是 safetensors 或 pickle (.bin) 检查点。
  • map_location (str 或 torch.device, 可选) — 一个 torch.device 对象、字符串或字典，用于指定如何重新映射存储位置。它指示所有张量应加载到的位置。
  • weights_only (bool, 可选, 默认为 False) — 如果为 True，则仅加载模型权重，而不加载优化器状态和其他元数据。仅适用于 PyTorch >= 1.13 的 pickle (.bin) 检查点。加载 safetensors 文件时无效。
  • mmap (bool, 可选, 默认为 False) — 是否使用内存映射文件加载。对于 PyTorch >= 2.1.0 中基于 zipfile 的检查点，内存映射可以提高大型模型的加载性能。加载 safetensors 文件时无效，因为 safetensors 库默认使用内存映射。

返回值

Union[Dict[str, "torch.Tensor"], Any]

加载的检查点。

  • 对于 safetensors 文件：始终返回一个字典，将参数名称映射到张量。
  • 对于 pickle 文件：返回任何被 pickle 化的 Python 对象（通常是状态字典，但也可能是整个模型、优化器状态或任何其他 Python 对象）。

引发

FileNotFoundError 或 ImportError 或 OSError 或 ValueError

  • FileNotFoundError — 如果检查点文件不存在。
  • ImportError — 如果尝试加载 .safetensors 文件或 PyTorch 检查点时，未安装 safetensors 或 torch。
  • OSError — 如果检查点文件格式无效，或者 git-lfs 文件未正确下载。
  • ValueError — 如果检查点文件路径为空或无效。

加载检查点文件，处理 safetensors 和 pickle 检查点格式。

示例

>>> from huggingface_hub import load_state_dict_from_file

# Load a PyTorch checkpoint
>>> state_dict = load_state_dict_from_file("path/to/model.bin", map_location="cpu")
>>> model.load_state_dict(state_dict)

# Load a safetensors checkpoint
>>> state_dict = load_state_dict_from_file("path/to/model.safetensors")
>>> model.load_state_dict(state_dict)

张量助手函数

get_torch_storage_id

huggingface_hub.get_torch_storage_id

< >

( tensor: torch.Tensor )

返回张量存储的唯一标识符。

多个不同的张量可以共享相同的底层存储。此标识符保证在此张量的存储在其生命周期内是唯一且恒定的。具有不重叠生命周期的两个张量存储可能具有相同的 ID。对于元张量，我们返回 None，因为我们无法判断它们是否共享相同的存储。

摘自 https://github.com/huggingface/transformers/blob/1ecf5f7c982d761b4daaa96719d162c324187c64/src/transformers/pytorch_utils.py#L278。

get_torch_storage_size

huggingface_hub.get_torch_storage_size

< >

( tensor: torch.Tensor )

摘自 https://github.com/huggingface/safetensors/blob/08db34094e9e59e2f9218f2df133b7b4aaff5a99/bindings/python/py_src/safetensors/torch.py#L31C1-L41C59

< > 更新 在 GitHub 上

←Webhooks 服务器