在KaimingNormal/Uniform中增加mode参数 (#73347)

Qin-sx · web-flow · commit 4b11dfdec9f0 · 2025-06-20T11:42:50.000+08:00
* added fan_out para in kaiming normal init

	modified:   python/paddle/nn/initializer/kaiming.py
	modified:   test/legacy_test/test_initializer.py

* updated func name in ValueError

	modified:   python/paddle/nn/initializer/kaiming.py

* added mode para in KaimingUniform for fan_out

	modified:   python/paddle/nn/initializer/kaiming.py
	modified:   test/legacy_test/test_initializer.py
diff --git a/python/paddle/nn/initializer/kaiming.py b/python/paddle/nn/initializer/kaiming.py
@@ -62,6 +62,7 @@ class MSRAInitializer(Initializer):
         seed (int32, optional): random seed. Default is 0.
         negative_slope (float, optional): negative_slope (only used with leaky_relu). Default is 0.0.
         nonlinearity(str, optional): the non-linear function. Default is relu.
+        mode(str, optional): the mode of initialization, can be 'fan_in' or 'fan_out'. When set to 'fan_in', the fan_in parameter is used for initialization. When set to 'fan_out', the out_features of trainable Tensor will be used. Default is 'fan_in'.
 
     Note:
         It is recommended to set fan_in to None for most cases.
@@ -75,6 +76,7 @@ def __init__(
         seed: int = 0,
         negative_slope: float = 0,
         nonlinearity: _NonLinearity = 'relu',
+        mode: str = 'fan_in',
     ) -> None:
         """Constructor for MSRAInitializer"""
         assert uniform is not None
@@ -85,6 +87,17 @@ def __init__(
         self._seed = seed
         self._negative_slope = negative_slope
         self._nonlinearity = nonlinearity
+        self._mode = mode
+        if self._mode not in ['fan_in', 'fan_out']:
+            raise ValueError(
+                "The mode of KaimingNormal/KaimingUniform should be 'fan_in' or 'fan_out', "
+                f"but received {self._mode}."
+            )
+        if self._mode == 'fan_out' and self._fan_in is not None:
+            raise ValueError(
+                "The mode of KaimingNormal/KaimingUniform is 'fan_out', "
+                "but fan_in is set. Please set fan_in to None."
+            )
 
     def forward(
         self, var: paddle.Tensor, block: paddle.pir.Block | None = None
@@ -110,7 +123,10 @@ def forward(
         f_in, f_out = self._compute_fans(var)
 
         # If fan_in is passed, use it
-        fan_in = f_in if self._fan_in is None else self._fan_in
+        if self._mode == 'fan_in':
+            fan_in = f_in if self._fan_in is None else self._fan_in
+        if self._mode == 'fan_out':
+            fan_in = f_out
 
         if self._seed == 0:
             self._seed = block.program.random_seed
@@ -273,6 +289,7 @@ class KaimingNormal(MSRAInitializer):
         fan_in (float32|None, optional): fan_in (in_features) of trainable Tensor, If None, it will be inferred automatically. If you don't want to use in_features of the Tensor, you can set the value of 'fan_in' smartly by yourself. Default is None.
         negative_slope (float, optional): negative_slope (only used with leaky_relu). Default is 0.0.
         nonlinearity(str, optional): the non-linear function. Default is relu.
+        mode(str, optional): the mode of initialization, can be 'fan_in' or 'fan_out'. When set to 'fan_in', the fan_in parameter is used for initialization. When set to 'fan_out', the out_features of trainable Tensor will be used. Default is 'fan_in'.
 
     Note:
         It is recommended to set fan_in to None for most cases.
@@ -294,13 +311,15 @@ def __init__(
         fan_in: float | None = None,
         negative_slope: float = 0.0,
         nonlinearity: str = 'relu',
+        mode: str = 'fan_in',
     ) -> None:
         super().__init__(
             uniform=False,
             fan_in=fan_in,
             seed=0,
             negative_slope=negative_slope,
             nonlinearity=nonlinearity,
+            mode=mode,
         )
 
 
@@ -324,6 +343,7 @@ class KaimingUniform(MSRAInitializer):
         fan_in (float32|None, optional): fan_in (in_features) of trainable Tensor, If None, it will be inferred automatically. If you don't want to use in_features of the Tensor, you can set the value of 'fan_in' smartly by yourself. Default is None.
         negative_slope (float, optional): negative_slope (only used with leaky_relu). Default is 0.0.
         nonlinearity(str, optional): the non-linear function. Default is relu.
+        mode(str, optional): the mode of initialization, can be 'fan_in' or 'fan_out'. When set to 'fan_in', the fan_in parameter is used for initialization. When set to 'fan_out', the out_features of trainable Tensor will be used. Default is 'fan_in'.
 
     Note:
         It is recommended to set fan_in to None for most cases.
@@ -345,11 +365,13 @@ def __init__(
         fan_in: float | None = None,
         negative_slope: float = 0.0,
         nonlinearity: str = 'relu',
+        mode: str = 'fan_in',
     ) -> None:
         super().__init__(
             uniform=True,
             fan_in=fan_in,
             seed=0,
             negative_slope=negative_slope,
             nonlinearity=nonlinearity,
+            mode=mode,
         )
diff --git a/test/legacy_test/test_initializer.py b/test/legacy_test/test_initializer.py
@@ -1671,6 +1671,68 @@ def test_msra_initializer(self, dtype="float32"):
         paddle.enable_static()
 
 
+class TestMSRAInitializerFanoutDygraph(unittest.TestCase):
+    def test_msra_fanout_initializer(self, dtype="float32"):
+        """
+        In dygraph mode, we can use initializer directly to initialize a tensor.
+        """
+        paddle.disable_static()
+
+        tensor = paddle.zeros([16, 1024])
+        tensor.stop_gradient = False
+
+        msra_ = paddle.nn.initializer.KaimingNormal(mode='fan_out')
+        msra_(tensor)
+
+        hist, _ = output_hist(tensor.numpy())
+
+        hist2, _ = output_hist(
+            np.random.normal(0, np.sqrt(2.0 / (1024)), [16, 1024])
+        )
+
+        np.testing.assert_allclose(hist, hist2, rtol=0, atol=0.01)
+        paddle.enable_static()
+
+    def test_msra_invalid_fanout_initializer(self, dtype="float32"):
+        """
+        In dygraph mode, we can use initializer directly to initialize a tensor.
+        """
+        paddle.disable_static()
+
+        tensor = paddle.zeros([16, 1024])
+        tensor.stop_gradient = False
+
+        with self.assertRaises(ValueError):
+            msra_ = paddle.nn.initializer.KaimingNormal(mode='fan')
+            msra_(tensor)
+
+        with self.assertRaises(ValueError):
+            msra_ = paddle.nn.initializer.KaimingNormal(
+                fan_in=1, mode='fan_out'
+            )
+            msra_(tensor)
+
+    def test_msra_uniform_fanout_initializer(self, dtype="float32"):
+        paddle.disable_static()
+
+        tensor = paddle.zeros([16, 1024])
+        tensor.stop_gradient = False
+
+        msra_ = paddle.nn.initializer.KaimingUniform(mode='fan_out')
+        msra_(tensor)
+
+        hist, _ = output_hist(tensor.numpy())
+
+        fan_out = tensor.shape[1]
+        limit = np.sqrt(6.0 / fan_out)
+        theory_data = np.random.uniform(-limit, limit, [16, 1024])
+
+        hist2, _ = output_hist(theory_data)
+
+        np.testing.assert_allclose(hist, hist2, rtol=0, atol=0.01)
+        paddle.enable_static()
+
+
 class TestConsistencyOfDynamicAndStaticGraph(unittest.TestCase):
     def test_order(self):
         paddle.set_device('cpu')
