# 序列化

`huggingface_hub` 提供了一系列辅助函数，用于以标准化的方式保存和加载机器学习模型权重。该库的这一部分仍在开发中，将在未来的版本中进一步完善。其目标是在 Hub 上统一权重的保存和加载方式，既消除各库之间的代码重复，又建立一致的规范。

## DDUF 文件格式

DDUF 是一种专为扩散模型设计的文件格式。它允许将运行模型所需的所有信息保存在单个文件中。这项工作受到 [GGUF](https://github.com/ggerganov/ggml/blob/master/docs/gguf.md) 格式的启发。`huggingface_hub` 提供了用于保存和加载 DDUF 文件的辅助工具，以确保文件格式得到遵守。

> [!WARNING]
> 这是解析器的一个非常早期的版本。其 API 和实现可能会在不久的将来发生变化。
>
> 解析器目前几乎不做任何验证。有关文件格式的更多详细信息，请访问 https://github.com/huggingface/huggingface.js/tree/main/packages/dduf。

### 如何编写 DDUF 文件？

以下是如何使用 [export_folder_as_dduf()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.export_folder_as_dduf) 导出包含扩散模型不同部分的文件夹：

```python
# 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()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.export_entries_as_dduf) 并传入一个文件列表，以将其包含在最终的 DDUF 文件中：

```python
# 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 的示例：

```python
# 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 文件？

```python
>>> 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]]

#### huggingface_hub.export_entries_as_dduf[[huggingface_hub.export_entries_as_dduf]]

```python
huggingface_hub.export_entries_as_dduf(dduf_path: str | os.PathLike, entries: Iterable)
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_dduf.py#L160)

**参数：**

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()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.export_folder_as_dduf) 更底层的辅助函数，在序列化数据时提供了更高的灵活性。特别是，您无需先将数据保存到磁盘，即可将其导出到 DDUF 文件中。

示例：
```python
# 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
...     ]
... )
```

```python
# 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[[huggingface_hub.export_folder_as_dduf]]

```python
huggingface_hub.export_folder_as_dduf(dduf_path: str | os.PathLike, folder_path: str | os.PathLike)
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_dduf.py#L249)

**参数：**

dduf_path (`str` 或 `os.PathLike`) : 要写入的 DDUF 文件的路径。

folder_path (`str` 或 `os.PathLike`) : 包含扩散模型的文件夹的路径。

将文件夹导出为 DDUF 文件。

A在底层使用 [export_entries_as_dduf()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.export_entries_as_dduf)。

示例：
```python
>>> 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[[huggingface_hub.read_dduf_file]]

```python
huggingface_hub.read_dduf_file(dduf_path: os.PathLike | str)
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_dduf.py#L91)

**参数：**

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

**返回：** `dict[str, DDUFEntry]`

一个以文件名为索引的 [DDUFEntry](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.DDUFEntry) 字典。

**引发：** `-`

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

读取 DDUF 文件并返回一个包含条目的字典。

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

示例：
```python
>>> 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)
```

#### huggingface_hub.DDUFEntry[[huggingface_hub.DDUFEntry]]

```python
huggingface_hub.DDUFEntry(filename: str, length: int, offset: int, dduf_path: Path)
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_dduf.py#L36)

**参数：**

filename (str) : DDUF 归档中文件的名称。

offset (int) : 文件在 DDUF 归档中的偏移量。

length (int)：DDUF 归档中文件的长度。

dduf_path (str) : DDUF 存档的路径（供内部使用）。

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

有关如何读取 DDUF 文件，请参阅 [read_dduf_file()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.read_dduf_file)。

#### as_mmap[[huggingface_hub.DDUFEntry.as_mmap]]

```python
as_mmap()
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_dduf.py#L58)

以内存映射文件的形式打开该文件。

用于直接从文件加载 safetensors。

示例：
```py
>>> import safetensors.torch
>>> with entry.as_mmap() as mm:
...     tensors = safetensors.torch.load(mm)
```

#### read_text[[huggingface_hub.DDUFEntry.read_text]]

```python
read_text(encoding: str = 'utf-8')
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_dduf.py#L75)

