特征提取

timm 中的所有模型都具有一致的机制，用于从模型中获取除分类任务之外的其他任务所需的各类特征。

倒数第二层特征（分类器前特征）

可以通过多种方式获取倒数第二层模型的特征，而无需对模型进行手术式修改（尽管你也可以随意进行手术）。首先，需要确定想要获取的是池化后的特征还是未池化的特征。

未池化

获取未池化特征有三种方式。最终的未池化特征有时被称为最后一个隐藏状态。在 timm 中，这包括直至最终归一化层（例如 ViT 风格的模型）的内容，但不包括池化 / 类别标记选择以及池化后的最终层。

在不修改网络的情况下，可以调用任意模型的 model.forward_features(input) 来替代通常的 model(input)。这将绕过网络的头部分类器和全局池化。

如果希望显式修改网络以返回未池化的特征，可以创建不带分类器和池化层的模型，或者在创建后再将其移除。这两种方式都会从网络中移除与分类器相关的参数。

forward_features()

>>> import torch
>>> import timm
>>> m = timm.create_model('xception41', pretrained=True)
>>> o = m(torch.randn(2, 3, 299, 299))
>>> print(f'Original shape: {o.shape}')
>>> o = m.forward_features(torch.randn(2, 3, 299, 299))
>>> print(f'Unpooled shape: {o.shape}')

输出：

Original shape: torch.Size([2, 1000])
Unpooled shape: torch.Size([2, 2048, 10, 10])

创建不带分类器和池化的

>>> import torch
>>> import timm
>>> m = timm.create_model('resnet50', pretrained=True, num_classes=0, global_pool='')
>>> o = m(torch.randn(2, 3, 224, 224))
>>> print(f'Unpooled shape: {o.shape}')

输出：

Unpooled shape: torch.Size([2, 2048, 7, 7])

稍后将其移除

>>> import torch
>>> import timm
>>> m = timm.create_model('densenet121', pretrained=True)
>>> o = m(torch.randn(2, 3, 224, 224))
>>> print(f'Original shape: {o.shape}')
>>> m.reset_classifier(0, '')
>>> o = m(torch.randn(2, 3, 224, 224))
>>> print(f'Unpooled shape: {o.shape}')

输出：

Original shape: torch.Size([2, 1000])
Unpooled shape: torch.Size([2, 1024, 7, 7])

将未池化的输出链接到分类器

可以使用 forward_head() 函数将最后一个隐藏状态反馈到模型的头部。

>>> model = timm.create_model('vit_medium_patch16_reg1_gap_256', pretrained=True)
>>> output = model.forward_features(torch.randn(2,3,256,256))
>>> print('Unpooled output shape:', output.shape)
>>> classified = model.forward_head(output)
>>> print('Classification output shape:', classified.shape)

输出：

Unpooled output shape: torch.Size([2, 257, 512])
Classification output shape: torch.Size([2, 1000])

池化

若要修改网络以返回池化特征，可以使用 forward_features() 并自行对结果进行池化/展平，或者像上面那样修改网络但保留池化操作。

创建无分类器

>>> import torch
>>> import timm
>>> m = timm.create_model('resnet50', pretrained=True, num_classes=0)
>>> o = m(torch.randn(2, 3, 224, 224))
>>> print(f'Pooled shape: {o.shape}')

输出：

Pooled shape: torch.Size([2, 2048])

稍后将其移除

>>> import torch
>>> import timm
>>> m = timm.create_model('ese_vovnet19b_dw', pretrained=True)
>>> o = m(torch.randn(2, 3, 224, 224))
>>> print(f'Original shape: {o.shape}')
>>> m.reset_classifier(0)
>>> o = m(torch.randn(2, 3, 224, 224))
>>> print(f'Pooled shape: {o.shape}')

输出：

Original shape: torch.Size([2, 1000])
Pooled shape: torch.Size([2, 1024])

多尺度特征图（特征金字塔）

目标检测、分割、关键点以及多种密集像素任务都需要访问骨干网络在多个尺度上的特征图。这通常是通过修改原始分类网络来实现的。由于每个网络的结构差异较大，在任意给定的目标检测或分割库中，通常只支持少数几种骨干网络。

timm 提供了一致的接口，用于将其中包含的任意模型创建为特征骨干网络，以输出选定层级的特征图。

可以通过在任意 create_model 调用中添加参数 features_only=True 来创建特征主干网络。默认情况下，大多数具有特征层级的模型会输出最多 5 个特征，直至下采样 32 倍。不过具体行为因模型而异：一些模型的层级较少，而另一些模型（如 ViT）则拥有更多非层级化的特征图，默认输出最后 3 个。可以向 create_model 传递 out_indices 参数，以指定需要输出的特征。

创建特征图提取模型

>>> import torch
>>> import timm
>>> m = timm.create_model('resnest26d', features_only=True, pretrained=True)
>>> o = m(torch.randn(2, 3, 224, 224))
>>> for x in o:
...     print(x.shape)

输出：

torch.Size([2, 64, 112, 112])
torch.Size([2, 256, 56, 56])
torch.Size([2, 512, 28, 28])
torch.Size([2, 1024, 14, 14])
torch.Size([2, 2048, 7, 7])

查询特征信息

在创建特征主干网络后，可以对其进行查询，以向下游头提供通道或分辨率缩减信息，而无需静态配置或硬编码常量。.feature_info 属性是一个类，用于封装有关特征提取点的信息。

