A API de Modelos do timm

Todo modelo do timm é um torch.nn.Module que compartilha uma interface comum sobre o forward(). A interface divide o modelo em um extrator de recursos e uma cabeça de classificador, expõe recursos intermediários e fornece hooks para treinamento (checkpointing de gradiente, agrupamento de parâmetros) e para alterar a resolução de entrada. O código escrito para essa interface funciona entre arquiteturas, que é a maneira como o script de treinamento do timm, os wrappers de extração de recursos e a fábrica de otimizadores tratam todos os modelos da mesma forma.

Esta página descreve a interface de um modelo construído: quais partes todos os modelos fornecem, quais a maioria fornece e quais são específicas de algumas famílias. Como os modelos são criados (argumentos de create_model, tags pré-treinadas, a configuração pré-treinada) é abordado na referência de modelos.

Extrator de recursos e cabeça

A passagem direta de um modelo é forward_head(forward_features(x)).

• forward_features(x) retorna a saída não agregada da última etapa, após a camada de normalização final, se o modelo tiver uma: um mapa de recursos NCHW para modelos convolucionais, tokens NLC (incluindo quaisquer tokens de prefixo de classe/registro) para modelos no estilo ViT.
• forward_head(x, pre_logits=False) aplica o pooling global (ou seleção de tokens), as camadas pré-logits e o classificador. Com pre_logits=True, ele para antes do classificador e retorna os recursos agregados.

Três atributos descrevem os tamanhos envolvidos:

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)

Alguns modelos se desviam do formato de entrada e saída com um único tensor: modelos de múltiplos ramos (CrossViT, EfficientViM) retornam uma lista de forward_features, e os ViTs NaFlex / Gemma4 aceitam entradas extras de coordenadas de patch e validade para resolução variável.

Classificador

• get_classifier() retorna o módulo do classificador.
• reset_classifier(num_classes, global_pool=None) substitui o classificador por um novo, inicializado aleatoriamente.
  • num_classes=0 substitui o classificador por uma identidade, de modo que model(x) retorna os recursos agregados quando a agregação permanece habilitada.
  • global_pool altera o pooling ao mesmo tempo. global_pool='' o desativa; combinado com num_classes=0, model(x) retorna recursos não poolados.
  • global_pool=None mantém o pooling atual.

Criar um modelo com num_classes / global_pool definidos tem o mesmo efeito que redefinir o classificador posteriormente, sem alocar o classificador original. Torres de visão apenas de codificador (os codificadores DeepSeek, Gemma4, MobileNetV5 e Qwen3) não possuem classificador e não implementam esses métodos.

Recursos intermediários

• forward_intermediates(x, indices=None, norm=False, stop_early=False, output_fmt='NCHW', intermediates_only=False) executa o modelo e retorna os recursos finais, além de uma lista de recursos intermediários.
  • indices seleciona os intermediários: None para todos, um int n para os últimos n, ou uma lista de índices (valores negativos contam a partir do final).
  • intermediates_only=True retorna apenas os recursos intermediários.
  • stop_early=True ignora os blocos após o último índice selecionado (alguns modelos, por exemplo, MobileNetV3, exigem intermediates_only=True).
  • norm=True aplica a camada de normalização final a cada intermediário.
• feature_info descreve os recursos intermediários disponíveis, com uma entrada por índice contendo o número de canais (num_chs), o stride em relação à entrada (reduction) e o nome do módulo (module).
• prune_intermediate_layers(indices=1, prune_norm=False, prune_head=True) remove os blocos após o último índice que será usado (e, opcionalmente, a normalização final e o head) e retorna os índices absolutos a serem usados. Use-o quando o modelo for utilizado apenas para extrair intermediários, como um backbone para um head de detecção ou segmentação.
  • prune_head=True remove a cabeça inteira (normalização da cabeça, camadas pré-logits, pooling de atenção e classificador), deixando no máximo o pooling. Passe prune_head=False para manter a cabeça quando a largura do recurso não muda.
  • Quando o pruning altera a largura dos recursos (pruning de um modelo hierárquico para um estágio anterior), o head e a normalização final dimensionados para o último estágio também são removidos. forward_features() então termina no último recurso mantido, e num_features / head_hidden_size são a sua largura.
  • Com os padrões, o modelo podado não possui parâmetros não utilizados (o DDP funciona sem find_unused_parameters), exceto CoaT (blocos paralelos), CrossViT e HRNet quando apenas algumas de suas escalas / ramificações paralelas são utilizadas.

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

Todos os recursos são retornados como mapas NCHW por padrão, incluindo os de modelos no estilo ViT, cujos tokens espaciais são remodelados em uma grade com os tokens de prefixo descartados. As outras opções de output_fmt são específicas da família (NLC para modelos no estilo ViT, NHWC para alguns), assim como argumentos extras, como return_prefix_tokens (família ViT). A criação de um modelo com features_only=True o encapsula como um backbone que retorna os mapas de recursos de model(x); consulte extração de recursos para detalhes sobre ambas as abordagens.

Tamanho da entrada

Os modelos se dividem em três grupos, de acordo com a forma como tratam a resolução de entrada:

