Support extract intermediate feature (#126)

Author: Zeqiang-Lai

classification/README.md

@@ -2,6 +2,17 @@
 
 This folder contains the implementation of the InternImage for image classification.
 
+<!-- TOC -->
+* [Install](#install)
+* [Data Preparation](#data-preparation)
+* [Evaluation](#evaluation)
+* [Training from Scratch on ImageNet-1K](#training-from-scratch-on-imagenet-1k)
+* [Manage Jobs with Slurm.](#manage-jobs-with-slurm)
+* [Training with Deepspeed](#training-with-deepspeed)
+* [Extracting Intermediate Features](#extracting-intermediate-features)
+* [Export](#export)
+<!-- TOC -->
+
 ## Usage
 
 ### Install
@@ -259,6 +270,18 @@ Then, you could use `best.pth` as usual, e.g., `model.load_state_dict(torch.load
 
 > Due to the lack of computational resources, the deepspeed training scripts are currently only verified for the first few epochs. Please fire an issue if you have problems for reproducing the whole training.
 
+### Extracting Intermediate Features
+
+To extract the features of an intermediate layer, you could use `extract_feature.py`. 
+
+For example, extract features of `b.png` from layers `patch_embed` and `levels.0.downsample` and save them to 'b.pth'.
+
+```bash
+python extract_feature.py --cfg configs/internimage_t_1k_224.yaml --img b.png --keys patch_embed levels.0.downsample --save --resume internimage_t_1k_224.pth
+```
+
+
+
 ### Export
 
 To export `InternImage-T` from PyTorch to ONNX, run:

classification/extract_feature.py

@@ -0,0 +1,128 @@
+import functools
+from collections import OrderedDict
+
+
+# using wonder's beautiful simplification:
+# https://stackoverflow.com/questions/31174295/getattr-and-setattr-on-nested-objects/31174427?noredirect=1#comment86638618_31174427
+def rgetattr(obj, attr, *args):
+    def _getattr(obj, attr):
+        return getattr(obj, attr, *args)
+
+    return functools.reduce(_getattr, [obj] + attr.split('.'))
+
+
+class IntermediateLayerGetter:
+    def __init__(self, model, return_layers, keep_output=True):
+        """Wraps a Pytorch module to get intermediate values
+
+        Arguments:
+            model {nn.module} -- The Pytorch module to call
+            return_layers {dict} -- Dictionary with the selected submodules
+            to return the output (format: {[current_module_name]: [desired_output_name]},
+            current_module_name can be a nested submodule, e.g. submodule1.submodule2.submodule3)
+
+        Keyword Arguments:
+            keep_output {bool} -- If True model_output contains the final model's output
+            in the other case model_output is None (default: {True})
+
+        Returns:
+            (mid_outputs {OrderedDict}, model_output {any}) -- mid_outputs keys are
+            your desired_output_name (s) and their values are the returned tensors
+            of those submodules (OrderedDict([(desired_output_name,tensor(...)), ...).
+            See keep_output argument for model_output description.
+            In case a submodule is called more than one time, all it's outputs are
+            stored in a list.
+        """
+        self._model = model
+        self.return_layers = return_layers
+        self.keep_output = keep_output
+
+    def __call__(self, *args, **kwargs):
+        ret = OrderedDict()
+        handles = []
+        for name, new_name in self.return_layers.items():
+            layer = rgetattr(self._model, name)
+
+            def hook(module, input, output, new_name=new_name):
+                if new_name in ret:
+                    if type(ret[new_name]) is list:
+                        ret[new_name].append(output)
+                    else:
+                        ret[new_name] = [ret[new_name], output]
+                else:
+                    ret[new_name] = output
+
+            try:
+                h = layer.register_forward_hook(hook)
+            except AttributeError as e:
+                raise AttributeError(f'Module {name} not found')
+            handles.append(h)
+
+        if self.keep_output:
+            output = self._model(*args, **kwargs)
+        else:
+            self._model(*args, **kwargs)
+            output = None
+
+        for h in handles:
+            h.remove()
+
+        return ret, output
+
+
+def main(args, config):
+    from models import build_model
+    import torchvision.transforms as T
+    from PIL import Image
+
+    model = build_model(config)
+    checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu')
+    model.load_state_dict(checkpoint['model'], strict=False)
+    model.cuda()
+
+    # examples:
+    # return_layers = {
+    #     'patch_embed': 'patch_embed',
+    #     'levels.0.downsample': 'levels.0.downsample',
+    #     'levels.0.blocks.0.dcn': 'levels.0.blocks.0.dcn',
+    # }
+    return_layers = {k: k for k in args.keys}
+    mid_getter = IntermediateLayerGetter(model, return_layers=return_layers, keep_output=True)
+
+    image = Image.open(args.img)
+
+    transforms = T.Compose([
+        T.Resize(config.DATA.IMG_SIZE),
+        T.ToTensor(),
+        T.Normalize(config.AUG.MEAN, config.AUG.STD)
+    ])
+    image = transforms(image)
+    image = image.unsqueeze(0)
+    image = image.cuda()
+
+    mid_outputs, model_output = mid_getter(image)
+
+    for k, v in mid_outputs.items():
+        print(k, v.shape)
+
+    return mid_outputs, model_output
+
+
+if __name__ == '__main__':
+    import argparse
+    import torch
+    from config import get_config
+
+    parser = argparse.ArgumentParser('Get Intermediate Layer Output')
+    parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='Path to config file')
+    parser.add_argument('--img', type=str, required=True, metavar="FILE", help='Path to img file')
+    parser.add_argument("--keys", default=None, nargs='+', help="The intermediate layer's keys you want to save.")
+    parser.add_argument('--resume', help='resume from checkpoint')
+    parser.add_argument('--save', action='store_true', help='Save the results.')
+    args = parser.parse_args()
+    config = get_config(args)
+
+    mid_outputs, model_output = main(args, config)
+
+    if args.save:
+        torch.save(mid_outputs, args.img[:-3] + '.pth')