[Typing] Update _ParameterConfig in Optimizer (#65277)

megemini · web-flow · commit cf2c785ddf5d · 2024-06-20T02:27:20.000+08:00
diff --git a/python/paddle/optimizer/adam.py b/python/paddle/optimizer/adam.py
@@ -18,8 +18,6 @@
 from collections import defaultdict
 from typing import TYPE_CHECKING, Sequence
 
-from typing_extensions import NotRequired
-
 import paddle
 from paddle import _C_ops, pir
 from paddle.base.libpaddle import DataType
@@ -33,20 +31,23 @@
     in_dynamic_or_pir_mode,
     in_pir_mode,
 )
-from .optimizer import Optimizer, _ParameterConfig
-
-
-class _AdamParameterConfig(_ParameterConfig):
-    beta1: NotRequired[float | Tensor]
-    beta2: NotRequired[float | Tensor]
-
+from .optimizer import Optimizer
 
 if TYPE_CHECKING:
+    from typing_extensions import NotRequired
+
     from paddle import Tensor
     from paddle.nn.clip import GradientClipBase
     from paddle.regularizer import WeightDecayRegularizer
 
     from .lr import LRScheduler
+    from .optimizer import _ParameterConfig
+
+    class _AdamParameterConfig(_ParameterConfig):
+        beta1: NotRequired[float | Tensor]
+        beta2: NotRequired[float | Tensor]
+        epsilon: NotRequired[float | Tensor]
+        lazy_mode: NotRequired[bool]
 
 
 __all__ = []
diff --git a/python/paddle/optimizer/adamax.py b/python/paddle/optimizer/adamax.py
@@ -21,17 +21,25 @@
 from ..base import core, framework
 from ..base.dygraph import no_grad
 from ..base.framework import name_scope
-from .adam import _AdamParameterConfig
 from .optimizer import Optimizer
 
-__all__ = []
-
 if TYPE_CHECKING:
+    from typing_extensions import NotRequired
+
     from paddle import Tensor
     from paddle.nn.clip import GradientClipBase
     from paddle.regularizer import WeightDecayRegularizer
 
     from .lr import LRScheduler
+    from .optimizer import _ParameterConfig
+
+    class _AdamaxParameterConfig(_ParameterConfig):
+        beta1: NotRequired[float | Tensor]
+        beta2: NotRequired[float | Tensor]
+        epsilon: NotRequired[float | Tensor]
+
+
+__all__ = []
 
 
 class Adamax(Optimizer):
@@ -69,7 +77,7 @@ class Adamax(Optimizer):
         beta2 (float|Tensor, optional): The exponential decay rate for the 2nd moment estimates.
             It should be a float number or a 0-D Tensor with shape [] and data type as float32.
             The default value is 0.999.
