Hub Python Library documentation
シリアライズ
シリアライズ
huggingface_hub は、ML モデルの重みを標準的な方法で保存および読み込むためのヘルパーを提供します。このライブラリの部分はまだ開発中で、将来のリリースで改善される予定です。その目標は、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
< source >( 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
< source >( dduf_path: str | os.PathLikefolder_path: str | os.PathLike )
フォルダーをDDUFファイルとしてエクスポートします。
Aは内部でexport_entries_as_dduf()を使用しています。
huggingface_hub.read_dduf_file
< source >( dduf_path: os.PathLike | str ) → dict[str, DDUFEntry]
パラメータ
戻り値
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
< source >( filename: strlength: intoffset: intdduf_path: Path )
DDUFファイル内のファイルエントリを表すオブジェクト。
DDUFファイルの読み取り方法については、read_dduf_file()を参照してください。
ファイルをメモリマップドファイルとして開きます。
ファイルからsafetensorsを直接読み込む際に便利です。
ファイルをテキストとして読み取ります。
「.txt」および「.json」の項目に有用です。
エラー
DDUF形式に関連するエラーの基底例外です。
DDUFファイルが破損している場合にスローされる例外。
DDUFのエクスポート中に発生するエラーの基底例外です。
エントリ名が無効な場合にスローされる例外。
テンソルの保存
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
< source >( 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を含める必要があります。safe_serializationパラメータに応じて、デフォルトは"model{suffix}.safetensors"またはpytorch_model{suffix}.binです。 - 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に保存されます。モデルが1つのシャードに収まらないほど大きい場合、各テンソルがどこに保存されているかを示すインデックスファイルが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
< source >( 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を含める必要があります。safe_serializationパラメータに応じて、デフォルトは"model{suffix}.safetensors"またはpytorch_model{suffix}.binです。 - 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], 省略可) — 共有テンソルを保存する際に破棄するテンソル名のリストです。指定されず、共有テンソルが検出された場合、アルファベット順で最初の名称が破棄されます。
モデルのステートディクショナリディスクに保存し、シャーディングや共有テンソルに関する問題を処理します。
PyTorchモデルを直接保存するには、save_torch_model()も参照してください。
テンソル共有の詳細については、このガイドを参照してください。
モデルのステートディクショナリは、各シャードが指定されたサイズより小さくなるように分割されます。シャードは、指定されたfilename_patternに従ってsave_directoryに保存されます。モデルが1つのシャードに収まらないほど大きい場合、各テンソルがどこに保存されているかを示すインデックスファイルが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
< source >( 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
< source >( state_dict: dictget_storage_size: Callablefilename_pattern: strget_storage_id: Callable = <function <lambda> at 0x7f56f657d510>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]]、省略可) — テンソルのストレージに対する一意の識別子を返す関数です。複数の異なるテンソルが同じ基盤ストレージを共有する場合があります。この識別子は、テンソルのストレージの生存期間中、一意かつ一定であることが保証されています。生存期間が重複しない2つのテンソルストレージは、同じ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形式のチェックポイントを読み込むには明示的なオプトインが必要です。これは、アンピクル処理が読み込み時に任意のコードを実行するためです。>>> 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
< source >( 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) — モデルの state dict のキーがチェックポイントのキーと一致することを厳密に強制するかどうかを指定します。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 の制限付きアンピクルアを使用し、オプティマイザーの状態やその他のメタデータなしにモデルの重みのみを読み込みます。pickle チェックポイント内の任意の Python オブジェクトを許可するには False に設定してください（これにより、読み込み時に任意のコードが実行されます）。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— チェックポイントファイルまたはディレクトリが存在しない場合。ImportError— .safetensorsファイルまたはPyTorchのチェックポイントをそれぞれ読み込もうとした際に、safetensorsまたはtorchがインストールされていない場合。ValueError— チェックポイントの形式を判定できない場合。
…
load_state_dict_from_file
huggingface_hub.load_state_dict_from_file
< source >( 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 の制限付きアンピクルアを使用し、オプティマイザーの状態やその他のメタデータなしにモデルの重みのみを読み込みます。pickle チェックポイント内の任意の Python オブジェクトを許可するには False に設定してください(これにより、読み込み時に任意のコードが実行されます)。safetensors ファイルの読み込み時には効果がなく、制限付きアンピクルアを持たない PyTorch < 1.13 にも影響しません。これらのバージョンでは常に制限なしでアンピクルされ、警告がログに記録されます。 - mmap (
bool、任意、デフォルトはFalse) — メモリマップドファイル読み込みを使用するかどうか。PyTorch >= 2.1.0 での zipfile ベースのチェックポイントにおいて、大規模モデルの読み込みパフォーマンスを向上させる可能性があります。safetensorsライブラリはデフォルトでメモリマッピングを使用するため、safetensors ファイルの読み込み時には影響しません。 - safe (
bool, optional, defaults toTrue) — If True, the checkpoint is always loaded as safetensors, whatever its name. Any other format (e.g. a pickle.binfile) raises aValueErrorinstead of being deserialized: pickle checkpoints can execute arbitrary code at load time. If False, the file is loaded as safetensors when its name says so, and withtorch.loadotherwise.
戻り値
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としてデシリアライズできない場合（pickle形式のチェックポイントを許可するにはsafe=Falseを渡します）。
チェックポイントファイルを読み込み、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
テンソルストレージの一意の識別子を返します。
複数の異なるテンソルが同じ基盤ストレージを共有できます。この識別子は、テンソルのストレージの生存期間中、一意かつ一定であることが保証されています。生存期間が重複しない2つのテンソルストレージは、同じidを持つ可能性があります。 メタテンソルの場合は、同じストレージを共有しているかどうかを判断できないため、Noneを返します。