Hub Python Library documentation

序列化

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

序列化

huggingface_hub 提供了一系列辅助函数,用于以标准化的方式保存和加载机器学习模型权重。该库的这一部分仍在开发中,将在未来的版本中进一步完善。其目标是在 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: str | os.PathLikeentries: Iterable )

参数

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

huggingface_hub.export_folder_as_dduf

< >

( dduf_path: str | os.PathLikefolder_path: str | os.PathLike )

参数

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

将文件夹导出为 DDUF 文件。

A在底层使用 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: 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: strlength: intoffset: intdduf_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())

错误

类 huggingface_hub.errors.DDUFError

< >

( )

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

类 huggingface_hub.errors.DDUFCorruptedFileError

< >

( )

当 DDUF 文件损坏时抛出异常。

类 huggingface_hub.errors.DDUFExportError

< >

( )

DDUF 导出过程中发生错误时的基础异常。

类 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.Modulesave_directory: str | pathlib.Pathfilename_pattern: str | None = Noneforce_contiguous: bool = Truemax_shard_size: int | str = '5GB'metadata: dict[str, str] | None = Nonesafe_serialization: bool = Trueis_main_process: bool = Trueshared_tensors_to_discard: list[str] | None = 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(),以保存具有更高灵活性的状态字典。

有关张量共享的更多信息,请参阅此指南。

模型状态字典会被拆分为多个分片,以确保每个分片的大小小于指定值。这些分片会以给定的 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。

如果您的模型是 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: dictsave_directory: str | pathlib.Pathfilename_pattern: str | None = Noneforce_contiguous: bool = Truemax_shard_size: int | str = '5GB'metadata: dict[str, str] | None = Nonesafe_serialization: bool = Trueis_main_process: bool = Trueshared_tensors_to_discard: list[str] | None = 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, 可选) — 强制将 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], 可选) — 保存共享张量时要丢弃的张量名称列表。如果未提供且检测到共享张量,则会按字母顺序丢弃第一个名称。

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

另请参阅 save_torch(),用于直接保存 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。

如果您的模型是 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 张量,并且设计为易于扩展到其他任何机器学习框架。

split_torch_state_dict_into_shards

huggingface_hub.split_torch_state_dict_into_shards

< >

