【Hackathon 7th No.23】NO.23 为 Paddle 新增 ParameterDict API -part

python/paddle/nn/__init__.py

@@ -77,7 +77,13 @@
 
 # TODO: import all neural network related api under this directory,
 # including layers, linear, conv, rnn etc.
-from .layer.container import LayerDict, LayerList, ParameterList, Sequential
+from .layer.container import (
+    LayerDict,
+    LayerList,
+    ParameterDict,
+    ParameterList,
+    Sequential,
+)
 from .layer.conv import (
     Conv1D,
     Conv1DTranspose,
@@ -243,6 +249,7 @@
     'TransformerEncoder',
     'Softmax',
     'Softmax2D',
+    'ParameterDict',
     'ParameterList',
     'Conv2D',
     'Softshrink',

python/paddle/nn/layer/container.py

@@ -319,6 +319,112 @@ def update(
                 self.add_sublayer(kv[0], kv[1])
 
 
+class ParameterDict(Layer):
+    """
+    Holds parameters in a dictionary.
+
+    ParameterDict can be indexed like a regular Python dictionary, but Parameters it contains are properly registered.
+
+    Parameters:
+        values (iterable, optional): a mapping (dictionary) of (string : Any) or an iterable of key-value pairs of type (string, Any)
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+
+            >>> class MyLayer(paddle.nn.Layer):
+            ...     def __init__(self, num_stacked_param):
+            ...         super().__init__()
+            ...         # create ParameterDict with iterable Parameters
+            ...         self.params = paddle.nn.ParameterDict(
+            ...             {f"t{i}": paddle.create_parameter(shape=[2, 2], dtype='float32') for i in range(num_stacked_param)})
+            ...
+            ...     def forward(self, x):
+            ...         for i, (key, p) in enumerate(self.params):
+            ...             x = paddle.matmul(x, p)
+            ...         return x
+            ...
+            >>> x = paddle.uniform(shape=[5, 2], dtype='float32')
+            >>> num_stacked_param = 4
+            >>> model = MyLayer(num_stacked_param)
+            >>> print(len(model.params))
+            4
+            >>> res = model(x)
+            >>> print(res.shape)
+            [5, 2]
+
+            >>> replaced_param = paddle.create_parameter(shape=[2, 3], dtype='float32')
+            >>> model.params['t3'] = replaced_param  # replace t3 param
+            >>> res = model(x)
+            >>> print(res.shape)
+            [5, 3]
+            >>> model.params['t4'] = paddle.create_parameter(shape=[3, 4], dtype='float32')  # append param
+            >>> print(len(model.params))
+            5
+            >>> res = model(x)
+            >>> print(res.shape)
+            [5, 4]
+    """
+
+    def __init__(
+        self,
+        parameters: (
+            ParameterDict
+            | Mapping[str, Tensor]
+            | Sequence[tuple[str, Tensor]]
+            | None
+        ) = None,
+    ) -> None:
+        super().__init__()
+        if parameters is not None:
+            self.update(parameters)
+
+    def __getitem__(self, key: str) -> Tensor:
+        with param_guard(self._parameters):
+            return self._parameters[key]
+
+    def __setitem__(self, key: str, param: Tensor) -> None:
+        assert isinstance(param, Parameter)
+        setattr(self, key, param)
+
+    def __len__(self) -> int:
+        return len(self._parameters)
+
+    def __iter__(self) -> Iterator[tuple[str, Tensor]]:
+        with param_guard(self._parameters):
+            return iter(self._parameters)
+
+    def update(
+        self,
+        parameters: (
+            ParameterDict
+            | Mapping[str, Tensor]
+            | Sequence[tuple[str, Tensor]]
+        ),
+    ) -> None:
+        """Update a given parameter at the end of the dict.
+
+        Parameters:
+            parameters (Parameter): parameter to update
+        """
+        assert isinstance(parameters, Iterable), (
+            "The type of parameters is not iterable of key/value pairs, the type of sublayers is "
+            + type(parameters).__name__
+        )
+
+        if isinstance(parameters, (OrderedDict, ParameterDict, Mapping)):
+            for key, parameter in parameters.items():
+                self.add_parameter(key, parameter)
+        else:
+            for i, kv in enumerate(parameters):
+                if len(kv) != 2:
+                    raise ValueError(
+                        f"The length of the {i}'s element in parameters is {len(kv)}, which must be 2."
+                    )
+                self.add_parameter(kv[0], kv[1])
+
+
 class ParameterList(Layer):
     """ParameterList Container.
 

test/legacy_test/test_imperative_container_parameterdict.py

@@ -0,0 +1,90 @@
+# Copyright (c) 2024 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.
+
+import unittest
+from collections import OrderedDict
+
+import numpy as np
+
+import paddle
+from paddle import base
+
+
+class MyLayer(paddle.nn.Layer):
+    def __init__(self, num_stacked_param):
+        super().__init__()
+        # create ParameterList with iterable Parameters
+        self.params = self.paddle_imperative_ParameterDict(num_stacked_param)
+
+    def paddle_imperative_ParameterDict(self, num_stacked_param):
+        return paddle.nn.ParameterDict(
+            [
+                (
+                    't' + str(i),
+                    paddle.create_parameter(shape=[2, 2], dtype='float32'),
+                )
+                for i in range(num_stacked_param)
+            ]
+        )
+
+    def forward(self, x):
+        for i, (key, _) in enumerate(self.params):
+            x = paddle.matmul(x, self.params[key])
+        return x
+
+
+class TestImperativeContainerParameterDict(unittest.TestCase):
+    def paramter_dict(self):
+        self.place = (
+            paddle.CUDAPlace(0)
+            if paddle.is_compiled_with_cuda()
+            else paddle.CPUPlace()
+        )
+        data_np = np.random.uniform(-1, 1, [5, 2]).astype('float32')
+        with base.dygraph.guard():
+            x = paddle.to_tensor(data_np)
+            num_stacked_param = 4
+            model = MyLayer(num_stacked_param)
+            self.assertEqual(len(model.params), num_stacked_param)
+            res = model(x)
+            self.assertListEqual(res.shape, [5, 2])
+            loss = paddle.mean(res)
+            loss.backward()
+
+            model.params['t' + str(num_stacked_param - 1)] = (
+                paddle.create_parameter(shape=[2, 3], dtype='float32')
+            )
+            res = model(x)
+            self.assertListEqual(res.shape, [5, 3])
+            parmeter = OrderedDict(
+                [
+                    (
+                        't' + str(num_stacked_param),
+                        paddle.create_parameter(shape=[3, 4], dtype='float32'),
+                    )
+                ]
+            )
+            model.params.update(parmeter)
+            self.assertEqual(len(model.params), num_stacked_param + 1)
+            res = model(x)
+            self.assertListEqual(res.shape, [5, 4])
+            loss = paddle.mean(res)
+            loss.backward()
+
+    def test_paramter_dict(self):
+        self.paramter_dict()
+
+
+if __name__ == '__main__':
+    unittest.main()