Hugging Face 的 LOGO

Hub Python 庫文件

序列化

Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

序列化

huggingface_hub 提供了一個標準化的方法來儲存和載入機器學習模型權重。該庫的這部分仍在開發中，將在未來的版本中進行改進。目標是統一模型權重在 Hub 上的儲存和載入方式，從而消除跨庫的程式碼重複並建立一致的約定。

DDUF 檔案格式

DDUF 是一種為擴散模型設計的[_檔案格式_](https://huggingface.co/docs/huggingface_hub/package_reference/serialization#dduf-file-format)。它允許將執行模型的所有資訊儲存在一個檔案中。這項工作受到了 [_GGUF_](https://github.com/ggerganov/ggml/blob/master/docs/gguf.md) 格式的啟發。huggingface_hub 提供了儲存和載入 DDUF 檔案的助手，確保檔案格式得到遵守。

這是解析器的一個非常早期的版本。API 和實現可能在不久的將來發生變化。

該解析器目前只進行了很少的驗證。有關檔案格式的更多詳細資訊，請檢視 [_https://github.com/huggingface/huggingface.js/tree/main/packages/dduf_](https://github.com/huggingface/huggingface.js/tree/main/packages/dduf)。

如何編寫 DDUF 檔案？

下面將介紹如何使用 [_export_folder_as_dduf()_](/docs/huggingface_hub/v1.4.0/en/package_reference/serialization#huggingface_hub.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()_](/docs/huggingface_hub/v1.4.0/en/package_reference/serialization#huggingface_hub.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[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()_](/docs/huggingface_hub/v1.4.0/en/package_reference/serialization#huggingface_hub.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: 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()_](/docs/huggingface_hub/v1.4.0/en/package_reference/serialization#huggingface_hub.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()_](/docs/huggingface_hub/v1.4.0/en/package_reference/serialization#huggingface_hub.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 解釋）以及將狀態字典分割成分片（shard）的邏輯，底層使用了 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[dict[str, str]] = None safe_serialization: bool = True is_main_process: bool = True shared_tensors_to_discard: typing.Optional[list[str]] = None )

引數

  • model (torch.nn.Module) — 要儲存到磁碟的模型。
  • save_directory (str or Path) — 模型將被儲存的目錄。
  • filename_pattern (str, optional) — 用於生成模型將要儲存的檔名的模式。模式必須是一個可以被 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 or str, optional) — 每個分片的 maximum size（以位元組為單位）。預設為 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) — 儲存共享張量時要丟棄的張量名稱列表。如果未提供且檢測到共享張量，則會按字母順序丟棄第一個名稱。

將給定的 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: dict 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[dict[str, str]] = None safe_serialization: bool = True is_main_process: bool = True shared_tensors_to_discard: typing.Optional[list[str]] = None )

引數

  • state_dict (dict[str, torch.Tensor]) — 要儲存的狀態字典。
  • save_directory (str or Path) — 模型將被儲存的目錄。
  • filename_pattern (str, optional) — 用於生成模型將要儲存的檔名的模式。模式必須是一個可以被 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 or str, optional) — 每個分片的 maximum size（以位元組為單位）。預設為 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], 可選) — 儲存共享張量時要丟棄的張量名稱列表。如果未提供且檢測到共享張量，它將丟棄按字母順序排列的第一個名稱。

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

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

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

模型狀態字典被分割成若干分片，使得每個分片的大小小於給定大小。分片將使用給定的 filename_pattern 儲存到 save_directory 中。如果模型太大而無法放入單個分片，則會在 save_directory 中儲存一個索引檔案，指示每個張量儲存在何處。此輔助函式底層使用了 split_torch_state_dict_into_shards()。如果 safe_serialization 為 True，則分片將儲存為 safetensors（預設）。否則，分片將儲存為 pickle。

在儲存模型之前，save_directory 會被清理掉任何先前存在的切片檔案。

如果模型中的某個張量大於 max_shard_size，它將落入自己的分片中，該分片的大小將大於 max_shard_size。

如果您的模型是 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 張量，並設計為易於擴充套件到任何其他 ML 框架。

split_torch_state_dict_into_shards

huggingface_hub.split_torch_state_dict_into_shards

< >

( state_dict: dict filename_pattern: str = 'model{suffix}.safetensors' max_shard_size: 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 物件，包含用於檢索分片的索引。

將模型狀態詞典分割成碎片，使每個碎片小於給定大小。

分片是透過按鍵的順序迭代 state_dict 來確定的。我們不會進行任何最佳化以使每個分片儘可能接近傳遞的最大大小。例如，如果限制是 10GB 且我們有大小為 [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] 的張量，它們將被分片為 [6GB], [6+2GB], [6+2+2GB]，而不是 [6+2+2GB], [6+2GB], [6GB]。

要將模型狀態詞典儲存到磁碟，請參閱 save_torch_state_dict()。此輔助函式在底層使用 split_torch_state_dict_into_shards。

如果模型中的某個張量大於 max_shard_size，它將落入自己的分片中，該分片的大小將大於 max_shard_size。

示例

>>> import json
>>> import os
>>> from safetensors.torch import save_file as safe_save_file
>>> from huggingface_hub import split_torch_state_dict_into_shards

>>> def save_state_dict(state_dict: dict[str, torch.Tensor], save_directory: str):
...     state_dict_split = split_torch_state_dict_into_shards(state_dict)
...     for filename, tensors in state_dict_split.filename_to_tensors.items():
...         shard = {tensor: state_dict[tensor] for tensor in tensors}
...         safe_save_file(
...             shard,
...             os.path.join(save_directory, filename),
...             metadata={"format": "pt"},
...         )
...     if state_dict_split.is_sharded:
...         index = {
...             "metadata": state_dict_split.metadata,
...             "weight_map": state_dict_split.tensor_to_filename,
...         }
...         with open(os.path.join(save_directory, "model.safetensors.index.json"), "w") as f:
...             f.write(json.dumps(index, indent=2))

split_state_dict_into_shards_factory

這是派生出每個框架特定幫助函式的底層工廠。在實踐中，您不應該直接使用此工廠，除非您需要將其適配到尚未支援的框架。如果是這種情況，請透過在 huggingface_hub 倉庫上開啟一個新問題來告知我們。

huggingface_hub.split_state_dict_into_shards_factory

< >

( state_dict: dict get_storage_size: typing.Callable[[~TensorT], int] filename_pattern: str get_storage_id: typing.Callable[[~TensorT], typing.Optional[typing.Any]] = <function <lambda> at 0x7ff1eb636a70> 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 物件，包含用於檢索分片的索引。

將模型狀態詞典分割成碎片，使每個碎片小於給定大小。

分片是透過按鍵的順序迭代 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) — 用於載入檢查點的模型。
  • 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 物件（通常是 state_dict，但也可以是整個模型、最佳化器狀態或任何其他 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。對於 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

< >

( tensor: torch.Tensor )

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

在 GitHub 上更新

←Webhooks server Strict dataclasses→
© . This site is unofficial and not affiliated with Hugging Face, Inc.