Hub Python Library documentation

직렬화

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

직렬화

huggingface_hub은 ML 모델 가중치를 표준화된 방식으로 저장하고 로드하는 데 사용할 수 있는 헬퍼를 제공합니다. 이 라이브러리 부분은 아직 개발 중이며 향후 릴리스에서 개선될 예정입니다. 목표는 라이브러리 간 코드 중복을 제거하고 일관된 규약을 확립하기 위해 Hub 전체에서 가중치가 저장되고 로드되는 방식을 조화시키는 것입니다.

DDUF 파일 형식

DDUF은 diffusion 모델을 위한 파일 형식입니다. 모델 실행에 필요한 모든 정보를 단일 파일로 저장할 수 있습니다. 이 작업은 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: 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’ 누락 등).

entries 이터러블에서 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.PathLikefolder_path: str | os.PathLike )

매개변수

  • dduf_path (str 또는 os.PathLike) — 쓰기 위한 DDUF 파일의 경로입니다.
  • folder_path (str 또는 os.PathLike) — 확산 모델이 포함된 폴더의 경로입니다.

DDUF 파일로 폴더를 내보냅니다.

A는 내부적으로 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: os.PathLike | str ) → dict[str, DDUFEntry]

매개변수

  • dduf_path (str 또는 os.PathLike) — 읽을 DDUF 파일의 경로입니다.

반환값

dict[str, DDUFEntry]

파일 이름으로 인덱싱된 DDUFEntry의 사전입니다.

다음 예외가 발생합니다

    • — DDUFCorruptedFileError: DDUF 파일이 손상된 경우(즉, DDUF 형식을 따르지 않는 경우).

DDUF 파일을 읽고 항목의 사전(dict)을 반환합니다.

메타데이터만 읽으며 데이터는 메모리에 로드되지 않습니다.

예:

>>> 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: strlength: intoffset: intdduf_path: Path )

매개변수

  • filename (str) — DDUF 아카이브 내 파일 이름.
  • offset (int) — DDUF 아카이브에서 파일의 오프셋입니다.
  • length (int) — DDUF 아카이브에서 파일의 길이입니다.
  • dduf_path (str) — DDUF 아카이브의 경로(내부 사용용).

DDUF 파일의 파일 항목을 나타내는 객체.

DDUF 파일을 읽는 방법은 read_dduf_file()을 참조하세요.

as_mmap

< >

( )

memory-mapped 파일로 파일을 엽니다.

…

예:

>>> import safetensors.torch
>>> with entry.as_mmap() as mm:
...     tensors = safetensors.torch.load(mm)

read_text

< >

( encoding: str = 'utf-8' )

Read the file as text.

‘.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()를 사용하여 상태 사전(state dictionary)을 샤드로 분할하는 로직을 처리합니다. 현재 torch 프레임워크만 지원됩니다.

nn.Module 대신 상태 사전(예: 레이어 이름과 관련 텐서 간의 매핑)을 저장하려면 동일한 기능을 제공하는 save_torch_state_dict()를 사용할 수 있습니다. 이는 저장 전에 상태 사전에 사용자 정의 로직을 적용하려는 경우에 유용합니다.

save_torch_model

huggingface_hub.save_torch_model

< >