以文本形式读取该文件。

适用于 .txt 和 .json 条目。

示例：
```py
>>> import json
>>> index = json.loads(entry.read_text())
```

### 错误[[huggingface_hub.errors.DDUFError]]

#### huggingface_hub.errors.DDUFError[[huggingface_hub.errors.DDUFError]]

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/errors.py#L521)

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

#### huggingface_hub.errors.DDUFCorruptedFileError[[huggingface_hub.errors.DDUFCorruptedFileError]]

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/errors.py#L525)

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

#### huggingface_hub.errors.DDUFExportError[[huggingface_hub.errors.DDUFExportError]]

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/errors.py#L529)

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

#### huggingface_hub.errors.DDUFInvalidEntryNameError[[huggingface_hub.errors.DDUFInvalidEntryNameError]]

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/errors.py#L533)

条目名称无效时抛出异常。

## 保存张量

`serialization` 模块的主要辅助函数接收一个 torch `nn.Module` 作为输入，并将其保存到磁盘。它处理保存共享张量的逻辑（参见 [safetensors 说明](https://huggingface.co/docs/safetensors/torch_shared_tensors)），以及使用 [split_torch_state_dict_into_shards()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.split_torch_state_dict_into_shards) 将状态字典拆分为分片的逻辑。目前仅支持 `torch` 框架。

如果您想保存状态字典（例如层名称与相关张量之间的映射）而不是 `nn.Module`，可以使用 [save_torch_state_dict()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.save_torch_state_dict)，它提供相同的功能。例如，如果您想在保存之前对状态字典应用自定义逻辑，这就会很有用。

### save_torch_model[[huggingface_hub.save_torch_model]]

#### huggingface_hub.save_torch_model[[huggingface_hub.save_torch_model]]

```python
huggingface_hub.save_torch_model(model: torch.nn.Module, save_directory: str | pathlib.Path, filename_pattern: str | None = None, force_contiguous: bool = True, max_shard_size: int | str = '5GB', metadata: dict[str, str] | None = None, safe_serialization: bool = True, is_main_process: bool = True, shared_tensors_to_discard: list[str] | None = None)
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_torch.py#L47)

**参数：**

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()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.save_torch_state_dict)，以保存具有更高灵活性的状态字典。

有关张量共享的更多信息，请参阅[此指南](https://huggingface.co/docs/safetensors/torch_shared_tensors)。

模型状态字典会被拆分为多个分片，以确保每个分片的大小小于指定值。这些分片会以给定的 `filename_pattern` 保存在 `save_directory` 中。如果模型过大而无法容纳在单个分片中，则会在 `save_directory` 中保存一个索引文件，以指示每个张量的保存位置。此辅助函数底层使用了 [split_torch_state_dict_into_shards()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.split_torch_state_dict_into_shards)。如果 `safe_serialization` 为 `True`，分片将以 safetensors 格式保存（默认行为）。否则，分片将以 pickle 格式保存。

在保存模型之前，`save_directory` 中任何之前的分片文件都会被清除。

> [!WARNING]
> 如果模型的某个张量大于 `max_shard_size`，它将被单独放入一个分片中，该分片的大小将超过 `max_shard_size`。

> [!WARNING]
> 如果您的模型是 `transformers.PreTrainedModel`，您应该将 `model._tied_weights_keys` 作为 `shared_tensors_to_discard` 传入，以正确处理共享张量的保存。这可以确保在保存期间正确丢弃重复的张量。

示例：

```py
>>> 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]]

#### huggingface_hub.save_torch_state_dict[[huggingface_hub.save_torch_state_dict]]

