Hub Python Library documentation
Serialización
Serialización
huggingface_hub ofrece utilidades para guardar y cargar los pesos de modelos de ML de manera estandarizada. Esta parte de la biblioteca aún está en desarrollo y se mejorará en futuras versiones. El objetivo es armonizar cómo se guardan y cargan los pesos en todo el Hub, tanto para eliminar la duplicación de código entre bibliotecas como para establecer convenciones consistentes.
Formato de archivo DDUF
DDUF es un formato de archivo diseñado para modelos de difusión. Permite guardar toda la información necesaria para ejecutar un modelo en un único archivo. Este trabajo se inspira en el formato GGUF. huggingface_hub proporciona utilidades para guardar y cargar archivos DDUF, garantizando que se respete el formato de archivo.
Esta es una versión muy temprana del parser. La API y la implementación pueden evolucionar en un futuro próximo.
El analizador actualmente realiza muy poca validación. Para más detalles sobre el formato de archivo, consulta https://github.com/huggingface/huggingface.js/tree/main/packages/dduf.
Cómo escribir un archivo DDUF
A continuación se muestra cómo exportar una carpeta que contiene diferentes partes de un modelo de difusión usando 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")Para mayor flexibilidad, puede usar export_entries_as_dduf() y pasar una lista de archivos para incluirlos en el archivo DDUF final:
# 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
... ]
... )El parámetro entries también admite pasar un iterable de rutas o bytes. Esto puede resultar útil si tienes un modelo cargado y quieres serializarlo directamente en un archivo DDUF, en lugar de tener que serializar primero cada componente en el disco y luego como archivo DDUF. A continuación se muestra un ejemplo de cómo se puede serializar un StableDiffusionPipeline como 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))Nota: en la práctica, diffusers proporciona un método para serializar directamente un pipeline en un archivo DDUF. El fragmento anterior se muestra solo como ejemplo.
Cómo leer un archivo 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)Funciones auxiliares
huggingface_hub.export_entries_as_dduf
< source >( dduf_path: str | os.PathLikeentries: Iterable )
Parámetros
- dduf_path (
stroos.PathLike) — La ruta al archivo DDUF que se escribirá. - entries (
Iterable[tuple[str, Union[str, Path, bytes]]]) — Un iterable de entradas para escribir en el archivo DDUF. Cada entrada es una tupla con el nombre del archivo y su contenido. El nombre del archivo debe ser la ruta del archivo dentro del archivo DDUF. El contenido puede ser una cadena o un pathlib.Path que represente una ruta a un archivo en el disco local, o directamente el contenido como bytes.
Lanza
-
- —
DDUFExportError: Si algo sale mal durante la exportación (p. ej. nombre de entrada no válido, falta de ‘model_index.json’, etc.).
- —
Escribe un archivo DDUF a partir de un iterable de entradas.
Este es un auxiliar de nivel inferior que export_folder_as_dduf() permite mayor flexibilidad al serializar datos. En particular, no es necesario guardar los datos en el disco antes de exportarlos en el archivo DDUF.
Ejemplo:
# 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 )
Exporta una carpeta como un archivo DDUF.
AUtiliza export_entries_as_dduf() internamente.
huggingface_hub.read_dduf_file
< source >( dduf_path: os.PathLike | str ) → dict[str, DDUFEntry]
Parámetros
Devuelve
dict[str, DDUFEntry]
Un diccionario de DDUFEntry indexado por nombre de archivo.
Lanza
-
- —
DDUFCorruptedFileError: Si el archivo DDUF está corrupto (es decir, no sigue el formato DDUF).
- —
Lee un archivo DDUF y devuelve un diccionario de entradas.
Solo se lee los metadatos, los datos no se cargan en memoria.
Ejemplo:
>>> 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)clase huggingface_hub.DDUFEntry
< source >( filename: strlength: intoffset: intdduf_path: Path )
Objeto que representa una entrada de archivo en un archivo DDUF.
Consulte read_dduf_file() para saber cómo leer un archivo DDUF.
Abra el archivo como un archivo mapeado en memoria.
Útil para cargar safetensors directamente desde el archivo.
Lee el archivo como texto.
Útil para entradas ‘.txt’ y ‘.json’.
Errores
Excepción base para errores relacionados con el formato DDUF.
Excepción lanzada cuando el archivo DDUF está corrupto.
Excepción base para errores durante la exportación de DDUF.
Excepción lanzada cuando el nombre de la entrada no es válido.
Guardar tensores
El principal helper del módulo serialization toma un nn.Module de torch como entrada y lo guarda en disco. Maneja la lógica para guardar tensores compartidos (ver explicación de safetensors) así como la lógica para dividir el state dictionary en shards, usando split_torch_state_dict_into_shards() por debajo. Por el momento, solo se admite el framework torch.
Si desea guardar un diccionario de estado (p. ej., un mapeo entre nombres de capas y los tensores relacionados) en lugar de un nn.Module, puede usar save_torch_state_dict(), que ofrece las mismas funcionalidades. Esto es útil, por ejemplo, si quiere aplicar lógica personalizada al diccionario de estado antes de guardarlo.
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 )
Parámetros
- model (
torch.nn.Module) — El modelo que se guardará en el disco. - save_directory (
stroPath) — El directorio en el que se guardará el modelo. - filename_pattern (
str, opcional) — El patrón para generar los nombres de los archivos en los que se guardará el modelo. El patrón debe ser una cadena de caracteres que se pueda formatear confilename_pattern.format(suffix=...)y debe contener la palabra clavesuffixPor defecto es"model{suffix}.safetensors"opytorch_model{suffix}.bindependiendo del parámetrosafe_serialization. - force_contiguous (
boolean, opcional) — Fuerza que el state_dict se guarde como tensores contiguos. Esto no afecta la corrección del modelo, pero podría cambiar el rendimiento si el diseño del tensor se eligió específicamente por esa razón. El valor predeterminado esTrue. - max_shard_size (
intostr, opcional) — El tamaño máximo de cada shard, en bytes. El valor predeterminado es 5GB. - metadata (
dict[str, str], opcional) — Información adicional para guardar junto con el modelo. Se añadirá metadatos para cada tensor descartado. Esta información no será suficiente para recuperar toda la estructura compartida, pero podría ayudar a comprender las cosas. - safe_serialization (
bool, opcional) — Si se debe guardar como safetensors, que es el comportamiento predeterminado. Si esFalse, los fragmentos se guardan como pickle. Se recomienda la serialización segura por razones de seguridad. Guardar como pickle está en desuso y se eliminará en una versión futura. - is_main_process (
bool, opcional) — Indica si el proceso que realiza esta llamada es el proceso principal o no. Útil en el entrenamiento distribuido, como con TPUs, cuando es necesario llamar a esta función desde todos los procesos. En este caso, establezcais_main_process=Truesolo en el proceso principal para evitar condiciones de carrera. El valor predeterminado es True. - shared_tensors_to_discard (
list[str], opcional) — Lista de nombres de tensores que se descartarán al guardar los tensores compartidos. Si no se proporciona y se detectan tensores compartidos, se descartará el primer nombre en orden alfabético.
Guarda un modelo de torch dado en el disco, gestionando problemas de fragmentación y tensores compartidos.
Véase también save_torch_state_dict() para guardar un state dict con más flexibilidad.
Para obtener más información sobre el intercambio de tensores, consulte esta guía.
El diccionario de estado del modelo se divide en fragmentos para que cada fragmento sea menor que un tamaño dado. Los fragmentos se
guardan en save_directory con el filename_pattern indicado. Si el modelo es demasiado grande para caber en un solo fragmento,
se guarda un archivo de índice en save_directory para indicar dónde se guarda cada tensor. Este helper utiliza split_torch_state_dict_into_shards() internamente. Si safe_serialization es True, los fragmentos se guardan como
safetensors (el valor predeterminado). De lo contrario, los fragmentos se guardan como pickle.
Antes de guardar el modelo, el save_directory se limpia de cualquier archivo de fragmento anterior.
Si uno de los tensores del modelo es mayor que
max_shard_size, terminará en su propio fragmento, que tendrá un tamaño mayor quemax_shard_size.
Si tu modelo es un
transformers.PreTrainedModel, debes pasarmodel._tied_weights_keyscomoshared_tensors_to_discardpara manejar correctamente el guardado de tensores compartidos. Esto garantiza que se descarten los tensores duplicados correctos durante el guardado.
Ejemplo:
>>> 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 )
Parámetros
- state_dict (
dict[str, torch.Tensor]) — El diccionario de estado a guardar. - save_directory (
stroPath) — El directorio en el que se guardará el modelo. - filename_pattern (
str, opcional) — El patrón para generar los nombres de los archivos en los que se guardará el modelo. El patrón debe ser una cadena de caracteres que se pueda formatear confilename_pattern.format(suffix=...)y debe contener la palabra clavesuffixPor defecto es"model{suffix}.safetensors"opytorch_model{suffix}.bindependiendo del parámetrosafe_serialization. - force_contiguous (
boolean, opcional) — Fuerza que el state_dict se guarde como tensores contiguos. Esto no afecta la corrección del modelo, pero podría cambiar el rendimiento si el diseño del tensor se eligió específicamente por esa razón. El valor predeterminado esTrue. - max_shard_size (
intostr, opcional) — El tamaño máximo de cada shard, en bytes. El valor predeterminado es 5GB. - metadata (
dict[str, str], opcional) — Información adicional para guardar junto con el modelo. Se añadirá metadatos para cada tensor descartado. Esta información no será suficiente para recuperar toda la estructura compartida, pero podría ayudar a comprender las cosas. - safe_serialization (
bool, opcional) — Si se debe guardar como safetensors, que es el comportamiento predeterminado. Si esFalse, los fragmentos se guardan como pickle. Se recomienda la serialización segura por razones de seguridad. Guardar como pickle está en desuso y se eliminará en una versión futura. - is_main_process (
bool, opcional) — Indica si el proceso que realiza esta llamada es el proceso principal o no. Útil en el entrenamiento distribuido, como con TPUs, cuando es necesario llamar a esta función desde todos los procesos. En este caso, establezcais_main_process=Truesolo en el proceso principal para evitar condiciones de carrera. El valor predeterminado es True. - shared_tensors_to_discard (
list[str], opcional) — Lista de nombres de tensores que se descartarán al guardar los tensores compartidos. Si no se proporciona y se detectan tensores compartidos, se descartará el primer nombre en orden alfabético.
Guardar un diccionario de estado de modelo en el disco, gestionando problemas de fragmentación y tensores compartidos.
Véase también save_torch_model() para guardar directamente un modelo de PyTorch.
Para obtener más información sobre el intercambio de tensores, consulte esta guía.
El diccionario de estado del modelo se divide en fragmentos para que cada fragmento sea menor que un tamaño dado. Los fragmentos se
guardan en save_directory con el filename_pattern indicado. Si el modelo es demasiado grande para caber en un solo fragmento,
se guarda un archivo de índice en save_directory para indicar dónde se guarda cada tensor. Este helper utiliza split_torch_state_dict_into_shards() internamente. Si safe_serialization es True, los fragmentos se guardan como
safetensors (el valor predeterminado). De lo contrario, los fragmentos se guardan como pickle.
Antes de guardar el modelo, el save_directory se limpia de cualquier archivo de fragmento anterior.
Si uno de los tensores del modelo es mayor que
max_shard_size, terminará en su propio fragmento, que tendrá un tamaño mayor quemax_shard_size.
Si tu modelo es un
transformers.PreTrainedModel, debes pasarmodel._tied_weights_keyscomoshared_tensors_to_discardpara manejar correctamente el guardado de tensores compartidos. Esto garantiza que se descarten los tensores duplicados correctos durante el guardado.
Ejemplo:
>>> 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")El módulo serialization también contiene utilidades de bajo nivel para dividir un diccionario de estado en varios fragmentos, creando un índice adecuado en el proceso. Estas utilidades están disponibles para tensores de torch y están diseñadas para extenderse fácilmente a cualquier otro marco de aprendizaje automático.
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
Parámetros
- state_dict (
dict[str, torch.Tensor]) — El diccionario de estado a guardar. - filename_pattern (
str, opcional) — El patrón para generar los nombres de los archivos en los que se guardará el modelo. El patrón debe ser una cadena de caracteres que pueda formatearse confilename_pattern.format(suffix=...)y debe contener la palabra clavesuffix. Por defecto es"model{suffix}.safetensors". - max_shard_size (
intostr, opcional) — El tamaño máximo de cada shard, en bytes. El valor predeterminado es 5GB.
Devuelve
StateDictSplit
Un objeto StateDictSplit que contiene los fragmentos y el índice para recuperarlos.
Divide el diccionario de estado de un modelo en fragmentos para que cada fragmento sea menor que un tamaño dado.
Los fragmentos se determinan iterando sobre el state_dict en el orden de sus claves. No se realiza ninguna optimización para que cada fragmento se acerque lo máximo posible al tamaño máximo indicado. Por ejemplo, si el límite es de 10 GB y tenemos tensores de tamaños [6 GB, 6 GB, 2 GB, 6 GB, 2 GB, 2 GB], se fragmentarán como [6 GB], [6+2 GB], [6+2+2 GB] y no como [6+2+2 GB], [6+2 GB], [6 GB].
Para guardar un diccionario de estado de modelo en el disco, consulte save_torch_state_dict(). Este helper utiliza
split_torch_state_dict_into_shardsinternamente.
Si uno de los tensores del modelo es mayor que
max_shard_size, terminará en su propio fragmento, que tendrá un tamaño mayor quemax_shard_size.
Ejemplo:
>>> 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
Esta es la fábrica subyacente de la que se deriva cada helper específico de un framework. En la práctica, no se espera que utilices esta fábrica directamente, a menos que necesites adaptarla a un framework que aún no está soportado. Si ese es el caso, infórmalo abriendo un nuevo issue en el repositorio 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 0x7f2781b88b80>max_shard_size: int | str = '5GB' ) → StateDictSplit
Parámetros
- state_dict (
dict[str, Tensor]) — El diccionario de estado a guardar. - get_storage_size (
Callable[[Tensor], int]) — Una función que devuelve el tamaño de un tensor al guardarlo en disco, en bytes. - get_storage_id (
Callable[[Tensor], Optional[Any]], opcional) — Una función que devuelve un identificador único para el almacenamiento de un tensor. Varios tensores diferentes pueden compartir el mismo almacenamiento subyacente. Este identificador está garantizado para ser único y constante para el almacenamiento de este tensor durante su ciclo de vida. Dos almacenamientos de tensores con ciclos de vida no superpuestos pueden tener el mismo id. - filename_pattern (
str, opcional) — El patrón para generar los nombres de los archivos en los que se guardará el modelo. El patrón debe ser una cadena de caracteres que pueda formatearse confilename_pattern.format(suffix=...)y debe contener la palabra clavesuffix - max_shard_size (
intostr, opcional) — El tamaño máximo de cada shard, en bytes. El valor predeterminado es 5GB.
Devuelve
StateDictSplit
Un objeto StateDictSplit que contiene los fragmentos y el índice para recuperarlos.
Divide el diccionario de estado de un modelo en fragmentos para que cada fragmento sea menor que un tamaño dado.
Los fragmentos se determinan iterando sobre el state_dict en el orden de sus claves. No se realiza ninguna optimización para que cada fragmento se acerque lo máximo posible al tamaño máximo indicado. Por ejemplo, si el límite es de 10 GB y tenemos tensores de tamaños [6 GB, 6 GB, 2 GB, 6 GB, 2 GB, 2 GB], se fragmentarán como [6 GB], [6+2 GB], [6+2+2 GB] y no como [6+2+2 GB], [6+2 GB], [6 GB].
Si uno de los tensores del modelo es mayor que
max_shard_size, terminará en su propio fragmento, que tendrá un tamaño mayor quemax_shard_size.
Carga de tensores
Los helpers de carga admiten puntos de control de archivo único y fragmentados, en formato safetensors o pickle. load_torch_model() toma un nn.Module y una ruta de punto de control (ya sea un archivo único o un directorio) como entrada y carga los pesos en el modelo.
Ambos helpers tienen
safe=Truepor defecto, lo que significa que el checkpoint siempre se deserializa con el cargador de safetensors, sin importar el nombre del archivo. Cargar un checkpoint de pickle requiere habilitarlo explícitamente, porque deserializar con pickle ejecuta código arbitrario en el momento de la carga:>>> 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)La ruta de pickle además tiene
weights_only=Truepor defecto, es decir, el deserializador restringido de torch. Un punto de control que contenga objetos que no sean tensores (estados del optimizador, un modelo completo, …) requiereweights_only=Falseademás desafe=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
Parámetros
- model (
torch.nn.Module) — El modelo en el que se cargará el punto de control. - checkpoint_path (
stroos.PathLike) — Ruta al archivo de checkpoint o al directorio que contiene los checkpoints. - strict (
bool, opcional, valor predeterminadoFalse) — Si se debe aplicar estrictamente que las claves del state dict del modelo coincidan con las claves del checkpoint. Al igual que contorch.nn.Module.load_state_dict, la comprobación se realiza una vez que el checkpoint se ha cargado, por lo que el modelo puede haberse actualizado parcialmente cuando se lanza el error. - safe (
bool, opcional, valor predeterminadoTrue) — Sisafees True, se cargarán los archivos safetensors. Sisafees False, la función intentará primero cargar los archivos safetensors si están disponibles; de lo contrario, recurrirá a cargar los archivos pickle. Unfilename_patternque no describa archivos safetensors se rechaza cuandosafe=True. - weights_only (
bool, opcional, valor predeterminadoTrue) — Si es True, solo carga los pesos del modelo sin los estados del optimizador ni otra metadatos, utilizando el deserializador restringido de torch. Establezca False para permitir objetos Python arbitrarios en un punto de control pickle (esto ejecuta código arbitrario en el momento de la carga). No tiene efecto en PyTorch < 1.13, que no tiene un deserializador restringido y siempre deserializa sin restricciones. - map_location (
strotorch.device, opcional) — Un objetotorch.device, una cadena o un diccionario que especifica cómo remapear las ubicaciones de almacenamiento. Indica la ubicación donde se deben cargar todos los tensores. - mmap (
bool, opcional, valor predeterminadoFalse) — Si se debe usar la carga de archivos con memoria mapeada. El mapeo de memoria puede mejorar el rendimiento de carga para modelos grandes en PyTorch >= 2.1.0 con puntos de control basados en zipfile. - filename_pattern (
str, opcional) — El patrón para buscar el archivo de índice. El patrón debe ser una cadena que pueda formatearse confilename_pattern.format(suffix=...)y debe contener la palabra clavesuffixPor defecto es"model{suffix}.safetensors".
Devuelve
NamedTuple
Una tupla nombrada con los campos missing_keys y unexpected_keys.
missing_keyses una lista de str que contiene las claves faltantes, es decir, las claves que están en el modelo pero no en el checkpoint.unexpected_keyses una lista de str que contiene las claves inesperadas, es decir, las claves que están en el checkpoint pero no en el modelo.
Lanza
FileNotFoundError o ImportError o ValueError
FileNotFoundError— Si el archivo o directorio de punto de control no existe.ImportError— Si no se ha instalado safetensors o torch al intentar cargar un archivo .safetensors o un checkpoint de PyTorch, respectivamente.ValueError— Si no se puede determinar el formato del punto de control.
Cargar un punto de control en un modelo, gestionando tanto puntos de control fragmentados como no fragmentados.
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]
Parámetros
- checkpoint_file (
stroos.PathLike) — Ruta al archivo de checkpoint que se va a cargar. Puede ser un checkpoint de safetensors o pickle (.bin). - map_location (
strotorch.device, opcional) — Un objetotorch.device, una cadena o un diccionario que especifica cómo remapear las ubicaciones de almacenamiento. Indica la ubicación donde se deben cargar todos los tensores. - weights_only (
bool, opcional, valor predeterminadoTrue) — Si es True, solo carga los pesos del modelo sin los estados del optimizador ni otra metadatos, utilizando el deserializador restringido de torch. Establezca False para permitir objetos Python arbitrarios en un punto de control pickle (esto ejecuta código arbitrario en el momento de la carga). No tiene efecto al cargar archivos safetensors, ni en PyTorch < 1.13, que no tiene un deserializador restringido: esas versiones siempre deserializan sin restricciones y se registra una advertencia. - mmap (
bool, opcional, valor predeterminadoFalse) — Si se debe usar la carga de archivos con memoria mapeada. El mapeo de memoria puede mejorar el rendimiento de carga para modelos grandes en PyTorch >= 2.1.0 con puntos de control basados en zipfile. No tiene efecto al cargar archivos safetensors, ya que la bibliotecasafetensorsutiliza el mapeo de memoria de forma predeterminada. - safe (
bool, opcional, valor predeterminadoTrue) — Si es True, el punto de control siempre se carga como safetensors, independientemente de su nombre. Cualquier otro formato (p. ej., un archivo pickle.bin) genera unValueErroren lugar de deserializarse: los puntos de control pickle pueden ejecutar código arbitrario en el momento de la carga. Si es False, el archivo se carga como safetensors si su nombre lo indica, y contorch.loaden caso contrario.
Devuelve
Union[dict[str, "torch.Tensor"], Any]
El punto de control cargado.
- Para los archivos safetensors: siempre devuelve un diccionario que mapea los nombres de los parámetros a los tensores.
- Para archivos pickle: devuelve cualquier objeto de Python que haya sido serializado con pickle (comúnmente un state dict, pero podría ser un modelo completo, el estado de un optimizador o cualquier otro objeto de Python).
Lanza
FileNotFoundError o ImportError o OSError o ValueError
FileNotFoundError— Si el archivo de punto de control no existe.ImportError— Si no se ha instalado safetensors o torch al intentar cargar un archivo .safetensors o un checkpoint de PyTorch, respectivamente.OSError— Si el formato del archivo de checkpoint no es válido o si los archivos de git-lfs no se descargan correctamente.ValueError— Si la ruta del archivo de checkpoint está vacía o es inválida, o si no se puede deserializar como safetensors mientrassafe=True(pasesafe=Falsepara permitir checkpoints de pickle).
Carga un archivo de punto de control, gestionando tanto el formato de puntos de control de safetensors como el de pickle.
Ejemplo:
>>> 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)Utilidades de tensores
get_torch_storage_id
Devuelve un identificador único para el almacenamiento de un tensor.
Varios tensores distintos pueden compartir el mismo almacenamiento subyacente. Este identificador está garantizado para ser único y constante para el almacenamiento de este tensor durante su ciclo de vida. Dos almacenamientos de tensores con ciclos de vida no superpuestos pueden tener el mismo id. En el caso de los meta tensores, devolvemos None ya que no podemos determinar si comparten el mismo almacenamiento.