| """ | |
| https://github.com/pytorch/vision/blob/main/torchvision/models/_utils.py | |
| Copyright(c) 2023 lyuwenyu. All Rights Reserved. | |
| """ | |
| from collections import OrderedDict | |
| from typing import Dict, List | |
| import torch.nn as nn | |
| 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`. | |
| """ | |
| _version = 3 | |
| def __init__(self, model: nn.Module, return_layers: List[str]) -> None: | |
| if not set(return_layers).issubset([name for name, _ in model.named_children()]): | |
| raise ValueError( | |
| "return_layers are not present in model. {}".format( | |
| [name for name, _ in model.named_children()] | |
| ) | |
| ) | |
| orig_return_layers = return_layers | |
| return_layers = {str(k): str(k) for k in return_layers} | |
| 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): | |
| outputs = [] | |
| for name, module in self.items(): | |
| x = module(x) | |
| if name in self.return_layers: | |
| outputs.append(x) | |
| return outputs | |