[Enhance] Enhance feature extraction function.

[mzr1996]

Motivation

Feature extraction is an important function of the backbone. We already have the extract_feat method in ImageClassifier, but it can only extract the feature map after neck.

Modification

Here we split the whole network into six parts and enable our API to get all of their outputs.

Use cases

Middle layer output, with model.extract_feat

① backbone output

import torch
from mmcv import Config
from mmcls.models import build_classifier

cfg = Config.fromfile('configs/resnet/resnet18_8xb32_in1k.py').model
cfg.backbone.out_indices = (0, 1, 2, 3)  # Output multi-scale feature maps
model = build_classifier(cfg)

outs = model.extract_feat(torch.rand(1, 3, 224, 224), stage='backbone')
for out in outs:
    print(out.shape)
# torch.Size([1, 64, 56, 56])
# torch.Size([1, 128, 28, 28])
# torch.Size([1, 256, 14, 14])
# torch.Size([1, 512, 7, 7])

② neck output

The neck is usually GlobalAveragePooling, and some networks don't have a neck.

import torch
from mmcv import Config
from mmcls.models import build_classifier

cfg = Config.fromfile('configs/resnet/resnet18_8xb32_in1k.py').model
cfg.backbone.out_indices = (0, 1, 2, 3)  # Output multi-scale feature maps
model = build_classifier(cfg)

outs = model.extract_feat(torch.rand(1, 3, 224, 224), stage='neck')
for out in outs:
    print(out.shape)
# torch.Size([1, 64])
# torch.Size([1, 128])
# torch.Size([1, 256])
# torch.Size([1, 512])

③ Pre-logits output (without the final linear classifier head)

Some heads have not only a single linear layer classifier but also some other processing, like VisionTransformerClsHead and StackedLinearClsHead. Now we can extract features before the final linear classifier.

import torch
from mmcv import Config
from mmcls.models import build_classifier

cfg = Config.fromfile('configs/vision_transformer/vit-base-p16_pt-64xb64_in1k-224.py').model
model = build_classifier(cfg)

out = model.extract_feat(torch.rand(1, 3, 224, 224), stage='pre_logits')
print(out.shape)  # The hidden dims in head is 3072
# torch.Size([1, 3072])

Final layer output, with model.simple_test

④ Head linear output (without softmax)

import torch
from mmcv import Config
from mmcls.models import build_classifier

cfg = Config.fromfile('configs/vision_transformer/vit-base-p16_pt-64xb64_in1k-224.py').model
model = build_classifier(cfg)

out = model.simple_test(torch.rand(1, 3, 224, 224), softmax=False, post_process=False)
print(out.shape)
# torch.Size([1, 1000])

⑤ Softmax output (without post-processing)

import torch
from mmcv import Config
from mmcls.models import build_classifier

cfg = Config.fromfile('configs/vision_transformer/vit-base-p16_pt-64xb64_in1k-224.py').model
model = build_classifier(cfg)

out = model.simple_test(torch.rand(1, 3, 224, 224), post_process=False)
print(out.shape)
# torch.Size([1, 1000])

In multi-label tasks, the softmax is changed to sigmoid.

⑥ Post-processing output

In post-processing, we will convert the tensor output into a list.

import torch
from mmcv import Config
from mmcls.models import build_classifier

cfg = Config.fromfile('configs/vision_transformer/vit-base-p16_pt-64xb64_in1k-224.py').model
model = build_classifier(cfg)

out = model.simple_test(torch.rand(1, 3, 224, 224))
print(f'{type(out)} ({len(out)}, {len(out[0])})')
# <class 'list'> (1, 1000)

The post-processing doesn't depend on softmax, you can also post-process the logits output.

out = model.simple_test(torch.rand(1, 3, 224, 224), softmax=False)

Checklist

Before PR:

• Pre-commit or other linting tools are used to fix the potential lint issues.
• Bug fixes are fully covered by unit tests, the case that causes the bug should be added in the unit tests.
• The modification is covered by complete unit tests. If not, please add more unit test to ensure the correctness.
• The documentation has been modified accordingly, like docstring or example tutorials.

After PR:

• If the modification has potential influence on downstream or other related projects, this PR should be tested with those projects, like MMDet or MMSeg.
• CLA has been signed and all committers have signed the CLA in this PR.

[codecov]

Codecov Report

Merging #593 (f6242a7) into master (f9a2b04) will increase coverage by 0.97%.
The diff coverage is 95.13%.

@@            Coverage Diff             @@
##           master     #593      +/-   ##
==========================================
+ Coverage   80.34%   81.31%   +0.97%     
==========================================
  Files         115      115              
  Lines        6585     6654      +69     
  Branches     1125     1142      +17     
==========================================
+ Hits         5291     5411     +120     
+ Misses       1151     1095      -56     
- Partials      143      148       +5     

