Hub Python 函式庫文件
序列化
並獲得增強的文件體驗
開始使用
序列化
huggingface_hub 提供了以標準化方式儲存與載入機器學習模型權重的輔助工具。函式庫的這部分功能仍在開發中，未來版本將持續改進。其目標是統一 Hub 上權重的儲存與載入方式，既能消除各函式庫間重複的程式碼，也能建立一致的規範。
DDUF 檔案格式
DDUF 是一種專為擴散模型（diffusion models）設計的檔案格式。它允許將執行模型所需的所有資訊儲存在單一檔案中。這項工作受到 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 參數也支援傳入一個路徑或位元組（bytes）的迭代器。如果您已經載入模型並希望直接將其序列化為 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.PathLike entries: Iterable )
從條目迭代器寫入 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.PathLike folder_path: str | os.PathLike )
將資料夾匯出為 DDUF 檔案。
內部呼叫 export_entries_as_dduf()。
huggingface_hub.read_dduf_file
< 原始碼 >( dduf_path: os.PathLike | str ) → dict[str, DDUFEntry]
參數
返回
dict[str, DDUFEntry]
以檔名索引的 DDUFEntry 字典。
引發
-
- —
DDUFCorruptedFileError：若 DDUF 檔案損毀（即不符合 DDUF 格式）。
- —
讀取 DDUF 檔案並返回條目字典。
僅讀取中繼資料（metadata），不會將資料載入至記憶體。
範例
>>> 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 )
代表 DDUF 檔案中檔案條目的物件。
請參閱 read_dduf_file() 以了解如何讀取 DDUF 檔案。
將檔案作為記憶體映射檔案（memory-mapped file）開啟。
適用於直接從檔案載入 safetensors。
將檔案讀取為文字。
適用於 ‘.txt’ 和 ‘.json’ 條目。
錯誤
與 DDUF 格式相關錯誤的基礎例外類別。
當 DDUF 檔案損毀時拋出的例外。
DDUF 匯出過程中發生錯誤的基礎例外類別。
當條目名稱無效時拋出的例外。
儲存張量（tensors）
serialization 模組的主要輔助工具接收一個 torch nn.Module 作為輸入並將其儲存到磁碟。它處理儲存共享張量的邏輯（請參閱 safetensors 說明），以及使用 split_torch_state_dict_into_shards() 將狀態字典（state dictionary）切分為分片（shards）的邏輯。目前僅支援 torch 框架。
如果您想儲存狀態字典（例如層名稱與相關張量之間的映射）而非 nn.Module，可以使用 save_torch_state_dict()，它提供了相同的功能。例如，如果您想在儲存狀態字典之前應用自定義邏輯，這會非常有用。
save_torch_model
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 )
參數
- 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 儲存為連續（contiguous）張量。這對模型的正確性沒有影響，但如果張量佈局是基於特定原因選擇的，這可能會改變效能。預設為True。 - max_shard_size (
int或str，選填) — 每個分片的最大位元組數。預設為 5GB。 - metadata (
dict[str, str]，選填) — 隨模型一起儲存的額外資訊。會為每個被丟棄的張量新增一些中繼資料。這些資訊不足以恢復整個共享結構，但可能有助於理解相關內容。 - safe_serialization (
bool，選填) — 是否儲存為 safetensors（這是預設行為）。若設為False，分片將以 pickle 格式儲存。基於安全性考量，推薦使用安全序列化（safe serialization）。儲存為 pickle 已被棄用，將在未來的版本中移除。 - is_main_process (
bool，選填) — 呼叫此函式的程序是否為主要程序。在 TPU 等分散式訓練中，需要從所有程序呼叫此函式時非常有用。在此情況下，僅在主要程序上設定is_main_process=True以避免競爭條件（race conditions）。預設為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: 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 )
參數
- 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_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 張量，並設計為易於擴充至任何其他機器學習框架。
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: 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: dict get_storage_size: Callable filename_pattern: str get_storage_id: Callable = <function <lambda> at 0x7f3f168ce680> max_shard_size: int | str = '5GB' ) → StateDictSplit
參數
- state_dict (
dict[str, Tensor]) — 要儲存的狀態字典 (state dictionary)。 - 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, 選填) — 每個分片 (shard) 的最大大小（以位元組為單位）。預設為 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 格式的單一檔案與分片檢查點 (checkpoints)。load_torch_model() 接受一個 nn.Module 與檢查點路徑（單一檔案或目錄）作為輸入，並將權重載入模型中。
load_torch_model
huggingface_hub.load_torch_model
< 原始碼 >( model: torch.nn.Module checkpoint_path: 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: str | None = 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) — 是否使用記憶體映射 (memory-mapped) 檔案載入。在 PyTorch >= 2.1.0 中，對於基於 zipfile 的檢查點，記憶體映射可以提升大型模型的載入效能。 - filename_pattern (
str, 選填) — 用於搜尋索引檔案的模式。模式必須是一個可以透過filename_pattern.format(suffix=...)進行格式化的字串，且必須包含suffix關鍵字。預設為"model{suffix}.safetensors"。
返回
NamedTuple
一個具有 missing_keys 和 unexpected_keys 欄位的命名元組 (NamedTuple)。
missing_keys是一個包含缺漏鍵的字串列表，即存在於模型中但不在檢查點中的鍵。unexpected_keys是一個包含非預期鍵的字串列表，即存在於檢查點中但不在模型中的鍵。
引發
FileNotFoundError 或 ImportError 或 ValueError
FileNotFoundError— 如果檢查點檔案或目錄不存在。ImportError— 如果在嘗試載入 .safetensors 檔案時未安裝 safetensors，或在嘗試載入 PyTorch 檢查點時未安裝 torch。ValueError— 如果檢查點路徑無效，或無法確定檢查點格式。
將檢查點載入到模型中，處理分片與未分片的檢查點。
load_state_dict_from_file
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 = 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 檔案：返回任何已被序列化的 Python 物件（通常是狀態字典，但也可能是整個模型、優化器狀態或其他任何 Python 物件）。
引發
FileNotFoundError 或 ImportError 或 OSError 或 ValueError
FileNotFoundError— 如果檢查點檔案不存在。ImportError— 如果在嘗試載入 .safetensors 檔案時未安裝 safetensors，或在嘗試載入 PyTorch 檢查點時未安裝 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
返回張量儲存的唯一識別碼。
多個不同的張量可以共用相同的底層儲存。此識別碼保證在張量儲存的生命週期內是唯一且恆定的。兩個生命週期不重疊的張量儲存可能擁有相同的 ID。對於 meta 張量，我們返回 None，因為我們無法判定它們是否共用相同的儲存。