( 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로 저장하는 기능은 비추천(deprecated) 상태이며, 향후 버전에서 제거될 예정입니다.
  • is_main_process (bool, 선택 사항) — 이 함수를 호출하는 프로세스가 메인 프로세스인지 여부입니다. TPU와 같은 분산 학습 환경에서 모든 프로세스에서 이 함수를 호출해야 할 때 유용합니다. 이 경우 레이스 조건을 방지하기 위해 메인 프로세스에서만 is_main_process=True로 설정해야 합니다. 기본값은 True입니다.
  • shared_tensors_to_discard (list[str], 선택 사항) — 공유 텐서를 저장할 때 제거할 텐서 이름 목록입니다. 제공되지 않고 공유 텐서가 감지되면, 알파벳 순서로 첫 번째 이름을 제거합니다.

torch 모델을 디스크에 저장하며, 샤딩 및 공유 텐서 문제를 처리합니다.

더 유연하게 state dict를 저장하려면 save_torch_state_dict()도 참고하세요.

텐서 공유에 대한 자세한 내용은 이 가이드를 확인하세요.

모델 상태 사전은 각 샤드가 주어진 크기보다 작도록 샤드로 분할됩니다. 샤드는 save_directory에 주어진 filename_pattern으로 저장됩니다. 모델이 단일 샤드에 담기에는 너무 크면, 각 텐서가 어디에 저장되었는지 나타내는 인덱스 파일이 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: 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로 저장하는 기능은 비추천(deprecated) 상태이며, 향후 버전에서 제거될 예정입니다.
  • is_main_process (bool, 선택 사항) — 이 함수를 호출하는 프로세스가 메인 프로세스인지 여부입니다. TPU와 같은 분산 학습 환경에서 모든 프로세스에서 이 함수를 호출해야 할 때 유용합니다. 이 경우 레이스 조건을 방지하기 위해 메인 프로세스에서만 is_main_process=True로 설정해야 합니다. 기본값은 True입니다.
  • shared_tensors_to_discard (list[str], 선택 사항) — 공유 텐서를 저장할 때 제거할 텐서 이름 목록입니다. 제공되지 않고 공유 텐서가 감지되면, 알파벳 순서로 첫 번째 이름을 제거합니다.

model state dictionary를 디스크에 저장하며, 샤딩 및 공유 텐서 문제를 처리합니다.

PyTorch 모델을 직접 저장하려면 save_torch_model()도 참조하세요.

텐서 공유에 대한 자세한 내용은 이 가이드를 확인하세요.

모델 상태 사전은 각 샤드가 주어진 크기보다 작도록 샤드로 분할됩니다. 샤드는 save_directory에 주어진 filename_pattern으로 저장됩니다. 모델이 단일 샤드에 담기에는 너무 크면, 각 텐서가 어디에 저장되었는지 나타내는 인덱스 파일이 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: 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 객체로, 샤드와 이를 가져오기 위한 인덱스를 포함합니다.

model state dictionary를 여러 개의 샤드로 분할하여 각 샤드의 크기가 주어진 크기보다 작도록 합니다.

state_dict의 키 순서를 반복하여 샤드가 결정됩니다. 전달된 최대 크기에 각 샤드가 최대한 가깝도록 최적화되지 않습니다. 예를 들어, 제한이 10GB이고 텐서 크기가 [6GB, 6GB, 2GB, 6GB, 2GB, 2GB]인 경우, 샤드는 [6GB], [6+2GB], [6+2+2GB]로 나뉘며 [6+2+2GB], [6+2GB], [6GB]로 나뉘지 않습니다.

모델 상태 사전(state dictionary)을 디스크에 저장하려면 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: dictget_storage_size: Callablefilename_pattern: strget_storage_id: Callable = <function <lambda> at 0x7f9673968670>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]], 선택 사항) — 텐서 스토리지에 대한 고유 식별자를 반환하는 함수입니다. 서로 다른 여러 텐서가 동일한 기반 스토리지를 공유할 수 있습니다. 이 식별자는 해당 텐서의 스토리지가 존재하는 동안 고유하고 일정하게 유지되는 것이 보장됩니다. 수명이 겹치지 않는 두 텐서 스토리지는 동일한 id를 가질 수 있습니다.
  • filename_pattern (str, 선택 사항) — 모델이 저장될 파일 이름을 생성하는 패턴입니다. 패턴은 filename_pattern.format(suffix=...)으로 포맷할 수 있는 문자열이어야 하며 suffix 키워드를 포함해야 합니다.
  • max_shard_size (int 또는 str, 선택 사항) — 바이트 단위로 각 샤드의 최대 크기입니다. 기본값은 5GB입니다.

반환값

StateDictSplit

StateDictSplit 객체로, 샤드와 이를 가져오기 위한 인덱스를 포함합니다.

