Modelos

Esta página contiene la documentación de referencia de la API para crear y descubrir modelos de timm, sus configuraciones preentrenadas y los helpers de checkpoint / Hub, así como una descripción general de la interfaz compartida por los modelos de timm. Las arquitecturas individuales se tratan en los resúmenes de modelos y en las páginas de cada modelo.

Uso

Los modelos se crean por nombre con create_model. El nombre selecciona una arquitectura registrada en timm, opcionalmente con una etiqueta de preentrenamiento después de un . (resnet50.a1_in1k), o un modelo en el Hugging Face Hub o en una carpeta local mediante un prefijo de origen. Los argumentos de palabra clave que create_model no utiliza en sí mismo se pasan a través de la función de entrada del modelo al constructor de la clase del modelo, por lo que se puede establecer cualquier argumento específico del modelo.

import torch
import timm

# Random init (pretrained=False is the default). The tag selects the pretrained cfg, and the weights loaded with
# pretrained=True. Without a tag the architecture's default tag is used.
model = timm.create_model('vit_tiny_patch16_224.augreg_in21k_ft_in1k')
print(model.pretrained_cfg['mean'], model.pretrained_cfg['crop_pct'])  # (0.5, 0.5, 0.5) 0.9

# Model arguments pass through to the model class, a 10 class head, single channel input and regularization.
model = timm.create_model('resnet50', num_classes=10, in_chans=1, drop_rate=0.2, drop_path_rate=0.05)

# Discover models, filters are fnmatch wildcards.
timm.list_models('convnext*')  # architecture names
timm.list_pretrained('resnet50.*')  # architecture.tag names w/ pretrained weights
timm.list_models(module='vision_transformer', pretrained=True)

# Image preprocessing to match the model's pretrained cfg (input size, interpolation, crop, mean / std).
model = timm.create_model('vit_tiny_patch16_224.augreg_in21k_ft_in1k')
data_config = timm.data.resolve_model_data_config(model)
transform = timm.data.create_transform(**data_config, is_training=False)

x = torch.randn(2, 3, 224, 224)
features = model.forward_features(x)  # unpooled features, (2, 197, 192)
pre_logits = model.forward_head(features, pre_logits=True)  # pooled features, (2, 192)
logits = model.forward_head(features)  # same as model(x), (2, 1000)

model.reset_classifier(5)  # replace the classifier w/ a new 5 class one
model.set_input_size(img_size=(384, 384))  # resample the position embedding for a new resolution
logits = model(torch.randn(2, 3, 384, 384))  # (2, 5)

El script de entrenamiento (train.py) expone los argumentos comunes a través de estas banderas:

• --model selecciona el modelo.
• --pretrained carga los pesos preentrenados.
• --num-classes establece el tamaño del clasificador (num_classes).
• --in-chans establece los canales de entrada (in_chans).
• --gp establece el tipo de agrupación global (global_pool).
• --drop establece el dropout del clasificador (drop_rate).
• --drop-path establece la profundidad estocástica (drop_path_rate).
• --pretrained-path carga pesos preentrenados desde un archivo local con adaptación del clasificador / convolución de entrada (pretrained_cfg_overlay=dict(file=...)).
• --initial-checkpoint carga un checkpoint después de crear el modelo, sin adaptación (checkpoint_path).
• --model-kwargs key=value ... pasa cualquier otro argumento del modelo.

Consulte el quickstart para una introducción.

Nombres de modelos

El Hub y las fuentes locales leen la configuración incluso con pretrained=False; los pesos solo se cargan con pretrained=True. Un nombre sin prefijo de fuente que contenga /, \ o : lanza un ValueError en lugar de recurrir a un modelo del registro con el mismo nombre. hf_hub: es un alias obsoleto de hf-hub:.

Argumentos comunes

• pretrained carga los pesos de la configuración preentrenada, desde state_dict o file si están definidos, en caso contrario desde el Hub (hf_hub_id) cuando huggingface_hub está instalado, o desde la url. La adaptación del clasificador / convolución de entrada a continuación permite crear un modelo preentrenado para el ajuste fino en una sola llamada.
• num_classes define el tamaño del clasificador, 0 elimina el clasificador. El valor predeterminado proviene de la configuración preentrenada. Al cargar pesos preentrenados, un número de clases diferente descarta los pesos del clasificador (se dejan en la inicialización aleatoria).
• global_pool establece el tipo de agrupación; las opciones son específicas del modelo (p. ej. avg, max, avgmax, catavgmax para la mayoría de las convnets, token, avg, map para ViT), y '' desactiva la agrupación.
• in_chans establece los canales de entrada. El valor predeterminado proviene de la cfg preentrenada. Al cargar pesos preentrenados, los pesos de la primera conv se adaptan, si se admite, cuando este valor difiere del número de canales del checkpoint de origen (input_size[0] en la cfg preentrenada).
• drop_rate establece el dropout del clasificador, compatible con la mayoría de los modelos.
• drop_path_rate establece la profundidad estocástica, admitida por la mayoría de los modelos.
• img_size es un argumento de modelos de estilo ViT y otros con un tamaño de entrada fijo; están construidos para la configuración preentrenada input_size a menos que se establezca.
• dynamic_img_size=True en las familias ViT y EVA vuelve a muestrear el embedding de posición al tamaño de entrada en cada pasada hacia adelante.
• features_only (con out_indices, output_stride, feature_cls) envuelve el modelo como un backbone de características que devuelve mapas de características, consulte extracción de características.
• pretrained_cfg reemplaza la configuración preentrenada (una etiqueta, un diccionario o PretrainedCfg).
• pretrained_cfg_overlay sobrescribe algunos de sus campos, por ejemplo, pretrained_cfg_overlay=dict(file='weights.pth') carga pesos de pretrained=True desde un archivo local, con la adaptación del clasificador / conv de entrada anterior (como lo hace --pretrained-path en train.py).
• checkpoint_path carga un checkpoint con load_checkpoint después de que se crea el modelo. Es una carga estricta sin adaptación; se usan los pesos EMA si el checkpoint los tiene.
• cache_dir reemplaza la caché de descarga para los pesos preentrenados y las configuraciones del Hub.
• device se pasa a la clase del modelo, los parámetros se crean directamente en el dispositivo. device='meta' crea un modelo sin asignar memoria para sus pesos.
• dtype se pasa a la clase del modelo y los parámetros se crean directamente en ese dtype.
• scriptable establece las banderas de configuración de capas para el scripting de TorchScript.
• exportable establece las banderas de configuración de capas para la exportación / el trazado.
• no_jit establece las banderas de configuración de capa para deshabilitar las capas con guiones de jit.