```python
huggingface_hub.save_torch_state_dict(state_dict: dict, save_directory: str | pathlib.Path, filename_pattern: str | None = None, force_contiguous: bool = True, max_shard_size: int | str = '5GB', metadata: dict[str, str] | None = None, safe_serialization: bool = True, is_main_process: bool = True, shared_tensors_to_discard: list[str] | None = None)
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_torch.py#L141)

**参数：**

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()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.save_torch_model)，用于直接保存 PyTorch 模型。

有关张量共享的更多信息，请参阅[此指南](https://huggingface.co/docs/safetensors/torch_shared_tensors)。

模型状态字典会被拆分为多个分片，以确保每个分片的大小小于指定值。这些分片会以给定的 `filename_pattern` 保存在 `save_directory` 中。如果模型过大而无法容纳在单个分片中，则会在 `save_directory` 中保存一个索引文件，以指示每个张量的保存位置。此辅助函数底层使用了 [split_torch_state_dict_into_shards()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.split_torch_state_dict_into_shards)。如果 `safe_serialization` 为 `True`，分片将以 safetensors 格式保存（默认行为）。否则，分片将以 pickle 格式保存。

在保存模型之前，`save_directory` 中任何之前的分片文件都会被清除。

> [!WARNING]
> 如果模型的某个张量大于 `max_shard_size`，它将被单独放入一个分片中，该分片的大小将超过 `max_shard_size`。

> [!WARNING]
> 如果您的模型是 `transformers.PreTrainedModel`，您应该将 `model._tied_weights_keys` 作为 `shared_tensors_to_discard` 传入，以正确处理共享张量的保存。这可以确保在保存期间正确丢弃重复的张量。

示例：

```py
>>> 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]]

#### huggingface_hub.split_torch_state_dict_into_shards[[huggingface_hub.split_torch_state_dict_into_shards]]

```python
huggingface_hub.split_torch_state_dict_into_shards(state_dict: dict, filename_pattern: str = 'model{suffix}.safetensors', max_shard_size: int | str = '5GB')
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_torch.py#L302)

**参数：**

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]。

> [!TIP]
> 要将模型状态字典保存到磁盘，请参阅 [save_torch_state_dict()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.save_torch_state_dict)。该辅助函数在底层使用了 `split_torch_state_dict_into_shards`。

> [!WARNING]
> 如果模型的某个张量大于 `max_shard_size`，它将被单独放入一个分片中，该分片的大小将超过 `max_shard_size`。

示例：
```py
>>> 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.split_state_dict_into_shards_factory]]

这是每个特定于框架的辅助函数所基于的底层工厂。在实际使用中，除非您需要将其适配到尚未支持的框架，否则不建议您直接使用此工厂。如果确实如此，请在 `huggingface_hub` 仓库中通过[提交一个新 issue](https://github.com/huggingface/huggingface_hub/issues/new) 告知我们。

#### huggingface_hub.split_state_dict_into_shards_factory[[huggingface_hub.split_state_dict_into_shards_factory]]

```python
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 0x7f9673968670>, max_shard_size: int | str = '5GB')
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_base.py#L50)

**参数：**

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]。

> [!WARNING]
> 如果模型的某个张量大于 `max_shard_size`，它将被单独放入一个分片中，该分片的大小将超过 `max_shard_size`。

## 加载张量

加载辅助函数支持单文件和分片检查点，格式可以是 safetensors 或 pickle。[load_torch_model()](/docs/huggingface_hub/v2.4.0/zh/package_reference/serialization#huggingface_hub.load_torch_model) 接收一个 `nn.Module` 和一个检查点路径（单个文件或目录）作为输入，并将权重加载到模型中。

> [!WARNING]
> 这两个辅助函数默认均为 `safe=True`，这意味着无论文件名称如何，检查点始终通过 safetensors 加载器进行反序列化。加载 pickle 检查点需要显式选择，因为反序列化 pickle 会在加载时执行任意代码：
> ```py
> >>> 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]]

#### huggingface_hub.load_torch_model[[huggingface_hub.load_torch_model]]

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

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_torch.py#L375)