Flag	Coverage Δ
unittests	81.31% <95.13%> (+0.97%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
mmcls/models/classifiers/image.py	84.88% <76.92%> (-1.96%)	⬇️
mmcls/models/heads/conformer_head.py	72.13% <86.66%> (+47.60%)	⬆️
mmcls/models/heads/vision_transformer_head.py	96.15% <94.44%> (+2.97%)	⬆️
mmcls/models/heads/deit_head.py	97.36% <95.65%> (-2.64%)	⬇️
mmcls/models/classifiers/base.py	74.39% <100.00%> (+4.76%)	⬆️
mmcls/models/heads/cls_head.py	96.42% <100.00%> (+13.09%)	⬆️
mmcls/models/heads/linear_head.py	93.54% <100.00%> (+8.36%)	⬆️
mmcls/models/heads/multi_label_head.py	100.00% <100.00%> (+28.94%)	⬆️
mmcls/models/heads/multi_label_linear_head.py	93.75% <100.00%> (+25.89%)	⬆️
mmcls/models/heads/stacked_head.py	100.00% <100.00%> (+2.89%)	⬆️
... and 5 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update f9a2b04...f6242a7. Read the comment docs.

[fcakyon]

@mzr1996 this is an awesome PR. Just what I needed! Great job. Cant wait to get it merged @Ezra-Yu 👍

[Ezra-Yu]

1. remove se_cfg in the docstring in StackedLinearClsHead
2. add some docstring and for some functions and class

[Ezra-Yu]

LGTM.

[greattBILL]

Hello, are all the codes based on the default, why is there no information about the trained model and the relative data set? I want to use the model I have trained(such as resnet101) and the matching dataset picture(imagenet) for neck output, can you provide a code demonstration. Thank you very much.

[mzr1996]

Hello, are all the codes based on the default, why is there no information about the trained model and the relative data set? I want to use the model I have trained(such as resnet101) and the matching dataset picture(imagenet) for neck output, can you provide a code demonstration. Thank you very much.

These codes are just simple examples. If you want to use the real dataset, you also need to build the dataset to do some pre-processing. Here is an example:

import mmcv
from mmcv import Config
from mmcls.apis import init_model
from mmcls.datasets import build_dataset

cfg = Config.fromfile('configs/resnet/resnet18_8xb32_in1k.py')
cfg.model.backbone.out_indices = (0, 1, 2, 3)  # Output multi-scale feature maps
model = init_model(cfg, 'your_checkpoint_path', device='cpu')  # Use the API `init_model` to initialize and load checkpoint

dataset = build_dataset(cfg.data.test)  # Use test dataset and test pipeline
img = dataset[0]['img'][None, :]    # Get one input image
outs = model.extract_feat(img, stage='neck')
for out in outs:
    print(out.shape)
# torch.Size([1, 256])
# torch.Size([1, 512])
# torch.Size([1, 1024])
# torch.Size([1, 2048])

[jingmingtao]

import mmcv
from mmcv import Config
from mmcls.apis import init_model
from mmcls.datasets import build_dataset

cfg = Config.fromfile('configs/resnet/resnet18_8xb32_in1k.py')
cfg.model.backbone.out_indices = (0, 1, 2, 3) # Output multi-scale feature maps
model = init_model(cfg, 'your_checkpoint_path', device='cpu') # Use the API init_model to initialize and load checkpoint

dataset = build_dataset(cfg.data.test) # Use test dataset and test pipeline
img = dataset[0]['img'][None, :] # Get one input image
outs = model.extract_feat(img, stage='neck')
for out in outs:
print(out.shape)

torch.Size([1, 256])

torch.Size([1, 512])

torch.Size([1, 1024])

torch.Size([1, 2048])

请教在新版的mmpretrain中如何完成该特征提取的功能，尝试了很久无法完成：
import mmengine
from mmengine import Config
from mmpretrain.apis import init_model
from mmpretrain.datasets import build_dataset

from mmpretrain.models import build_classifier
import torch

cfg = Config.fromfile('configs/resnet/resnet18_8xb32_in1k.py')
cfg.model.backbone.out_indices = (0, 1, 2, 3) # Output multi-scale feature maps
model = init_model(cfg, 'your_checkpoint_path', device='cpu') # Use the API init_model to initialize and load checkpoint

dataset = build_dataset(cfg.data.test) # Use test dataset and test pipeline
img = dataset[0]['img'][None, :] # Get one input image
outs = model.extract_feat(img, stage='neck')
for out in outs:
print(out.shape)

[greattBILL]

这是来自QQ邮箱的假期自动回复邮件。   您好，我最近正在休假中，无法亲自回复您的邮件。我将在假期结束后，尽快给您回复。

[watertianyi]

请教在新版本的 mmpretrain 中如何完成该特征提取的功能，尝试了很久无法完成：

@mzr1996
Please tell me how to complete this feature extraction function in the new version of mmpretrain. I tried for a long time and couldn't complete it:

[greattBILL]

这是来自QQ邮箱的假期自动回复邮件。   您好，我最近正在休假中，无法亲自回复您的邮件。我将在假期结束后，尽快给您回复。