• Qualquer tamanho de entrada. A maioria dos modelos convolucionais e híbridos convolucionais sem embeddings de posição de tamanho fixo ou partições de janela funciona em qualquer tamanho de entrada maior que o stride total, sem alterações.
• Tamanho de entrada fixo, redimensionável. Modelos com embeddings de posição, partições de janela de atenção ou tabelas de posição relativa são construídos para um tamanho específico, definido pelo argumento img_size na criação (padrão: o input_size do cfg pré-treinado). Nas famílias que implementam isso (ViT, DeiT, BEiT, EVA, Swin, MaxViT e outras, veja a tabela abaixo), set_input_size(img_size=...) altera esse tamanho posteriormente, reamostrando os embeddings de posição e as tabelas de posição relativa e atualizando as janelas. Vários aceitam argumentos adicionais, por exemplo, patch_size para ViT, window_size / window_ratio para Swin. Chamar a função no tamanho atual não altera o modelo. Os demais devem ser criados com o img_size em que serão utilizados.
• Tamanho de entrada dinâmico. Modelos das famílias ViT e EVA criados com dynamic_img_size=True reamostram seu embedding de posição para o tamanho de entrada a cada passagem de frente, com algum custo em velocidade. Modelos Swin, SwinV2 e SwinV2-CR criados com dynamic_img_size=True (ou o equivalente strict_img_size=False) calculam suas máscaras de atenção para o tamanho de entrada em cada passagem de frente, em vez disso. ViTs NaFlex lidam com resolução variável 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))

Hooks de treinamento

• set_grad_checkpointing(enable=True) habilita o checkpointing de ativações (gradientes), recalculando as ativações do forward de cada bloco no backward pass para economizar memória (--grad-checkpointing em train.py). A maioria dos modelos o suporta, mas algumas arquiteturas mais antigas (VGG, Inception, NASNet, …) geram um erro.
• group_matcher(coarse=False) retorna expressões regulares que atribuem cada parâmetro a um grupo de caule ou bloco. A fábrica de otimizadores usa isso para escalar as taxas de aprendizado por camada para o decaimento da taxa de aprendizado por camada (layer_decay, consulte a referência de otimizadores).
• no_weight_decay(), quando implementado, retorna os nomes dos parâmetros que não devem ter weight decay aplicado (embeddings de posição, tokens de classe, …). A fábrica de otimizador exclui esses parâmetros e todos os parâmetros 1D (biases, escalas de normalização) do weight decay.

Configuração pré-treinada

model.pretrained_cfg (e seu alias mais antigo model.default_cfg) é um dicionário que descreve o checkpoint de origem: tamanho de entrada, interpolação, recorte, normalização, nomes do classificador e do primeiro módulo de convolução, licença e fonte. Ele é usado para adaptar o classificador e a convolução de entrada quando pesos pré-treinados são carregados, e como padrões de pré-processamento. Consulte a referência de modelos para os campos.

O cfg pré-treinado não é alterado quando um modelo é criado ou modificado para diferir de seu checkpoint (in_chans, img_size, set_input_size()). A entrada atual do modelo é rastreada separadamente, timm.models.get_model_input_config(model) retorna seus canais, tamanho e restrição de tamanho fixo, e timm.data.resolve_input_data_config(model) resolve uma configuração de dados para as transformações em relação a ela. timm.models.get_model_args(model) retorna os argumentos para reconstruir o modelo como ele está agora, e save_for_hf os grava na configuração exportada, para que um modelo redimensionado ou adaptado possa ser recarregado. Veja adaptar e exportar modelos.

Suporte da família

A tabela abaixo lista quais partes opcionais da interface cada família de modelos suporta. Ela é gerada por hfdocs/generate_model_docs.py, que constrói e executa um pequeno modelo de cada classe de modelo em cada família. O suporte pode variar entre as arquiteturas de uma família; consulte as páginas das famílias para mais detalhes.

✓ suportado, ✓ (quadrado) funciona apenas para tamanhos quadrados, ✗ não suportado, – não necessário (o modelo já aceita qualquer tamanho de entrada). Cada classe de modelo de uma família é verificada executando cada recurso em um modelo inicializado aleatoriamente da menor arquitetura (não de teste) da família. Classes com os mesmos resultados compartilham uma linha, e a página da família lista a arquitetura usada para cada uma.

Implementando um modelo

Um novo modelo deve implementar a interface para funcionar com o script de treinamento, a extração de recursos e a fábrica de otimizadores, sem casos especiais:

• forward_features, forward_head(x, pre_logits=False) e forward como sua composição, com num_classes, num_features e head_hidden_size definidos.
• get_classifier e reset_classifier(num_classes, global_pool=None).
• forward_intermediates, prune_intermediate_layers e uma lista feature_info.
• set_grad_checkpointing, group_matcher e no_weight_decay para qualquer parâmetro que não deve ser decaído.
• set_input_size se o modelo for construído para um tamanho de entrada fixo.
• Os argumentos padrão do construtor: num_classes, in_chans, global_pool, drop_rate, drop_path_rate, e device / dtype para a criação dos parâmetros.

tests/test_models.py verifica as partes de forward, backward, features, intermediates, classifier reset e pretrained cfg da interface para os modelos registrados (os maiores são excluídos de alguns testes para manter o tempo de CI baixo).