model state dictionary를 여러 개의 샤드로 분할하여 각 샤드의 크기가 주어진 크기보다 작도록 합니다.

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 체크포인트를 로드하려면 명시적으로 선택해야 합니다. unpickling은 로드 시점에 임의의 코드를 실행하기 때문입니다:

>>> 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의 제한된 unpickler를 사용합니다. 텐서 이외의 객체(옵티마이저 상태, 전체 모델 등)를 포함하는 체크포인트의 경우 safe=False와 함께 weights_only=False가 필요합니다.

load_torch_model

huggingface_hub.load_torch_model

< >

( 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) — 모델 상태 딕셔너리의 키가 체크포인트의 키와 정확히 일치하는지 엄격하게 강제할지 여부입니다. 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인 경우, PyTorch의 제한된 unpickler를 사용하여 옵티마이저 상태 및 기타 메타데이터 없이 모델 가중치만 로드합니다. pickle 체크포인트에 임의의 Python 객체를 허용하려면 False로 설정하십시오(이 경우 로드 시점에 임의의 코드가 실행됩니다). 제한된 unpickler가 없으며 항상 제한 없이 unpickle하는 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 필드를 가진 named tuple입니다.

  • missing_keys는 모델에는 있지만 체크포인트에는 없는 키, 즉 누락된 키를 포함하는 str 목록입니다.
  • unexpected_keys는 체크포인트에는 있지만 모델에는 없는 키, 즉 예상치 못한 키를 포함하는 str 목록입니다.

다음 예외가 발생합니다

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: 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의 제한된 unpickler를 사용하여 옵티마이저 상태 및 기타 메타데이터 없이 모델 가중치만 로드합니다. pickle 체크포인트에 임의의 Python 객체를 허용하려면 False로 설정하십시오(이 경우 로드 시점에 임의의 코드가 실행됩니다). 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를 나타내면 해당 형식으로, 그렇지 않으면 torch.load로 로드됩니다.

반환값

Union[dict[str, "torch.Tensor"], Any]

로드된 체크포인트.

  • safetensors 파일의 경우: 매개변수 이름과 텐서를 매핑하는 사전(dictionary)을 항상 반환합니다.
  • pickle 파일의 경우: pickled된 모든 Python 객체를 반환합니다(일반적으로 state dict이지만, 전체 모델, 옵티마이저 상태 또는 기타 Python 객체일 수도 있습니다).

다음 예외가 발생합니다

FileNotFoundError 또는 ImportError 또는 OSError 또는 ValueError

  • FileNotFoundError — 체크포인트 파일이 존재하지 않는 경우.
  • ImportError — .safetensors 파일 또는 PyTorch 체크포인트를 로드할 때 각각 safetensors 또는 torch가 설치되어 있지 않은 경우.
  • OSError — 체크포인트 파일 형식이 잘못되었거나 git-lfs 파일이 올바르게 다운로드되지 않은 경우.
  • ValueError — 체크포인트 파일 경로가 비어 있거나 유효하지 않은 경우, 또는 safe=True 상태에서 safetensors로 역직렬화할 수 없는 경우 (safe=False를 전달하여 pickle 체크포인트를 허용할 수 있습니다).

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)

Tensor 헬퍼

get_torch_storage_id

huggingface_hub.get_torch_storage_id

< >

( tensor: torch.Tensor )

tensor 저장소에 고유 식별자를 반환합니다.

여러 서로 다른 tensor가 동일한 기본 저장소(storage)를 공유할 수 있습니다. 이 식별자는 해당 tensor의 저장소가 존재하는 동안 유일하고 일정하게 유지됩니다. 수명이 겹치지 않는 두 tensor 저장소는 동일한 id를 가질 수 있습니다. meta tensor의 경우, 동일한 저장소를 공유하는지 여부를 알 수 없으므로 None을 반환합니다.

https://github.com/huggingface/transformers/blob/1ecf5f7c982d761b4daaa96719d162c324187c64/src/transformers/pytorch_utils.py#L278에서 가져왔습니다.

get_torch_storage_size

Update on GitHub