Refine sequential (#67878)

HydrogenSulfate · web-flow · commit 3862f8303d27 · 2024-09-09T20:23:22.000+08:00
* add append/extend/insert for nn.Sequential

* update unitest for converage

* Update container.py

* update docstring of Sequential and add type hint for arg 'sequential'

* change Sequential to Iterable, support __iter__ and add unitest

* use assertTrue

* use np.allclose
diff --git a/python/paddle/nn/layer/container.py b/python/paddle/nn/layer/container.py
@@ -595,6 +595,20 @@ class Sequential(Layer):
             >>> model2.add_sublayer('l3', paddle.nn.Linear(3, 3))  # add sublayer
             >>> res2 = model2(data)  # sequential execution
 
+            >>> # append single layer at the end of sequential
+            >>> model2 = paddle.nn.Sequential(paddle.nn.Linear(10, 20))
+            >>> model2.append(paddle.nn.Linear(20, 30))
+            >>> res2 = model2(data)  # [30, 30]
+
+            >>> # insert single layer at the given position
+            >>> model2 = paddle.nn.Sequential(paddle.nn.Linear(20, 30))
+            >>> model2.insert(0, paddle.nn.Linear(10, 20))
+            >>> res2 = model2(data)  # [30, 30]
+
+            >>> # extend sequential with given sequence of layer(s) at the end
+            >>> model2 = paddle.nn.Sequential()
+            >>> model2.extend([paddle.nn.Linear(10, 20), paddle.nn.Linear(20, 30)])
+            >>> res2 = model2(data)  # [30, 30]
     """
 
     def __init__(self, *layers: Layer | tuple[str, Layer] | list[Any]) -> None:
@@ -636,3 +650,28 @@ def forward(self, input: Any) -> Any:
         for layer in self._sub_layers.values():
             input = layer(input)
         return input
+
+    def append(self, module: Layer) -> Sequential:
+        self.add_sublayer(str(len(self)), module)
+        return self
+
+    def insert(self, index: int, module: Layer) -> Sequential:
+        if not isinstance(module, Layer):
+            raise AssertionError(f'module should be of type: {Layer}')
+        n = len(self._sub_layers)
+        if not (-n <= index <= n):
+            raise IndexError(f'Index out of range: {index}')
+        if index < 0:
+            index += n
+        for i in range(n, index, -1):
+            self._sub_layers[str(i)] = self._sub_layers[str(i - 1)]
+        self._sub_layers[str(index)] = module
+        return self
+
+    def extend(self, sequential: Iterable[Layer]) -> Sequential:
+        for layer in sequential:
+            self.append(layer)
+        return self
+
+    def __iter__(self) -> Iterator[Layer]:
+        return iter(self._sub_layers.values())
diff --git a/test/legacy_test/test_imperative_container_sequential.py b/test/legacy_test/test_imperative_container_sequential.py
@@ -56,6 +56,58 @@ def test_sequential(self):
             loss2 = paddle.mean(res2)
             loss2.backward()
 
+    def test_append_insert_extend(self):
+        data = np.random.uniform(-1, 1, [5, 10]).astype('float32')
+        with base.dygraph.guard():
+            data = paddle.to_tensor(data)
+
+            model1 = paddle.nn.Sequential()
+            # test append
+            model1.append(Linear(10, 1))
+            model1.append(Linear(1, 2))
+            res1 = model1(data)
+            self.assertListEqual(res1.shape, [5, 2])
+
+            # test insert
+            model1.insert(0, Linear(10, 10))
+            res1 = model1(data)
+
+            # test insert type error(non nn.Layer type)
+            model2 = paddle.nn.Sequential()
+            self.assertRaises(AssertionError, model2.insert, 0, 1)
+
+            # test insert index error(1)
+            model2 = paddle.nn.Sequential()
+            self.assertRaises(IndexError, model2.insert, 1, Linear(10, 10))
+
+            # test insert at negtive index -1
+            model2 = paddle.nn.Sequential()
+            model2.insert(0, Linear(10, 10))
+            self.assertEqual(len(model2), 1)
+
+            # res1 = model1(data)
+
+            # test extend
+            model1.extend([Linear(2, 3), Linear(3, 4)])
+            res1 = model1(data)
+            self.assertListEqual(res1.shape, [5, 4])
+
+            loss1 = paddle.mean(res1)
+            loss1.backward()
+
+            # test __iter__
+            model3 = paddle.nn.Sequential(
+                Linear(10, 1),
+                Linear(1, 2),
+            )
+            output1 = model3(data)
+            output2 = data
+            for layer in model3:
+                output2 = layer(output2)
+            np.testing.assert_allclose(
+                output1.numpy(), output2.numpy(), equal_nan=True
+            )
+
     def test_sequential_list_params(self):
         data = np.random.uniform(-1, 1, [5, 10]).astype('float32')
         with base.dygraph.guard():
