对 IntermediateLayerGetter 理解

初看这个IntermediateLayerGetter名字，（调节中间层获取，自己翻译的），不是很能理解，只有看它的解释再结合给的样例才能理解。

理解：这个类就是获取一个Model中你指定要获取的哪些层的输出，然后这些层的输出会在一个有序的字典中，字典中的key就是刚开始初始化这个类传进去的，value就是feature经过指定需要层的输出。

这是源码中的样例：

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()])


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

代码第四行先传入一个Model，第二个OrderedDict参数，key是指定要获取Model中哪些层的输出，value是这些层的输出会放在一个OrderedDict中，其中这个有序字典中的key就是前面的value，就相当于可以按照自己喜好来设置输出的key。

下面来看看代码是怎么实现的：

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

我们来分析 def __init__(self, model: nn.Module, return_layers: Dict[str, str]) -> None:

这个方法其实就是先遍历Model，一直在里面找我们传入的return_layers层，直到找到return_layers里面最后一个，然后遍历就结束。这里面需要注意，return_layers里面的key的顺序要跟Model中定义的一致。这样找到return_layers中所有key在Model中对应的层（这个层可能时卷卷积，也可能时ReLu，也可能时Sequential，就看Model中怎么定义的）。找到的层存放在

layers = OrderedDict()中，然后用这个layers初始化这个类IntermediateLayerGetter（这个layers中其实存放了Model中从第一层一直到找到的return_layers中的最后一层，就像Model有1-10层，return_layers是{'1':'one', '5':'five'}，这样layers中内容是{'1':'one', '2':'two','3':'three','4':'four','5':'five'}，这时有的同学就问了，我只要1和5啊，别急，关键就是下面的forward了）

再分析 forward(self, x)，在这个方法中定义了一个 out = OrderedDict() 有序字典

虽然上面说layers存放1-5层，但是在forward中，传入一个feature时也是从layers中依次取一层进行卷积啥的，只有对应的层在return_layers才存放在out中，这样out中只有我们需要的层的输出。

一开始以为就是获取我需要的层，比如第三层卷积，第四层ReLu，但是通过init才知道，里面包含从第一层到第四层，只有在前向传播才输出需要的层的结果。这个名字让我以为抽取Model中需要的层，但是看完forward以及样例才理解，是获取需要层的输出结果，就是IntermediateLayerGetter并不是给单独的第三层卷积，第四层ReLu，而是给的feature经过这两个层的输出结果。

当然，我们也可以获取 IntermediateLayerGetter 返回的对象，然后通过逻辑获取对应的层。