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)класс huggingface_hub.DDUFEntry
< source >( filename: strlength: intoffset: intdduf_path: Path )
Объект, представляющий запись файла в файле DDUF.
См. read_dduf_file() для получения информации о чтении файла DDUF.
Откройте файл как файл с отображением в памяти (memory-mapped file).
Полезно для загрузки safetensors непосредственно из файла.
Прочитайте файл как текст.
Полезно для записей в форматах ‘.txt’ и ‘.json’.
Ошибки
Базовое исключение для ошибок, связанных с форматом DDUF.
Исключение, возникающее при повреждении файла DDUF.
Базовое исключение для ошибок при экспорте DDUF.
Исключение, возникающее при некорректном имени элемента.
Сохранение тензоров
Основная вспомогательная функция модуля serialization принимает в качестве входного параметра torch nn.Module и сохраняет его на диск. Она обрабатывает логику сохранения общих тензоров (см. описание safetensors), а также логику разбиения state dictionary на фрагменты, используя под капотом 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По умолчанию используется"model{suffix}.safetensors"илиpytorch_model{suffix}.binв зависимости от параметраsafe_serialization. - force_contiguous (
boolean, необязательный) — Принудительно сохраняет state_dict в виде непрерывных тензоров. Это не влияет на корректность модели, но может потенциально изменить производительность, если расположение тензора было выбрано специально для этой цели. По умолчаниюTrue. - max_shard_size (
intилиstr, необязательный) — Максимальный размер каждого шарда в байтах. По умолчанию — 5 ГБ. - 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() для сохранения 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
< 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По умолчанию используется"model{suffix}.safetensors"илиpytorch_model{suffix}.binв зависимости от параметраsafe_serialization. - force_contiguous (
boolean, необязательный) — Принудительно сохраняет state_dict в виде непрерывных тензоров. Это не влияет на корректность модели, но может потенциально изменить производительность, если расположение тензора было выбрано специально для этой цели. По умолчаниюTrue. - max_shard_size (
intилиstr, необязательный) — Максимальный размер каждого шарда в байтах. По умолчанию — 5 ГБ. - 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.
Для получения дополнительной информации о шаринге тензоров ознакомьтесь с этим руководством.
Словарь состояния модели разбивается на фрагменты так, чтобы каждый фрагмент был меньше заданного размера. Фрагменты сохраняются в 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
< 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, необязательный) — Максимальный размер каждого шарда в байтах. По умолчанию — 5 ГБ.
Возвращает
StateDictSplit
Объект StateDictSplit, содержащий фрагменты и индекс для их получения.
Разделите словарь состояния модели на фрагменты так, чтобы каждый фрагмент был меньше заданного размера.
Шарды определяются путём итерации по state_dict в порядке его ключей. Никакая оптимизация для того, чтобы каждый шард был как можно ближе к максимальному размеру, не выполняется. Например, если лимит составляет 10 ГБ, а размеры тензоров равны [6 ГБ, 6 ГБ, 2 ГБ, 6 ГБ, 2 ГБ, 2 ГБ], они будут разбиты на [6 ГБ], [6+2 ГБ], [6+2+2 ГБ], а не на [6+2+2 ГБ], [6+2 ГБ], [6 ГБ].
Чтобы сохранить словарь состояния модели на диск, см. 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
Это базовый фабрика, из которой производятся все специфичные для фреймворков помощники. На практике вы не должны использовать эту фабрику напрямую, за исключением случаев, когда вам необходимо адаптировать её под фреймворк, который пока не поддерживается. Если это ваш случай, пожалуйста, сообщите нам, создав новую issue в репозитории huggingface_hub.
huggingface_hub.split_state_dict_into_shards_factory
< source >( 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]], необязательный) — Функция, возвращающая уникальный идентификатор хранилища тензора. Несколько различных тензоров могут использовать одно и то же базовое хранилище. Этот идентификатор гарантированно уникален и неизменен для хранилища данного тензора в течение всего времени его существования. Два хранилища тензоров с непересекающимися временами жизни могут иметь одинаковый идентификатор. - filename_pattern (
str, необязательный) — Шаблон для генерации имён файлов, в которых будет сохранена модель. Шаблон должен быть строкой, которую можно форматировать с помощьюfilename_pattern.format(suffix=...)и которая должна содержать ключевое словоsuffix - max_shard_size (
intилиstr, необязательный) — Максимальный размер каждого шарда в байтах. По умолчанию — 5 ГБ.
Возвращает
StateDictSplit
Объект StateDictSplit, содержащий фрагменты и индекс для их получения.
Разделите словарь состояния модели на фрагменты так, чтобы каждый фрагмент был меньше заданного размера.
Шарды определяются путём итерации по state_dict в порядке его ключей. Никакая оптимизация для того, чтобы каждый шард был как можно ближе к максимальному размеру, не выполняется. Например, если лимит составляет 10 ГБ, а размеры тензоров равны [6 ГБ, 6 ГБ, 2 ГБ, 6 ГБ, 2 ГБ, 2 ГБ], они будут разбиты на [6 ГБ], [6+2 ГБ], [6+2+2 ГБ], а не на [6+2+2 ГБ], [6+2 ГБ], [6 ГБ].
Если один из тензоров модели больше
max_shard_size, он окажется в собственном шарде, размер которого превыситmax_shard_size.
Загрузка тензоров
Вспомогательные функции загрузки поддерживают как одиночные, так и разбитые на части контрольные точки в формате safetensors или pickle. load_torch_model() принимает на вход nn.Module и путь к контрольной точке (файл или каталог) и загружает веса в модель.
Оба хелпера по умолчанию используют
safe=True, что означает, что чекпоинт всегда десериализуется с помощью загрузчика safetensors — независимо от имени файла. Загрузка чекпоинта в формате pickle требует явного включения, так как распаковка 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.filename_pattern, не описывающий файлы safetensors, отклоняется, когдаsafe=True. - weights_only (
bool, необязательный, по умолчаниюTrue) — Если True, загружаются только веса модели без состояний оптимизатора и других метаданных, с использованием ограниченного unpickler’а torch. Установите False, чтобы разрешить произвольные объекты Python в pickle-чекпоинте (это приводит к выполнению произвольного кода при загрузке). Не влияет на PyTorch < 1.13, в котором нет ограниченного unpickler’а и распаковка всегда выполняется без ограничений. - map_location (
strилиtorch.device, необязательный) — Объектtorch.device, строка или словарь, задающие способ переназначения мест хранения. Указывает расположение, в которое следует загрузить все тензоры. - mmap (
bool, необязательный, по умолчаниюFalse) — Использовать ли загрузку файла через отображение в память (memory mapping). Отображение в память может повысить производительность загрузки для больших моделей в 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. Установите False, чтобы разрешить произвольные объекты Python в pickle-чекпоинте (это приводит к выполнению произвольного кода при загрузке). Не влияет при загрузке файлов safetensors, а также на PyTorch < 1.13, в котором нет ограниченного разупаковщика — эти версии всегда разупаковывают без ограничений, и при этом записывается предупреждение. - mmap (
bool, необязательный, по умолчаниюFalse) — Использовать ли загрузку файла через отображение в память (memory mapping). Отображение в память может повысить производительность загрузки для больших моделей в 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: всегда возвращает словарь, сопоставляющий имена параметров с тензорами.
- Для файлов pickle: возвращает любой объект Python, который был сериализован (обычно это state dict, но это может быть и вся модель, состояние оптимизатора, или любой другой объект Python).
Исключения
FileNotFoundError или ImportError или OSError или ValueError
FileNotFoundError— если файл контрольной точки не существует.ImportError— Если при попытке загрузить файл .safetensors или чекпоинт PyTorch не установлены safetensors или torch соответственно.OSError— если формат файла контрольной точки недействителен или файлы git-lfs не были корректно загружены.ValueError— Если путь к файлу контрольной точки пуст или недействителен, или если его невозможно десериализовать как safetensors приsafe=True(передайте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)Вспомогательные функции для тензоров
get_torch_storage_id
Возвращает уникальный идентификатор хранилища тензора.
Несколько различных тензоров могут использовать одно и то же базовое хранилище. Этот идентификатор гарантированно уникален и неизменен для хранилища данного тензора на протяжении всего его жизненного цикла. Два хранилища тензоров с непересекающимися жизненными циклами могут иметь одинаковый id. В случае мета-тензоров мы возвращаем None, поскольку невозможно определить, разделяют ли они одно и то же хранилище.