( state_dict: dictfilename_pattern: str = 'model{suffix}.safetensors'max_shard_size: int | str = '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 仓库中通过提交一个新 issue 告知我们。

huggingface_hub.split_state_dict_into_shards_factory

< >

( state_dict: dictget_storage_size: Callablefilename_pattern: strget_storage_id: Callable = <function <lambda> at 0x7f9673968670>max_shard_size: int | str = '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。

加载张量

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

这两个辅助函数默认均为 safe=True,这意味着无论文件名称如何,检查点始终通过 safetensors 加载器进行反序列化。加载 pickle 检查点需要显式选择,因为反序列化 pickle 会在加载时执行任意代码:

>>> from huggingface_hub import load_state_dict_from_file, load_torch_model

# Raises: the file is not safetensors
>>> load_state_dict_from_file("path/to/pytorch_model.bin")
>>> load_torch_model(model, "path/to/pytorch_model.bin")

# Explicit opt-in
>>> load_state_dict_from_file("path/to/pytorch_model.bin", safe=False)
>>> load_torch_model(model, "path/to/pytorch_model.bin", safe=False)

pickle 路径还默认使用 weights_only=True,即 torch 的受限反序列化器。如果检查点中包含非张量对象(优化器状态、整个模型等),则需要在 safe=False 的基础上额外设置 weights_only=False。

load_torch_model

huggingface_hub.load_torch_model

< >

( model: torch.nn.Modulecheckpoint_path: str | os.PathLikestrict: bool = Falsesafe: bool = Trueweights_only: bool = Truemap_location: typing.Union[str, ForwardRef('torch.device'), NoneType] = Nonemmap: bool = Falsefilename_pattern: str | None = None ) → NamedTuple

参数

  • model (torch.nn.Module) — 用于加载检查点的模型。
  • checkpoint_path (str 或 os.PathLike) — 检查点文件或包含检查点的目录的路径。
  • strict (bool,可选,默认为 False) — 是否严格强制要求模型状态字典中的键与检查点中的键相匹配。 与 torch.nn.Module.load_state_dict 一样,该检查在检查点加载完成后执行,因此 在抛出错误时,模型可能已被部分更新。
  • safe (bool,可选,默认为 True) — 如果 safe 为 True,则会加载 safetensors 文件。如果 safe 为 False,该函数会首先尝试加载可用的 safetensors 文件,否则回退到加载 pickle 文件。当 safe=True 时,任何未描述 safetensors 文件的 filename_pattern 都会被拒绝。
  • weights_only (bool,可选,默认为 True) — 如果为 True,则仅加载模型权重,不加载优化器状态和其他元数据,并使用 torch 的受限 unpickler。设置为 False 以允许 pickle 检查点中包含任意 Python 对象(这会在加载时执行任意代码)。对 PyTorch < 1.13 没有影响,因为该版本没有受限 unpickler,并且始终无限制地执行 unpickle。
  • 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 是一个包含缺失键的 str 列表,即模型中存在但检查点中不存在的键。
  • unexpected_keys 是一个包含意外键的 str 列表,即存在于检查点中但不在模型中的键。

引发

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: str | os.PathLikemap_location: typing.Union[str, ForwardRef('torch.device'), NoneType] = Noneweights_only: bool = Truemmap: bool = Falsesafe: bool = True ) → Union[dict[str, "torch.Tensor"], Any]

参数

  • checkpoint_file (str 或 os.PathLike) — 要加载的检查点文件的路径。可以是 safetensors 或 pickle (.bin) 检查点。
  • map_location (str 或 torch.device,可选) — torch.device 对象、字符串或字典,用于指定如何重新映射存储位置。它指示所有张量应加载到的位置。
  • weights_only (bool,可选,默认为 True) — 如果为 True,则仅加载模型权重,不加载优化器状态和其他元数据,并使用 torch 的受限反序列化器。设置为 False 以允许 pickle 检查点中包含任意 Python 对象(这会在加载时执行任意代码)。在加载 safetensors 文件时没有效果,对于没有受限反序列化器的 PyTorch < 1.13 也没有效果——这些版本始终不受限制地进行反序列化,并会记录一条警告。
  • mmap (bool,可选,默认为 False) — 是否使用内存映射文件加载。对于 PyTorch >= 2.1.0 中基于 zipfile 的检查点,内存映射可以提升大模型的加载性能。在加载 safetensors 文件时此参数无效,因为 safetensors 库默认使用内存映射。
  • safe (bool,可选,默认为 True) — 如果为 True,无论检查点文件名是什么,都会始终将其作为 safetensors 加载。任何其他格式(例如 pickle .bin 文件)都会引发 ValueError,而不是被反序列化:pickle 检查点在加载时可能会执行任意代码。如果为 False,则当文件名表明其为 safetensors 格式时按该格式加载,否则使用 torch.load 加载。

返回值

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

加载的检查点。

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

引发

FileNotFoundError 或 ImportError 或 OSError 或 ValueError

  • FileNotFoundError — 如果检查点文件不存在。
  • ImportError — 当尝试加载 .safetensors 文件或 PyTorch 检查点时,如果分别未安装 safetensors 或 torch,则会抛出此错误。
  • OSError — 如果检查点文件格式无效,或者 git-lfs 文件未正确下载。
  • ValueError — 如果检查点文件路径为空或无效,或者在 safe=True 时无法将其反序列化为 safetensors(传入 safe=False 以允许使用 pickle 检查点)。

加载检查点文件,同时支持 safetensors 和 pickle 检查点格式。

示例:

>>> from huggingface_hub import load_state_dict_from_file

# Load a safetensors checkpoint (safe by default)
>>> state_dict = load_state_dict_from_file("path/to/model.safetensors", safe=True)
>>> model.load_state_dict(state_dict)

# Load a pickle checkpoint. `safe=False` is required: pickle files can execute arbitrary code.
>>> state_dict = load_state_dict_from_file("path/to/model.bin", safe=False, map_location="cpu")
>>> model.load_state_dict(state_dict)

张量辅助

get_torch_storage_id

huggingface_hub.get_torch_storage_id

< >

( tensor: torch.Tensor )

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

多个不同的张量可以共享同一底层存储。此标识符保证在该张量存储的生命周期内是唯一的且保持不变的。两个生命周期不重叠的张量存储可能具有相同的 id。 对于 meta 张量,由于无法判断它们是否共享同一存储,我们返回 None。

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

get_torch_storage_size

Update on GitHub