序列化

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 提供了一種直接將管道序列化為 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)

類 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()。

作為記憶體對映

< 源 >

( )

以記憶體對映檔案開啟檔案。

對於直接從檔案載入 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()，它提供了相同的功能。例如，如果您想在儲存之前對狀態字典應用自定義邏輯，這會很有用。

儲存 PyTorch 模型

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() 以更靈活地儲存狀態字典。

有關張量共享的更多資訊，請檢視此指南。

模型狀態字典被拆分為分片，以便每個分片都小於給定大小。分片將以給定的 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")
>>>

儲存 PyTorch 狀態字典

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 (布林值, 可選) — 強制狀態字典儲存為連續張量。這不會影響模型的正確性，但如果張量的佈局是專門為此目的選擇的，則可能會改變效能。預設為 True。
max_shard_size (int 或 str, 可選) — 每個分片的最大大小，以位元組為單位。預設為 5GB。
metadata (Dict[str, str], 可選) — 隨模型一起儲存的額外資訊。對於每個丟棄的張量，都會新增一些元資料。此資訊不足以恢復整個共享結構，但可能有助於理解。
safe_serialization (布林值, 可選) — 是否以 safetensors 格式儲存，這是預設行為。如果為 False，則分片將儲存為 pickle 格式。出於安全考慮，建議使用安全序列化。儲存為 pickle 格式已棄用，並將在未來版本中移除。
is_main_process (布林值, 可選) — 呼叫此函式的程序是否為主程序。在分散式訓練（如 TPUs）中需要從所有程序呼叫此函式時很有用。在這種情況下，只有主程序才將 is_main_process 設定為 True 以避免競態條件。預設為 True。
shared_tensors_to_discard (List[str], 可選) — 儲存共享張量時要丟棄的張量名稱列表。如果未提供且檢測到共享張量，則將按字母順序丟棄第一個名稱。

將模型狀態字典儲存到磁碟，處理分片和共享張量問題。

另請參閱 save_torch_model() 以直接儲存 PyTorch 模型。

有關張量共享的更多資訊，請檢視此指南。

在儲存模型之前，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 張量，旨在輕鬆擴充套件到任何其他 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, 可選) — 用於生成模型儲存檔名的模式。模式必須是可使用 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: 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, 可選) — 用於生成模型儲存檔名的模式。模式必須是可使用 filename_pattern.format(suffix=...) 格式化的字串，並且必須包含關鍵字 suffix。預設為 "model{suffix}.safetensors"。
max_shard_size (int 或 str, 可選) — 每個分片的最大大小，以位元組為單位。預設為 5GB。

StateDictSplit

一個包含分片和檢索它們的索引的 StateDictSplit 物件。

將模型狀態字典分片，使每個分片小於給定大小。

要將模型狀態字典儲存到磁碟，請參閱 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: 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 0x7fd63c3b1e10> 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]], 可選) — 一個函式，返回張量儲存的唯一識別符號。多個不同的張量可以共享相同的底層儲存。此識別符號在張量的生命週期內保證是唯一的和恆定的。兩個具有非重疊生命週期的張量儲存可能具有相同的 ID。
filename_pattern (str, 可選) — 用於生成模型儲存檔名的模式。模式必須是可使用 filename_pattern.format(suffix=...) 格式化的字串，並且必須包含關鍵字 suffix。
max_shard_size (int 或 str, 可選) — 每個分片的最大大小，以位元組為單位。預設為 5GB。

StateDictSplit

一個包含分片和檢索它們的索引的 StateDictSplit 物件。

將模型狀態字典分片，使每個分片小於給定大小。

如果模型中的某個張量大於 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) — 要載入檢查點的模型。
checkpoint_path (str 或 os.PathLike) — 檢查點檔案或包含檢查點(s)的目錄的路徑。
strict (布林值, 可選, 預設為 False) — 是否嚴格執行模型狀態字典中的鍵與檢查點中的鍵匹配。
safe (布林值, 可選, 預設為 True) — 如果 safe 為 True，則將載入 safetensors 檔案。如果 safe 為 False，函式將首先嚐試載入 safetensors 檔案（如果可用），否則將回退到載入 pickle 檔案。filename_pattern 引數優先於 safe 引數。
weights_only (布林值, 可選, 預設為 False) — 如果為 True，則僅載入模型權重，不包括最佳化器狀態和其他元資料。僅在 PyTorch >= 1.13 中受支援。
map_location (str 或 torch.device, 可選) — 一個 torch.device 物件、字串或字典，指定如何重新對映儲存位置。它指示所有張量應載入到的位置。
mmap (布林值, 可選, 預設為 False) — 是否使用記憶體對映檔案載入。記憶體對映可以改善 PyTorch >= 2.1.0 中基於 zipfile 的檢查點的大型模型載入效能。
filename_pattern (str, 可選) — 用於查詢索引檔案的模式。模式必須是可使用 filename_pattern.format(suffix=...) 格式化的字串，並且必須包含關鍵字 suffix。預設為 "model{suffix}.safetensors"。

命名元組

一個包含 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 (布林值, 可選, 預設為 False) — 如果為 True，則僅載入模型權重，不包括最佳化器狀態和其他元資料。僅支援 PyTorch >= 1.13 的 pickle (.bin) 檢查點。載入 safetensors 檔案時無效。
mmap (布林值, 可選, 預設為 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 上更新

Hub Python 庫

序列化

DDUF 檔案格式

如何編寫 DDUF 檔案？

如何讀取 DDUF 檔案？

輔助函式

huggingface_hub.export_entries_as_dduf

huggingface_hub.export_folder_as_dduf

huggingface_hub.read_dduf_file

類 huggingface_hub.DDUFEntry

作為記憶體對映

read_text

錯誤

類 huggingface_hub.errors.DDUFError

類 huggingface_hub.errors.DDUFCorruptedFileError

類 huggingface_hub.errors.DDUFExportError

類 huggingface_hub.errors.DDUFInvalidEntryNameError

儲存張量

儲存 PyTorch 模型

huggingface_hub.save_torch_model

儲存 PyTorch 狀態字典

huggingface_hub.save_torch_state_dict

split_tf_state_dict_into_shards

huggingface_hub.split_tf_state_dict_into_shards

split_torch_state_dict_into_shards

huggingface_hub.split_torch_state_dict_into_shards

split_state_dict_into_shards_factory

huggingface_hub.split_state_dict_into_shards_factory

載入張量

load_torch_model

huggingface_hub.load_torch_model

load_state_dict_from_file

huggingface_hub.load_state_dict_from_file

張量輔助函式

get_torch_storage_id

huggingface_hub.get_torch_storage_id

get_torch_storage_size

huggingface_hub.get_torch_storage_size