timm 模型 API

每个 timm 模型都是一个 torch.nn.Module，在 forward() 之上共享一个通用接口。该接口将模型划分为特征提取器和分类器头部，暴露中间特征，并为训练（梯度检查点、参数分组）和更改输入分辨率提供钩子。基于该接口编写的代码可以跨架构工作，这正是 timm 的训练脚本、特征提取封装器和优化器工厂能够以相同方式处理所有模型的原因。

本页介绍已构建模型的接口：每个模型都提供哪些部分、大多数模型提供哪些部分，以及哪些部分是特定于某些模型家族的。模型的创建方式（create_model 参数、预训练标签、预训练配置）在模型参考中介绍。

特征提取器和头部

模型的前向传播是 forward_head(forward_features(x))。

• forward_features(x) 返回最后一个阶段的未池化输出，如果模型存在最终的 norm 层，则返回经过该层处理后的结果：对于卷积模型，返回 NCHW 特征图；对于 ViT 风格的模型，返回 NLC 个 token（包括任何 class / register 前缀 token）。
• forward_head(x, pre_logits=False) 应用全局池化（或 token 选择）、任何 pre-logits 层以及分类器。当 pre_logits=True 时，它会在分类器之前停止并返回池化后的特征。

三个属性描述了所涉及的尺寸：

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)

少数模型偏离了单张量输入、单张量输出的形式：多分支模型（CrossViT、EfficientViM）的 forward_features 返回一个列表，而 NaFlex / Gemma4 ViT 则接受额外的 patch 坐标和有效性输入，以支持可变分辨率。

分类器

• get_classifier() 返回分类器模块。
• reset_classifier(num_classes, global_pool=None) 用一个新的、随机初始化的分类器替换现有分类器。
  • num_classes=0 会用恒等映射替换分类器，因此在池化保持启用时，model(x) 会返回池化后的特征。
  • global_pool 会同时更改池化方式。global_pool='' 将其禁用；与 num_classes=0 结合使用时，model(x) 将返回未池化的特征。
  • global_pool=None 保持当前的池化方式。

将 num_classes / global_pool 设置为特定值来创建模型，其效果等同于在创建后重置分类器，且无需分配原始分类器。仅包含编码器的视觉塔（DeepSeek、Gemma4、MobileNetV5 和 Qwen3 编码器）没有分类器，也不实现这些方法。

中间特征

• forward_intermediates(x, indices=None, norm=False, stop_early=False, output_fmt='NCHW', intermediates_only=False) 运行模型并返回最终特征以及中间特征列表。
  • indices 用于选择中间结果：None 表示全部，int 类型的 n 表示最后 n 个，或一个索引列表（负值从末尾开始计数）。
  • intermediates_only=True 仅返回中间特征。
  • stop_early=True 会跳过最后一个选定索引之后的块（少数模型，例如 MobileNetV3，需要 intermediates_only=True）。
  • norm=True 会对每个中间层应用最终的 norm 层。
• feature_info 描述了可用的中间特征，每个索引对应一个条目，包含通道数（num_chs）、相对于输入的步幅（reduction）以及模块名称（module）。
• prune_intermediate_layers(indices=1, prune_norm=False, prune_head=True) 移除最后一个将被使用的索引之后的所有块（可选地移除最终的 norm 和 head），并返回需要提取的绝对索引。当模型仅用于提取中间特征，或作为检测或分割 head 的 backbone 时，请使用此方法。
  • prune_head=True 会移除整个头部（包括 head norm、pre-logits 层、attention pooling 和分类器），最多只保留 pooling。如果特征宽度没有变化，请传入 prune_head=False 以保留头部。
  • 当剪枝改变了特征宽度（即将分层模型剪枝到更早的阶段）时，为最后一个阶段配置的 head 和最终 norm 也会被移除。随后，forward_features() 将在最后一个保留的特征处结束，而 num_features / head_hidden_size 即为该特征的宽度。
  • 使用默认设置时，剪枝后的模型没有未使用的参数（DDP 无需使用 find_unused_parameters），但 CoaT（并行块）、CrossViT 和 HRNet 除外，当仅采用其部分并行尺度/分支时。

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

默认情况下，所有特征均以 NCHW 格式的映射返回，包括 ViT 风格模型的输出，其空间 token 会被重塑为网格，并丢弃前缀 token。其他 output_fmt 选项因模型家族而异（例如 ViT 风格模型使用 NLC，某些模型使用 NHWC），额外参数也是如此，例如 return_prefix_tokens（ViT 家族）。使用 features_only=True 创建模型会将其封装为骨干网络，返回 model(x) 的特征映射，有关这两种方法的详细信息，请参阅 特征提取。

输入大小

模型根据其对输入分辨率的处理方式分为三类：