-        epsilon (float, optional): A small float value for numerical stability.
+        epsilon (float|Tensor, optional): A small float value for numerical stability.
             The default value is 1e-08.
         parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` to update to minimize ``loss``.
             This parameter is required in dygraph mode. And you can specify different options for
@@ -154,9 +162,9 @@ def __init__(
         learning_rate: float | LRScheduler = 0.001,
         beta1: float | Tensor = 0.9,
         beta2: float | Tensor = 0.999,
-        epsilon: float = 1e-8,
+        epsilon: float | Tensor = 1e-8,
         parameters: Sequence[Tensor]
-        | Sequence[_AdamParameterConfig]
+        | Sequence[_AdamaxParameterConfig]
         | None = None,
         weight_decay: float | WeightDecayRegularizer | None = None,
         grad_clip: GradientClipBase | None = None,
diff --git a/python/paddle/optimizer/adamw.py b/python/paddle/optimizer/adamw.py
@@ -73,7 +73,7 @@ class AdamW(Optimizer):
         beta2 (float|Tensor, optional): The exponential decay rate for the 2nd moment estimates.
             It should be a float number or a 0-D Tensor with shape [] and data type as float32.
             The default value is 0.999.
-        epsilon (float, optional): A small float value for numerical stability.
+        epsilon (float|Tensor, optional): A small float value for numerical stability.
             The default value is 1e-08.
         parameters (list|tuple|None, optional): List/Tuple of ``Tensor`` names to update to minimize ``loss``.
             This parameter is required in dygraph mode. And you can specify different options for
@@ -170,7 +170,7 @@ def __init__(
         learning_rate: float | LRScheduler = 0.001,
         beta1: float | Tensor = 0.9,
         beta2: float | Tensor = 0.999,
-        epsilon: float = 1e-8,
+        epsilon: float | Tensor = 1e-8,
         parameters: Sequence[Tensor]
         | Sequence[_AdamParameterConfig]
         | None = None,
diff --git a/python/paddle/optimizer/lamb.py b/python/paddle/optimizer/lamb.py
@@ -24,11 +24,23 @@
 from .optimizer import Optimizer
 
 if TYPE_CHECKING:
+    from typing_extensions import NotRequired
+
     from paddle import Tensor
     from paddle.nn.clip import GradientClipBase
 
     from .optimizer import _ParameterConfig
 
+    class _LambParameterConfig(_ParameterConfig):
+        beta1: NotRequired[float | Tensor]
+        beta2: NotRequired[float | Tensor]
+        epsilon: NotRequired[float | Tensor]
+        lamb_weight_decay: NotRequired[float]
+        exclude_from_weight_decay_fn: NotRequired[
+            Callable[[Tensor], bool] | None
+        ]
+
+
 __all__ = []
 
 
@@ -62,14 +74,14 @@ class Lamb(Optimizer):
     learning rate, :math:`\\lambda` the LAMB weight decay rate.
 
     Args:
-        learning_rate (float|Variable, optional): the learning rate used to update parameters. \
+        learning_rate (float|Tensor, optional): the learning rate used to update parameters. \
             Can be a float value or a Variable with data type float32. Default 0.001.
         lamb_weight_decay (float, optional): The LAMB weight decay rate. Default 0.01. Remind that weight_decay should be None.
-        beta1 (float, optional): The exponential decay rate for the 1st moment estimates.
+        beta1 (float|Tensor, optional): The exponential decay rate for the 1st moment estimates.
             Default 0.9.
-        beta2 (float, optional): The exponential decay rate for the 2nd moment estimates.
+        beta2 (float|Tensor, optional): The exponential decay rate for the 2nd moment estimates.
             Default 0.999.
-        epsilon (float, optional): A small float value for numerical stability. Default 1e-6.
+        epsilon (float|Tensor, optional): A small float value for numerical stability. Default 1e-6.
         parameters (list|tuple|None, optional):  Iterable of ``Variable`` names to update to minimize ``loss``. \
             This parameter is required in dygraph mode. And you can specify different options for \
             different parameter groups such as the learning rate, weight decay, etc, \
@@ -98,7 +110,8 @@ class Lamb(Optimizer):
             >>> loss = paddle.mean(out)
             >>> beta1 = paddle.to_tensor([0.9], dtype="float32")
             >>> beta2 = paddle.to_tensor([0.85], dtype="float32")
-            >>> lamb = paddle.optimizer.Lamb(learning_rate=0.002, parameters=linear.parameters(), lamb_weight_decay=0.01)
+            >>> lamb = paddle.optimizer.Lamb(
+            ...     learning_rate=0.002, beta1=beta1, beta2=beta2, parameters=linear.parameters(), lamb_weight_decay=0.01)
             >>> back = out.backward()
             >>> lamb.step()
             >>> lamb.clear_grad()
@@ -113,10 +126,12 @@ def __init__(
         self,
         learning_rate: float | Tensor = 0.001,
         lamb_weight_decay: float = 0.01,
-        beta1: float = 0.9,
-        beta2: float = 0.999,
-        epsilon: float = 1e-6,
-        parameters: Sequence[Tensor] | Sequence[_ParameterConfig] | None = None,
+        beta1: float | Tensor = 0.9,
+        beta2: float | Tensor = 0.999,
+        epsilon: float | Tensor = 1e-6,
+        parameters: Sequence[Tensor]
+        | Sequence[_LambParameterConfig]
+        | None = None,
         grad_clip: GradientClipBase | None = None,
         exclude_from_weight_decay_fn: Callable[[Tensor], bool] | None = None,
         multi_precision: bool = False,
diff --git a/python/paddle/optimizer/optimizer.py b/python/paddle/optimizer/optimizer.py
@@ -20,7 +20,6 @@
 from typing import TYPE_CHECKING, Callable, Sequence
 
 import numpy as np
-from typing_extensions import NotRequired, TypedDict
 
 import paddle
 import paddle.autograd as imperative_base
@@ -50,19 +49,19 @@
 from ..base.layer_helper import LayerHelper, LayerHelperBase
 from .lr import LRScheduler
 
-
-class _ParameterConfig(TypedDict):
-    params: Sequence[Tensor]
-    weight_decay: NotRequired[float | WeightDecayRegularizer | None]
-    learning_rate: NotRequired[float | Tensor | LRScheduler | None]
-
-
 if TYPE_CHECKING:
+    from typing_extensions import NotRequired, TypedDict
+
     from paddle import Tensor
     from paddle.nn.clip import GradientClipBase
 
     from ..base.framework import Operator, Program
 
+    class _ParameterConfig(TypedDict):
+        params: Sequence[Tensor]
+        weight_decay: NotRequired[float | WeightDecayRegularizer | None]
+        learning_rate: NotRequired[float | Tensor | LRScheduler | None]
+
 
 __all__ = []
 
diff --git a/python/paddle/optimizer/sgd.py b/python/paddle/optimizer/sgd.py
@@ -31,6 +31,7 @@
 
     from .lr import LRScheduler
     from .optimizer import _ParameterConfig
+
 __all__ = []
 
 
