Output layer names for feature extractor

[code-demoe]

I want to get the names of layers/blocks giving the output in feature extractors.
E.g.: model=timm.create_model('resnet50',features_only=True,out_indices=(0,1,2,3))
out=model(torch.rand((1,3,32,32))
for k in out:
print(k.shape)

torch.Size([1, 64, 16, 16])
torch.Size([1, 256, 8, 8])
torch.Size([1, 512, 4, 4])
torch.Size([1, 1024, 2, 2])

Now, I want to get the layer/block/stage names giving these outputs. How can I do that?

[rwightman]

@code-demoe it's all in the feature_info object of the model after it's been created... any of the getters used below will by default return the values for the specified out_indices, but you can request the info for specific indices, including those that weren't selected if you want.

https://github.com/rwightman/pytorch-image-models/blob/f1f332fee50aa29c435b56cef16e137b6a6ca246/timm/models/features.py#L20-L82

model = timm.create_model('resnet50', features_only=True)

model.feature_info.channels()
[64, 256, 512, 1024, 2048]

model.feature_info.module_name()
['act1', 'layer1', 'layer2', 'layer3', 'layer4']

model.feature_info.get_dicts()
[{'num_chs': 64, 'reduction': 2, 'module': 'act1'}, {'num_chs': 256, 'reduction': 4, 'module': 'layer1'}, {'num_chs': 512, 'reduction': 8, 'module': 'layer2'}, {'num_chs': 1024, 'reduction': 16, 'module': 'layer3'}, {'num_chs': 2048, 'reduction': 32, 'module': 'layer4'}]

>>> m = timm.create_model('nf_resnet50', features_only=True, out_indices=(2,3))

>>> m.feature_info.get_dicts(keys=('reduction', 'num_chs'))
[{'reduction': 8, 'num_chs': 512}, {'reduction': 16, 'num_chs': 1024}]

>>> m.feature_info.module_name()
['stages.1', 'stages.2']

>>> m.feature_info.module_name(idx=(0, -1))
['stem.conv', 'stages.3']

[code-demoe]

Thank you so much, this saved me a lot of effort in reading code.

[rwightman]

This interface was designed so the models that consume these as backbones can be completely dynamic without requiring hard coded values, they can just query the backbone.

[rwightman]

One thing to be aware of though, the feature extraction model may be modified from the original, the names reference the original, I have thought of also including the modified names (many models have some of the sequential layers flattened to make the extraction more efficient, without requiring the use of the hook based mode).

[code-demoe]

Sounds good. Thanks!

[code-demoe]

Can you add a function in feature_info to return index of the feature extractor layers in full model, in some models the string literal returned by model.feature_info.module_name() doesn't match with the layer name in the model. There's a mismatch of '_'.
e.g.

model.feature_info.module_name()
stages.0

but layer name inside model is stages_0
Also, I wasn't able to find a decent way of accessing a model layer using layer name. I tried using model.named_modules(), model._modules, model.named_children(), but PyTorch doesn't seem to be good with accessing layers using string literals. If I am missing something simple please guide me towards the correct approach.