Los argumentos fijados en None se descartan antes de llegar al modelo, por lo que se aplican los valores predeterminados del modelo. Los argumentos que la clase del modelo no acepta provocan un TypeError en la mayoría de los modelos; algunos ignoran los argumentos desconocidos.

Configuración preentrenada

Las arquitecturas tienen uno o más cfgs preentrenados, uno por etiqueta, registrados con el modelo. Un cfg describe los pesos (ubicación, licencia, origen) y cómo se entrenó el modelo y cómo debería evaluarse con ellos. Los modelos creados por create_model tienen su cfg como un dict en model.pretrained_cfg (model.default_cfg es un alias), las versiones del registro se devuelven como instancias de PretrainedCfg por get_pretrained_cfg.

Los campos de entrada son lo que resolve_data_config / resolve_model_data_config resuelven en una configuración de datos para las transformaciones y el cargador del conjunto de datos. La cfg describe el punto de control de origen y no se actualiza cuando el modelo se crea o se modifica de manera diferente (num_classes, in_chans, img_size, set_input_size()). Los canales de entrada actuales, el tamaño y la restricción de tamaño fijo del modelo se rastrean por separado: get_model_input_config los devuelve, y resolve_input_data_config resuelve una configuración de datos en relación con ellos (como lo hacen los scripts de entrenamiento, validación e inferencia). get_model_args devuelve los argumentos para reconstruir el modelo tal como está ahora, que save_for_hf escribe en la configuración exportada. Consulte adapt and export models para más detalles.

Interfaz del modelo

Los modelos de timm comparten una interfaz común más allá de forward(): la separación en forward_features() / forward_head(), el restablecimiento del clasificador, las características intermedias, el redimensionamiento de la entrada y los hooks de entrenamiento, como el checkpointing de gradientes y el agrupamiento de parámetros. Esto se describe en la API de modelos de timm, incluyendo qué familias de modelos admiten qué partes.

Puntos de control y el Hub

• load_checkpoint carga un archivo de checkpoint (PyTorch o .safetensors) en un modelo.
• load_state_dict lee un archivo de punto de control y devuelve su diccionario de estado.
• resume_checkpoint restaura el estado del modelo, el optimizador y el escalador de pérdida de un punto de control de train.py y devuelve la época desde la que reanudar.

load_checkpoint y load_state_dict seleccionan los pesos EMA de un checkpoint de train.py de forma predeterminada (use_ema) y eliminan los prefijos de clave module. (DDP) y _orig_mod. (torch.compile).

• save_for_hf escribe los pesos y un config.json (arquitectura, cfg preentrenado, model_args, label_names) en una carpeta.
• push_to_hf_hub sube los pesos y la configuración a un repositorio del Hub con una tarjeta de modelo generada si el repositorio no tiene ninguna.

A continuación, el modelo se crea con hf-hub: o local-dir:, consulte compartir y cargar modelos.

import timm
from timm.models import save_for_hf

model = timm.create_model('resnet18', num_classes=4)
# ... fine-tune
save_for_hf(model, 'my_resnet18', model_config=dict(label_names=['a', 'b', 'c', 'd']), safe_serialization=True)
model = timm.create_model('local-dir:my_resnet18', pretrained=True)  # 4 classes, weights and labels restored
print(model.pretrained_cfg['label_names'])

Modelos personalizados

Una función de punto de entrada de modelo decorada con register_model se añade al registro, por lo que puede crearse con create_model y encontrarse con list_models. Los puntos de entrada en timm construyen su modelo con build_model_with_cfg, que gestiona la configuración preentrenada, la carga de pesos preentrenados y features_only.

API

Creación de modelos

>>> from timm import create_model

>>> # Create a MobileNetV3-Large model with no pretrained weights.
>>> model = create_model('mobilenetv3_large_100')

>>> # Create a MobileNetV3-Large model with pretrained weights.
>>> model = create_model('mobilenetv3_large_100', pretrained=True)
>>> model.num_classes
1000

>>> # Create a MobileNetV3-Large model with pretrained weights and a new head with 10 classes.
>>> model = create_model('mobilenetv3_large_100', pretrained=True, num_classes=10)
>>> model.num_classes
10

>>> # Create a Dinov2 small model with pretrained weights and save weights in a custom directory.
>>> model = create_model('vit_small_patch14_dinov2.lvd142m', pretrained=True, cache_dir="/data/my-models")
>>> # Data will be stored at */data/my-models/models--timm--vit_small_patch14_dinov2.lvd142m/*

Descubrimiento de modelos

Configuración preentrenada

Entrada y exportación del modelo actual

Puntos de control y Hub

Modelos personalizados