>>> import torch
>>> import timm
>>> m = timm.create_model('regnety_032', features_only=True, pretrained=True)
>>> print(f'Feature channels: {m.feature_info.channels()}')
>>> o = m(torch.randn(2, 3, 224, 224))
>>> for x in o:
...     print(x.shape)

输出：

Feature channels: [32, 72, 216, 576, 1512]
torch.Size([2, 32, 112, 112])
torch.Size([2, 72, 56, 56])
torch.Size([2, 216, 28, 28])
torch.Size([2, 576, 14, 14])
torch.Size([2, 1512, 7, 7])

选择特定的特征级别或限制步幅

还有两个额外的创建参数会影响输出特征。

out_indices 选择要输出的索引
output_stride 用于限制网络的特征输出步幅（顺便一提，在分类模式下同样有效）

输出索引选择

out_indices 参数为所有模型所支持，但并非所有模型都具有相同的索引到特征步幅的映射关系。请查阅代码或检查 feature_info 以进行比较。输出索引通常对应于第 C(i+1)th 个特征层级（即 2^(i+1) 倍的下采样）。对于大多数卷积神经网络模型，索引 0 对应步幅为 2 的特征，索引 4 对应步幅为 32 的特征。对于许多 ViT 或 ViT-Conv 混合模型，可能存在大量甚至全部具有相同形状的特征图，或者是层级与非层级特征图的组合。建议查看 feature_info 属性，以了解特征的数量、对应的通道数以及下采样级别。

out_indices 支持负索引，这使得获取最后一个、倒数第二个等特征图变得非常方便。out_indices=(-2,) 将返回任意模型的倒数第二个特征图。

输出步长（特征图膨胀）

output_stride 通过将层转换为使用空洞卷积来实现。这样做并不总是很直接，一些网络仅支持 output_stride=32。

>>> import torch
>>> import timm
>>> m = timm.create_model('ecaresnet101d', features_only=True, output_stride=8, out_indices=(2, 4), pretrained=True)
>>> print(f'Feature channels: {m.feature_info.channels()}')
>>> print(f'Feature reduction: {m.feature_info.reduction()}')
>>> o = m(torch.randn(2, 3, 320, 320))
>>> for x in o:
...     print(x.shape)

输出：

Feature channels: [512, 2048]
Feature reduction: [8, 8]
torch.Size([2, 512, 40, 40])
torch.Size([2, 2048, 40, 40])

灵活提取中间特征图

除了使用模型工厂中的 features_only 之外，所有模型都支持 forward_intermediates() 方法，该方法提供了一种灵活的机制，用于提取中间特征图和最后一个隐藏状态（可以链接到头部）。此外，此方法还支持一些特定于模型的功能，例如为某些模型返回类别或蒸馏前缀标记。

与 forward_intermediates 函数配套的是 prune_intermediate_layers 函数，它允许从模型中剪除不需要的层，包括头部、最终归一化层和/或末尾的块/阶段。

indices 参数同时用于 forward_intermediates() 和 prune_intermediate_layers()，以选择要返回的特征或要移除的层。与 features_only API 中的 out_indices 类似，indices 因模型而异，用于选择返回哪些中间结果。

在 ViT 等非层级块状模型中，索引对应于各个块；在具有层级阶段的模型中，它们通常对应于 stem 以及每个层级阶段的输出。正向索引（从起始位置开始）和负向索引（相对于末尾）均可使用，并且使用 None 可返回所有中间结果。

prune_intermediate_layers() 调用会返回一个 indices 变量，因为在裁剪模型时，负索引必须转换为绝对（正）索引。

model = timm.create_model('vit_medium_patch16_reg1_gap_256', pretrained=True)
output, intermediates = model.forward_intermediates(torch.randn(2,3,256,256))
for i, o in enumerate(intermediates):
    print(f'Feat index: {i}, shape: {o.shape}')

Feat index: 0, shape: torch.Size([2, 512, 16, 16])
Feat index: 1, shape: torch.Size([2, 512, 16, 16])
Feat index: 2, shape: torch.Size([2, 512, 16, 16])
Feat index: 3, shape: torch.Size([2, 512, 16, 16])
Feat index: 4, shape: torch.Size([2, 512, 16, 16])
Feat index: 5, shape: torch.Size([2, 512, 16, 16])
Feat index: 6, shape: torch.Size([2, 512, 16, 16])
Feat index: 7, shape: torch.Size([2, 512, 16, 16])
Feat index: 8, shape: torch.Size([2, 512, 16, 16])
Feat index: 9, shape: torch.Size([2, 512, 16, 16])
Feat index: 10, shape: torch.Size([2, 512, 16, 16])
Feat index: 11, shape: torch.Size([2, 512, 16, 16])

model = timm.create_model('vit_medium_patch16_reg1_gap_256', pretrained=True)
print('Original params:', sum([p.numel() for p in model.parameters()]))

indices = model.prune_intermediate_layers(indices=(-2,), prune_head=True, prune_norm=True)  # prune head, norm, last block
print('Pruned params:', sum([p.numel() for p in model.parameters()]))

intermediates = model.forward_intermediates(torch.randn(2,3,256,256), indices=indices, intermediates_only=True)  # return penultimate intermediate
for o in intermediates:    
    print(f'Feat shape: {o.shape}')

Original params: 38880232
Pruned params: 35212800
Feat shape: torch.Size([2, 512, 16, 16])

Update on GitHub