La API de modelos de timm

Cada modelo de timm es un torch.nn.Module que comparte una interfaz común sobre forward(). La interfaz divide el modelo en un extractor de características y una cabeza clasificadora, expone las características intermedias y proporciona ganchos para el entrenamiento (punto de control de gradientes, agrupación de parámetros) y para cambiar la resolución de entrada. El código escrito para ella funciona entre arquitecturas, que es la forma en que el script de entrenamiento de timm, los envoltorios de extracción de características y la fábrica de optimizadores manejan cada modelo de la misma manera.

Esta página describe la interfaz de un modelo construido: qué partes proporciona cada modelo, cuáles proporcionan la mayoría y cuáles son específicas de algunas familias. Cómo se crean los modelos (argumentos de create_model, etiquetas preentrenadas, la configuración preentrenada) se cubre en la referencia de modelos.

Extractor de características y cabeza

La pasada hacia adelante de un modelo es forward_head(forward_features(x)).

• forward_features(x) devuelve la salida no agrupada de la última etapa, después de la capa de normalización final si el modelo la tiene: un mapa de características NCHW para modelos convolucionales, NLC tokens (incluidos los tokens de prefijo de clase / registro) para modelos de estilo ViT.
• forward_head(x, pre_logits=False) aplica el agrupamiento global (o la selección de tokens), las capas previas a los logits y el clasificador. Con pre_logits=True se detiene antes del clasificador y devuelve las características agrupadas.

Tres atributos describen los tamaños involucrados:

import torch
import timm

model = timm.create_model('convnext_atto').eval()
x = torch.randn(2, 3, 224, 224)

features = model.forward_features(x)  # (2, 320, 7, 7), num_features = 320
pooled = model.forward_head(features, pre_logits=True)  # (2, 320), head_hidden_size = 320
logits = model.forward_head(features)  # (2, 1000), same as model(x)

Algunos modelos se desvían de la forma de entrada y salida de un solo tensor: los modelos de múltiples ramas (CrossViT, EfficientViM) devuelven una lista desde forward_features, y los ViTs de NaFlex / Gemma4 aceptan entradas adicionales de coordenadas de parche y validez para una resolución variable.

Clasificador

• get_classifier() devuelve el módulo clasificador.
• reset_classifier(num_classes, global_pool=None) reemplaza el clasificador por uno nuevo, inicializado aleatoriamente.
  • num_classes=0 reemplaza el clasificador con una identidad, por lo que model(x) devuelve las características agrupadas cuando el agrupamiento sigue habilitado.
  • global_pool modifica el pooling al mismo tiempo. global_pool='' lo desactiva; combinado con num_classes=0, model(x) devuelve características sin pooling.
  • global_pool=None mantiene el agrupamiento actual.

Crear un modelo con num_classes / global_pool establecidos tiene el mismo efecto que restablecer el clasificador después, sin asignar el clasificador original. Las torres de visión solo de codificador (los codificadores de DeepSeek, Gemma4, MobileNetV5 y Qwen3) no tienen clasificador y no implementan estos métodos.

Características intermedias

• forward_intermediates(x, indices=None, norm=False, stop_early=False, output_fmt='NCHW', intermediates_only=False) ejecuta el modelo y devuelve las características finales junto con una lista de características intermedias.
  • indices selecciona los intermediarios: None para todos, un int n para los últimos n, o una lista de índices (los valores negativos se cuentan desde el final).
  • intermediates_only=True devuelve solo las características intermedias.
  • stop_early=True omite los bloques posteriores al último índice seleccionado (algunos modelos, p. ej. MobileNetV3, requieren intermediates_only=True).
  • norm=True aplica la capa de normalización final a cada intermedio.
