From 9a6a4fbc04b5c5baa1602c64975dcdc41448cebe Mon Sep 17 00:00:00 2001
From: LielinJiang <50691816+LielinJiang@users.noreply.github.com>
Date: Sun, 16 Aug 2020 15:17:37 +0800
Subject: [PATCH] Add children and named_children for Layer (#26289)

* add children and named_children
---
 python/paddle/fluid/dygraph/layers.py         | 52 ++++++++++++++-
 .../test_imperative_layer_children.py         | 63 +++++++++++++++++++
 2 files changed, 114 insertions(+), 1 deletion(-)
 create mode 100644 python/paddle/fluid/tests/unittests/test_imperative_layer_children.py

diff --git a/python/paddle/fluid/dygraph/layers.py b/python/paddle/fluid/dygraph/layers.py
index 250e2b3b38..1bed04479f 100644
--- a/python/paddle/fluid/dygraph/layers.py
+++ b/python/paddle/fluid/dygraph/layers.py
@@ -161,7 +161,7 @@ class Layer(core.Layer):
 
               print(net.state_dict())
         """
-        for layer in self.sublayers():
+        for layer in self.children():
             layer.apply(fn)
 
         fn(self)
@@ -353,6 +353,56 @@ class Layer(core.Layer):
         ]
         return ret
 
+    def children(self):
+        """Returns an iterator over immediate children layers.
+
+        Yields:
+            Layer: a child layer
+
+        Examples:
+            .. code-block:: python
+
+                import paddle.fluid as fluid
+
+                with fluid.dygraph.guard():
+                    fc1 = fluid.Linear(10, 3)
+                    fc2 = fluid.Linear(3, 10, bias_attr=False)
+                    model = fluid.dygraph.Sequential(fc1, fc2)
+                    
+                    layer_list = list(model.children())
+
+                    print(layer_list)
+
+        """
+        for _, layer in self.named_children():
+            yield layer
+
+    def named_children(self):
+        """Returns an iterator over immediate children layers, yielding both
+        the name of the layer as well as the layer itself.
+
+        Yields:
+            (string, Layer): Tuple containing a name and child layer
+
+        Examples:
+            .. code-block:: python
+
+                import paddle.fluid as fluid
+
+                with fluid.dygraph.guard():
+                    fc1 = fluid.Linear(10, 3)
+                    fc2 = fluid.Linear(3, 10, bias_attr=False)
+                    model = fluid.dygraph.Sequential(fc1, fc2)
+                    for prefix, layer in model.named_children():
+                        print(prefix, layer)
+
+        """
+        memo = set()
+        for name, layer in self._sub_layers.items():
+            if layer is not None and layer not in memo:
+                memo.add(layer)
+                yield name, layer
+
     def sublayers(self, include_sublayers=True):
         """Returns a list of sub layers.
 
diff --git a/python/paddle/fluid/tests/unittests/test_imperative_layer_children.py b/python/paddle/fluid/tests/unittests/test_imperative_layer_children.py
new file mode 100644
index 0000000000..e6d8b052d7
--- /dev/null
+++ b/python/paddle/fluid/tests/unittests/test_imperative_layer_children.py
@@ -0,0 +1,63 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import unittest
+
+import paddle
+import paddle.nn as nn
+import paddle.fluid as fluid
+
+import numpy as np
+
+
+class LeNetDygraph(fluid.dygraph.Layer):
+    def __init__(self):
+        super(LeNetDygraph, self).__init__()
+        self.features = nn.Sequential(
+            nn.Conv2D(
+                1, 6, 3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Pool2D(2, 'max', 2),
+            nn.Conv2D(
+                6, 16, 5, stride=1, padding=0),
+            nn.ReLU(),
+            nn.Pool2D(2, 'max', 2))
+
+    def forward(self, inputs):
+        x = self.features(inputs)
+
+        return x
+
+
+class TestLayerChildren(unittest.TestCase):
+    def test_apply_init_weight(self):
+        with fluid.dygraph.guard():
+            net = LeNetDygraph()
+            net.eval()
+
+            net_layers = nn.Sequential(*list(net.children()))
+            net_layers.eval()
+
+            x = paddle.rand([2, 1, 28, 28])
+
+            y1 = net(x)
+            y2 = net_layers(x)
+
+            np.testing.assert_allclose(y1.numpy(), y2.numpy())
+
+
+if __name__ == '__main__':
+    unittest.main()
-- 
GitLab