Add IntermediateLayerGetter

python/paddle/vision/models/utils.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import paddle.nn as nn
+from collections import OrderedDict
+from typing import Dict
 
 def _make_divisible(v, divisor=8, min_value=None):
     """
@@ -30,3 +33,63 @@ def _make_divisible(v, divisor=8, min_value=None):
     if new_v < 0.9 * v:
         new_v += divisor
     return new_v
+
+
+class IntermediateLayerGetter(nn.LayerDict):
+    """
+    Layer wrapper that returns intermediate layers from a model.
+
+    It has a strong assumption that the layers have been registered into the model in the 
+    same order as they are used. This means that one should **not** reuse the same nn.Layer
+    twice in the forward if you want this to work.
+
+    Additionally, it is only able to query sublayer 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.Layer): model on which we will extract the features
+        return_layers (Dict[name, new_name]): a dict containing the names of the layers 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:
+        .. code-block:: python
+
+        >>> m = paddle.vision.models.resnet18(pretrained=True)
+        >>> # extract layer1 and layer3, giving as names `feat1` and feat2`
+        >>> new_m = paddle.vision.models.utils.IntermediateLayerGetter(m,
+        >>>     {'layer1': 'feat1', 'layer3': 'feat2'})
+        >>> out = new_m(paddle.rand([1, 3, 224, 224]))
+        >>> print([(k, v.shape) for k, v in out.items()])
+        >>>     [('feat1', [1, 64, 56, 56]), 
+        >>>      ('feat2', [1, 256, 14, 14])]
+    """
+
+    __annotations__ = {
+        "return_layers": Dict[str, str],
+    }
+
+    def __init__(self, model: nn.Layer, 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