• feature_info describe las características intermedias disponibles, una entrada por índice con el número de canales (num_chs), el stride relativo a la entrada (reduction) y el nombre del módulo (module).
• prune_intermediate_layers(indices=1, prune_norm=False, prune_head=True) elimina los bloques posteriores al último índice que se utilizará (y opcionalmente la normalización final y la cabeza) y devuelve los índices absolutos a tomar. Úsalo cuando el modelo solo se empleará para extraer intermedios, como columna vertebral para una cabeza de detección o segmentación.
  • prune_head=True elimina la cabeza completa (normalización de la cabeza, capas previas a los logits, pooling de atención y clasificador), dejando como máximo el pooling. Pasa prune_head=False para conservar la cabeza cuando el ancho de las características no cambia.
  • Cuando el recorte modifica el ancho de las características (recortar un modelo jerárquico a una etapa anterior), también se eliminan la cabeza y la normalización final dimensionadas para la última etapa. forward_features() entonces termina en la última característica conservada, y num_features / head_hidden_size son su ancho.
  • Con los valores predeterminados, el modelo podado no tiene parámetros sin usar (DDP funciona sin find_unused_parameters), excepto CoaT (bloques paralelos), CrossViT y HRNet cuando solo se toman algunas de sus escalas / ramas paralelas.

model = timm.create_model('resnet18').eval()
print([(f['num_chs'], f['reduction']) for f in model.feature_info])
# [(64, 2), (64, 4), (128, 8), (256, 16), (512, 32)]

final, intermediates = model.forward_intermediates(x)  # all 5 feature levels
last_two = model.forward_intermediates(x, indices=2, intermediates_only=True)  # (2, 256, 14, 14), (2, 512, 7, 7)

take_indices = model.prune_intermediate_layers(indices=2, prune_head=True)  # [3, 4], no classifier left

Todas las características se devuelven por defecto como mapas NCHW, incluidas las de los modelos de estilo ViT, cuyos tokens espaciales se reorganizan en una cuadrícula descartando los tokens de prefijo. Las otras opciones de output_fmt son específicas de cada familia (NLC para los modelos de estilo ViT, NHWC para algunos), al igual que los argumentos adicionales como return_prefix_tokens (familia ViT). Crear un modelo con features_only=True lo envuelve como un backbone que devuelve los mapas de características de model(x); consulte extracción de características para ver ambos enfoques en detalle.

Tamaño de entrada

Los modelos se dividen en tres grupos según cómo manejan la resolución de entrada:

• Cualquier tamaño de entrada. La mayoría de los modelos convolucionales, y los híbridos convolucionales sin embeddings de posición de tamaño fijo o particiones de ventana, funcionan con cualquier tamaño de entrada mayor que su paso total, sin cambios.
• Tamaño de entrada fijo, redimensionable. Los modelos con embeddings de posición, particiones de ventana de atención o tablas de posición relativa se construyen para un tamaño determinado, establecido por el argumento img_size en la creación (predeterminado: el input_size de la configuración preentrenada). En las familias que lo implementan (ViT, DeiT, BEiT, EVA, Swin, MaxViT y otras, ver la tabla inferior), set_input_size(img_size=...) lo modifica posteriormente, muestreando de nuevo los embeddings de posición y las tablas de posición relativa y actualizando las ventanas. Varios aceptan argumentos adicionales, p. ej. patch_size para ViT, window_size / window_ratio para Swin. Llamarlo con el tamaño actual deja el modelo sin cambios. Los demás deben crearse con el img_size con el que se van a utilizar.
• Tamaño de entrada dinámico. Los modelos de las familias ViT y EVA creados con dynamic_img_size=True re-muestrean su embedding de posición al tamaño de entrada en cada pasada de adelante, con cierto coste en velocidad. Los modelos Swin, SwinV2 y SwinV2-CR creados con dynamic_img_size=True (o el equivalente strict_img_size=False) calculan sus máscaras de atención para el tamaño de entrada en cada pasada de adelante. Los ViT NaFlex manejan la resolución variable de forma nativa.

model = timm.create_model('vit_tiny_patch16_224')
model.set_input_size(img_size=(384, 384))
model(torch.randn(1, 3, 384, 384))