**参数：**

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 的受限反序列化器。设置为 False 可允许 pickle 检查点中包含任意 Python 对象（这会在加载时执行任意代码）。对 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` 是一个包含缺失键的 str 列表，即模型中存在但检查点中不存在的键。
- `unexpected_keys` 是一个包含意外键的 str 列表，即存在于检查点中但不在模型中的键。

**引发：** ``FileNotFoundError`` 或 ``ImportError`` 或 ``ValueError``

- [`FileNotFoundError`](https://docs.python.org/3/library/exceptions.html#FileNotFoundError) -- 如果检查点文件或目录不存在。
- [`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError) -- 当尝试加载 .safetensors 文件或 PyTorch 检查点时，如果分别未安装 safetensors 或 torch，则会抛出此错误。
- [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) -- 如果无法确定检查点格式。

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

示例：
```python
>>> 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]]

#### huggingface_hub.load_state_dict_from_file[[huggingface_hub.load_state_dict_from_file]]

```python
huggingface_hub.load_state_dict_from_file(checkpoint_file: str | os.PathLike, map_location: typing.Union[str, ForwardRef('torch.device'), NoneType] = None, weights_only: bool = True, mmap: bool = False, safe: bool = True)
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_torch.py#L633)

**参数：**

checkpoint_file (`str` 或 `os.PathLike`) : 要加载的检查点文件的路径。可以是 safetensors 或 pickle (`.bin`) 检查点。

map_location (`str` 或 `torch.device`，*可选*)：一个 `torch.device` 对象、字符串或字典，用于指定如何重新映射存储位置。它指示所有张量应加载到的位置。

weights_only (`bool`, *可选*，默认为 `True`)：如果为 True，则仅加载模型权重，不加载优化器状态和其他元数据，并使用 torch 的受限 unpickler。设置为 False 以允许 pickle 检查点中包含任意 Python 对象（这会在加载时执行任意代码）。在加载 safetensors 文件时无效，对没有受限 unpickler 的 PyTorch < 1.13 也无效——这些版本始终无限制地执行 unpickle，并会记录一条警告。

mmap (`bool`, *可选*, 默认为 `False`) : 是否使用内存映射文件加载。对于 PyTorch >= 2.1.0 中基于 zipfile 的检查点，内存映射可以提升大模型的加载性能。在加载 safetensors 文件时此参数无效，因为 `safetensors` 库默认使用内存映射。

safe (`bool`，*可选*，默认为 `True`)：如果为 True，无论检查点文件名是什么，都会始终将其作为 safetensors 加载。任何其他格式（例如 pickle `.bin` 文件）都会引发 `ValueError`，而不是被反序列化：pickle 检查点在加载时可能会执行任意代码。如果为 False，则当文件名表明其为 safetensors 时按 safetensors 加载，否则使用 `torch.load` 加载。

**返回：** `Union[dict[str, "torch.Tensor"], Any]`

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

**引发：** ``FileNotFoundError`` 或 ``ImportError`` 或 ``OSError`` 或 ``ValueError``

- [`FileNotFoundError`](https://docs.python.org/3/library/exceptions.html#FileNotFoundError) -- 如果检查点文件不存在。
- [`ImportError`](https://docs.python.org/3/library/exceptions.html#ImportError) -- 当尝试加载 .safetensors 文件或 PyTorch 检查点时，如果分别未安装 safetensors 或 torch，则会抛出此错误。
- [`OSError`](https://docs.python.org/3/library/exceptions.html#OSError) -- 如果检查点文件格式无效，或者 git-lfs 文件未正确下载。
- [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) -- 如果检查点文件路径为空或无效，或者在 `safe=True` 时无法将其反序列化为 safetensors（请传入 `safe=False` 以允许使用 pickle 检查点）。

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

示例：
```python
>>> 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]]

#### huggingface_hub.get_torch_storage_id[[huggingface_hub.get_torch_storage_id]]

```python
huggingface_hub.get_torch_storage_id(tensor: torch.Tensor)
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_torch.py#L878)

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

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

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

### get_torch_storage_size[[huggingface_hub.get_torch_storage_size]]

#### huggingface_hub.get_torch_storage_size[[huggingface_hub.get_torch_storage_size]]

```python
huggingface_hub.get_torch_storage_size(tensor: torch.Tensor)
```

[源代码](https://github.com/huggingface/huggingface_hub/blob/v2.4.0/src/huggingface_hub/serialization/_torch.py#L895)

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