Understanding of intermediatelayergetter

First look at this IntermediateLayerGetter name ,（ Adjust the middle layer to obtain , Translated by myself ）, Not very understandable , Only by looking at its explanation and combining with the examples given can we understand .

understand ： This class is to get a Model In, you specify which layers of output you want to get , Then the output of these layers will be in an ordered Dictionary , In the dictionary key This is the class just passed in ,value Namely feature After specifying the output of the required layer .

This is the sample in the source code ：

import torchvision
import torch
m = torchvision.models.resnet18(pretrained=True)
new_m = torchvision.models._utils.IntermediateLayerGetter(m, {'layer1': 'feat1', 'layer3': 'feat2'})
out = new_m(torch.rand(1, 3, 224, 224))
print([(k, v.shape) for k, v in out.items()])


 Output ：
[('feat1', torch.Size([1, 64, 56, 56])), ('feat2', torch.Size([1, 256, 14, 14]))]

The fourth line of the code passes in a Model, the second OrderedDict Parameters ,key Is to specify to get Model Which layers in the output ,value Yes, the output of these layers will be placed in a OrderedDict in , In this ordered Dictionary key It's the one in front value, It is equivalent to setting the output according to your preferences key.

Let's see how the code is implemented ：

class IntermediateLayerGetter(nn.ModuleDict):
    """
    Module wrapper that returns intermediate layers from a model

    It has a strong assumption that the modules have been registered
    into the model in the same order as they are used.
    This means that one should **not** reuse the same nn.Module
    twice in the forward if you want this to work.

    Additionally, it is only able to query submodules that are directly
    assigned to the model. So if `model` is passed, `model.feature1` can
    be returned, but not `model.feature1.layer2`.

    Args:
        model (nn.Module): model on which we will extract the features
        return_layers (Dict[name, new_name]): a dict containing the names
            of the modules for which the activations will be returned as
            the key of the dict, and the value of the dict is the name
            of the returned activation (which the user can specify).

    Examples::

        >>> m = torchvision.models.resnet18(pretrained=True)
        >>> # extract layer1 and layer3, giving as names `feat1` and feat2`
        >>> new_m = torchvision.models._utils.IntermediateLayerGetter(m,
        >>>     {'layer1': 'feat1', 'layer3': 'feat2'})
        >>> out = new_m(torch.rand(1, 3, 224, 224))
        >>> print([(k, v.shape) for k, v in out.items()])
        >>>     [('feat1', torch.Size([1, 64, 56, 56])),
        >>>      ('feat2', torch.Size([1, 256, 14, 14]))]
    """
    _version = 2
    __annotations__ = {
        "return_layers": Dict[str, str],
    }

    def __init__(self, model: nn.Module, return_layers: Dict[str, str]) -> None:
        if not set(return_layers).issubset([name for name, _ in model.named_children()]):
            raise ValueError("return_layers are not present in model")
        orig_return_layers = return_layers
        return_layers = {str(k): str(v) for k, v in return_layers.items()}
        layers = OrderedDict()
        for name, module in model.named_children():
            layers[name] = module
            if name in return_layers:
                del return_layers[name]
            if not return_layers:
                break

        super(IntermediateLayerGetter, self).__init__(layers)
        self.return_layers = orig_return_layers

    def forward(self, x):
        out = OrderedDict()
        for name, module in self.items():
            x = module(x)
            if name in self.return_layers:
                out_name = self.return_layers[name]
                out[out_name] = x
        return out

Let's analyze def __init__(self, model: nn.Module, return_layers: Dict[str, str]) -> None:

This method is actually traversing Model, I've been looking for our incoming return_layers layer , Until I find return_layers The last one inside , Then the traversal ends . We need to pay attention to it ,return_layers Inside key In the same order Model As defined in . Find it this way return_layers All in key stay Model The corresponding layer in （ This layer may be convoluted , Or maybe ReLu, Or maybe Sequential, Just look Model How to define ）. The layers found are stored in

layers = OrderedDict() in , Then use this layers Initialize this class IntermediateLayerGetter（ This layers In fact, there are Model From the first floor to the return_layers The last layer in , It's like Model Yes 1-10 layer ,return_layers yes {'1':'one', '5':'five'}, such layers The content of is {'1':'one', '2':'two','3':'three','4':'four','5':'five'}, At this time, some students asked , I just 1 and 5 ah , Don't worry. , The key is the following forward 了 ）

reanalysis forward(self, x), In this method, a out = OrderedDict() Orderly dictionary

Although it says layers Deposit 1-5 layer , But in forward in , Pass in a feature Time is also from layers Take one layer in turn for convolution , Only the corresponding layer is in return_layers It's stored in out in , such out Only the output of the layer we need .

At first, I thought it was to get the layer I needed , For example, the third layer convolution , The fourth level ReLu, But by init not have understood until then , It contains from the first floor to the fourth floor , Only in forward propagation can the results of the required layer be output . The name makes me think Model Layers required in , But it's over forward And examples to understand , Is to get the output result of the required layer , Namely IntermediateLayerGetter It's not for a separate third layer convolution , The fourth level ReLu, It's for you feature The output results after these two layers .

Of course , We can also get  IntermediateLayerGetter Returned object , Then obtain the corresponding layer through logic .