support optimizer for LambdaDecay

Author: algorithm1832

python/paddle/optimizer/lr.py

@@ -1509,6 +1509,25 @@ class LambdaDecay(LRScheduler):
 
     lr_lambda: Callable[[int], float]
 
+    @overload
+    def __init__(
+        self,
+        learning_rate: float,
+        lr_lambda: Callable[[int], float],
+        last_epoch: int = -1,
+        verbose: bool = False,
+    ): ...
+
+    @overload
+    def __init__(
+        self,
+        optimizer: paddle.optimizer.Optimizer,
+        lr_lambda: Callable[[int], float],
+        last_epoch: int = -1,
+        verbose: bool = False,
+    ): ...
+
+    @lr_scheduler_decorator()
     def __init__(
         self,
         learning_rate: float,

test/legacy_test/test_lr_scheduler.py

@@ -1463,6 +1463,24 @@ def test_step_decay(self):
             )
         paddle.enable_static()
 
+    def test_lambda_decay(self):
+        paddle.disable_static()
+        linear = paddle.nn.Linear(10, 10)
+        base_lr = 0.5
+        lr_lambda = lambda epoch: 0.95**epoch
+        sgd = paddle.optimizer.SGD(
+            learning_rate=base_lr, parameters=linear.parameters()
+        )
+        scheduler = paddle.optimizer.lr.LambdaDecay(
+            optimizer=sgd, lr_lambda=lr_lambda
+        )
+        self.assertEqual(scheduler.base_lr, sgd.get_lr())
+        self.assertIs(sgd._learning_rate, scheduler)
+        lrs = self._test_network(linear, sgd, scheduler)
+        for i in range(len(lrs)):
+            np.testing.assert_allclose(lrs[i], base_lr * lr_lambda(i))
+        paddle.enable_static()
+
 
 if __name__ == '__main__':
     paddle.enable_static()