• 任意输入尺寸。大多数卷积模型，以及没有固定尺寸位置嵌入或窗口划分的卷积混合模型，只要输入尺寸大于其总步幅，即可在不作任何更改的情况下处理任意输入尺寸。
• 固定输入尺寸，可调整。 带有位置嵌入、注意力窗口分区或相对位置表的模型是为特定尺寸构建的，该尺寸在创建时由 img_size 参数设定（默认值：预训练配置中的 input_size）。在实现了此功能的模型系列中（ViT、DeiT、BEiT、EVA、Swin、MaxViT 等，详见下表），set_input_size(img_size=...) 可在之后更改该尺寸，它会重新采样位置嵌入和相对位置表，并更新窗口。部分模型接受额外参数，例如 ViT 的 patch_size，Swin 的 window_size / window_ratio。以当前尺寸调用此方法不会改变模型。其他模型必须在创建时指定其将使用的 img_size。
• 动态输入尺寸。使用 dynamic_img_size=True 创建的 ViT 和 EVA 系列模型会在每次前向传播时，将位置嵌入重采样到输入尺寸，这会带来一定的速度开销。使用 dynamic_img_size=True（或等效的 strict_img_size=False）创建的 Swin、SwinV2 和 SwinV2-CR 模型则会在每次前向传播时，针对输入尺寸计算其注意力掩码。NaFlex ViT 原生支持可变分辨率。

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

训练钩子

• set_grad_checkpointing(enable=True) 启用激活（梯度）检查点，在反向传播中重新计算每个块的前向激活以节省内存（train.py 中的 --grad-checkpointing）。大多数模型都支持此功能，但一些较旧的架构（VGG、Inception、NASNet 等）会引发错误。
• group_matcher(coarse=False) 返回正则表达式，用于将每个参数分配到一个主干或块组。 优化器工厂使用它来按层缩放学习率，以实现逐层学习率衰减 （layer_decay，参见 优化器参考）。
• no_weight_decay()（在已实现的情况下）返回不应应用权重衰减的参数名称（位置嵌入、类别标记等）。优化器工厂会将这些参数以及所有 1D 参数（偏置、归一化缩放）排除在权重衰减之外。

预训练配置

model.pretrained_cfg（及其旧别名 model.default_cfg）是一个描述源检查点的字典，包含输入尺寸、插值、裁剪、归一化、分类器和首个卷积模块的名称、许可证及来源。它在加载预训练权重时用于适配分类器和输入卷积，并作为预处理的默认设置。有关字段说明，请参阅 模型参考。

在创建模型或将其修改为与检查点不同的配置（in_chans、img_size、set_input_size()）时，预训练的 cfg 不会改变。模型的当前输入是单独跟踪的，timm.models.get_model_input_config(model) 返回其通道数、尺寸和固定尺寸约束，而 timm.data.resolve_input_data_config(model) 则针对该输入解析用于数据转换的配置。timm.models.get_model_args(model) 返回用于按当前状态重建模型的参数，save_for_hf 会将这些参数写入导出的配置中，从而可以重新加载经过调整大小或适配的模型。请参阅 适配和导出模型。

家庭支持

下表列出了每个模型系列支持的接口可选部分。该表由 hfdocs/generate_model_docs.py 生成，它会构建并运行每个系列中每个模型类的小型模型。支持情况可能因系列内的不同架构而异，详情请参阅系列页面。

✓ 支持，✓ (square) 仅支持方形尺寸，✗ 不支持，– 无需支持（该模型本身已接受任意输入尺寸）。对于每个模型家族中的每个模型类，都会通过在该家族中最小的（非测试）架构上随机初始化模型并运行所有功能来进行检查。具有相同结果的类共享一行，家族页面会列出每个类所使用的架构。

实现模型

新模型应实现该接口，以便与训练脚本、特征提取和优化器工厂配合使用，而无需特殊处理：

• forward_features、forward_head(x, pre_logits=False) 以及它们的组合 forward，其中设置了 num_classes、num_features 和 head_hidden_size。
• get_classifier 和 reset_classifier(num_classes, global_pool=None)。
• forward_intermediates、prune_intermediate_layers 以及一个 feature_info 列表。
• set_grad_checkpointing、group_matcher 以及 no_weight_decay，用于任何不应进行衰减的参数。
• 如果模型是为固定输入大小构建的，则使用set_input_size。
• 标准构造函数参数：num_classes、in_chans、global_pool、drop_rate、drop_path_rate，以及用于创建参数的 device / dtype。

tests/test_models.py 会检查已注册模型接口中的前向、反向、特征、中间层、分类器重置以及预训练配置部分（为了缩短 CI 时间，最大的模型会被排除在某些测试之外）。