Ganchos de entrenamiento

• set_grad_checkpointing(enable=True) habilita el checkpointing de activaciones (gradientes), que recalcula las activaciones hacia adelante de cada bloque en la pasada hacia atrás para ahorrar memoria (--grad-checkpointing en train.py). La mayoría de los modelos lo admiten, pero algunas arquitecturas más antiguas (VGG, Inception, NASNet, …) generan un error.
• group_matcher(coarse=False) devuelve expresiones regulares que asignan cada parámetro a un grupo de tallo o bloque. La fábrica de optimizadores lo utiliza para escalar las tasas de aprendizaje por capa para el decaimiento de la tasa de aprendizaje por capa (layer_decay, consulte la referencia de optimizadores).
• no_weight_decay(), cuando está implementado, devuelve los nombres de los parámetros que no deben tener decaimiento de peso (embeddings de posición, tokens de clase, …). La fábrica del optimizador excluye estos y todos los parámetros de 1D (sesgos, escalas de normalización) del decaimiento de peso.

Configuración preentrenada

model.pretrained_cfg (y su alias anterior model.default_cfg) es un diccionario que describe el punto de control de origen: tamaño de entrada, interpolación, recorte, normalización, nombres del clasificador y del primer módulo de convolución, licencia y fuente. Se utiliza para adaptar el clasificador y la convolución de entrada cuando se cargan pesos preentrenados, y como valores predeterminados de preprocesamiento. Consulte la referencia de modelos para ver los campos.

La configuración preentrenada no se modifica cuando un modelo se crea o se modifica para que difiera de su punto de control (in_chans, img_size, set_input_size()). La entrada actual del modelo se rastrea por separado, timm.models.get_model_input_config(model) devuelve sus canales, tamaño y restricción de tamaño fijo, y timm.data.resolve_input_data_config(model) resuelve una configuración de datos para las transformaciones en relación con ella. timm.models.get_model_args(model) devuelve los argumentos para reconstruir el modelo tal como está ahora, y save_for_hf los escribe en la configuración exportada, de modo que un modelo redimensionado o adaptado pueda recargarse. Consulte adaptar y exportar modelos.

Apoyo por familia

La tabla siguiente enumera qué partes opcionales de la interfaz soporta cada familia de modelos. Se genera mediante hfdocs/generate_model_docs.py, que construye y ejecuta un modelo pequeño de cada clase de modelo en cada familia. El soporte puede variar entre las arquitecturas de una familia; consulte las páginas de las familias para más detalles.

✓ compatible, ✓ (cuadrado) funciona solo para tamaños cuadrados, ✗ no compatible, – no necesario (el modelo ya acepta cualquier tamaño de entrada). Cada clase de modelo de una familia se comprueba ejecutando cada característica en un modelo inicializado aleatoriamente de su arquitectura más pequeña (no de prueba) de la familia. Las clases con los mismos resultados comparten una fila; la página de la familia enumera la arquitectura utilizada para cada una.

Implementar un modelo

Un nuevo modelo debe implementar la interfaz para que funcione con el script de entrenamiento, la extracción de características y la fábrica de optimizadores sin casos especiales:

• forward_features, forward_head(x, pre_logits=False) y forward como su composición, con num_classes, num_features y head_hidden_size establecidos.
• get_classifier y reset_classifier(num_classes, global_pool=None).
• forward_intermediates, prune_intermediate_layers y una lista feature_info.
• set_grad_checkpointing, group_matcher y no_weight_decay para cualquier parámetro que no deba decaer.
• set_input_size si el modelo está diseñado para un tamaño de entrada fijo.
• Los argumentos estándar del constructor: num_classes, in_chans, global_pool, drop_rate, drop_path_rate y device / dtype para la creación de parámetros.

tests/test_models.py comprueba las partes de la interfaz de forward, backward, features, intermediates, classifier reset y pretrained cfg de los modelos registrados (los más grandes se excluyen de algunas pruebas para reducir el tiempo de CI).