如何优雅地获取神经网络的中间特征？

1.使用方法

IntermediateLayerGetter类在torchvision/models/_utils.py中实现。

一个简单的使用案例如下：

import torch
import torchvision.models as models

original_model = models.resnet18(pretrained=True)
wrapped_model = models._utils.IntermediateLayerGetter(original_model, {'layer1': 'feat1', 'layer3': 'feat2'})
out = wrapped_model(torch.rand(1, 3, 224, 224))
print(out['feat1'].shape)
print(out['feat2'].shape)

IntermediateLayerGetter类在实例化时，对原来的模型类进行了一层封装，且需要传入字典来指示想返回的中间特征名和访问特征时使用的name。

构造IntermediateLayerGetter时需要传入字典，字典的key来源于dict(original_model.named_children()).keys()，对于上例，key来源于：

dict_keys(['conv1', 'bn1', 'relu', 'maxpool', 'layer1', 'layer2', 'layer3', 'layer4', 'avgpool', 'fc'])

传入字典的值是自己定义的字符串，在前向推理结束后的返回结果中，将自定义的字符串作为key来访问对应的中间变量。比如上例传入的字典是{'layer1': 'feat1', 'layer3': 'feat2'}，则得到前向推理输出结果out后，通过out['feat1']访问layer1的输出，通过out['feat2']访问layer3的输出。

2.原理

IntermediateLayerGetter类的源码比较简单，如下：

class IntermediateLayerGetter(nn.ModuleDict):
    _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().__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

本质上来讲，IntermediateLayerGetter的实例在初始化时，使用model.named_children()构造一个OrderDict，再用得到的OrderDict去初始化容器nn.ModuleDict()。

在前向计算时按照nn.ModuleDict()容器的内容，顺序执行里面的模块；只是在执行时，会判断容器中模块的名字【即model.named_children()的key】是否在指定的返回值名字列表中，若在列表中，则保存该中间结果到返回值字典中。

这就是在实例化IntermediateLayerGetter时传入字典的key来源于dict(original_model.named_children()).keys()的原因。

3.局限性

根据前文IntermediateLayerGetter的实现方法以及原理，可以很容易发现使用IntermediateLayerGetter获取网络推理中间结果的局限性：

（1）只能获取model.named_children()级别的模块的输出特征，对于更细分模块的输出特征则无法获取；

（2）模型的顶层必须是可以顺序执行的，因为只有这样才能将model.named_children()获取的模块存到OrderDict中并封装为nn.ModuleDict()。

4.开源工程使用案例

在DETR官方的开源代码中（链接：https://github.com/facebookresearch/detr），在文件models/backbone.py的BackboneBase类中使用了该方法获取其中model的中间结果。

若觉得还不错的话，请点个 “赞” 或 “在看” 吧

论文指导班

论文指导班面向那些没有导师指导、需要升学申博的朋友，指导学员从零开始调研相关方向研究、尝试idea、做实验、写论文，指导老师会提供一些idea、代码实现部分的指导、论文写作指导和修改，但整体仍然是由学员自主完成。需要说明的是，论文指导班并非帮你写论文，或者直接给一篇论文让你挂名，我们不会做任何灰色产业，因此，想直接买论文或挂名的朋友请勿联系。

指导老师：

海外QS Top-60某高校人工智能科学博士在读, 师从IEEE Fellow，曾在多家AI企业担任研究实习生和全职算法研究员，具备极强的学术届和工业界综合背景。研究领域主要包括通用计算机视觉模型的高效设计，训练，部署压缩以及在目标检测，语义分割等下游任务应用，具体包括模型压缩 (知识蒸馏，模型搜索量化剪枝), 通用视觉模型与应用(VIT, 目标检测，语义分割), AI基础理论(AutoML, 数据增广，无监督/半监督/长尾/噪声/联邦学习)等；共发表和审稿中的15余篇SCI国际期刊和顶级会议论文，包括NeurIPS，CVPR, ECCV，ICLR，AAAI, ICASSP等CCF-A/B类会议。发明专利授权2项。

长期担任计算机视觉、人工智能、多媒体领域顶级会议CVPR, ECCV, NeurIPS, AAAI,  ACM MM等审稿人。指导研究生本科生发表SCI, EI，CCF-C类会议和毕业论文累计30余篇，有丰富的保研,申博等方面经验，成功辅导学员赴南洋理工，北大，浙大等深造。

涉及范围：CCF会议A类/SCI一区、CCF会议B类/SCI二区、CCF会议C类/SCI三区、SCI四区、EI期刊、EI会议、核心期刊、研究生毕业设计

报名请扫描下方二维码了解详细情况，备注：“论文班报名”。

如果有其他想要当论文指导老师的朋友，请发简历给我，同样扫描上方二维码，备注：“论文指导老师”。基本条件：已发表两篇以上一作顶会，或3-5篇其他级别的一作论文，学历在985博士及以上。