add arg differentiable for Optimizer, SGD

Author: algorithm1832

python/paddle/optimizer/optimizer.py

@@ -218,6 +218,8 @@ def __init__(
         weight_decay: float | WeightDecayRegularizer | None = None,
         grad_clip: GradientClipBase | None = None,
         name: str | None = None,
+        *,
+        differentiable: bool = False,
     ) -> None:
         if parameters is not None:
             # paddle.Tensor is also iterable, so here we don't check whether
@@ -331,6 +333,13 @@ def __init__(
         self._fuse_buffer_version = 0
         self.merged_model_params = None
 
+        grad_decorator = (
+            imperative_base.enable_grad()
+            if differentiable
+            else imperative_base.no_grad()
+        )
+        self.step = grad_decorator(self.step)
+
     def _create_master_grad_states(self):
         # master gradients states
         if in_pir_mode():
@@ -2040,7 +2049,6 @@ def _declarative_step(self):
         params_grads = [(param, param.grad) for param in parameters]
         optimize_ops = self.apply_gradients(params_grads)
 
-    @imperative_base.no_grad()
     @framework.non_static_only
     def step(
         self, closure: Callable[[], Tensor] | None = None

python/paddle/optimizer/sgd.py

@@ -98,6 +98,8 @@ def __init__(
         grad_clip: GradientClipBase | None = None,
         multi_precision: bool = False,
         name: str | None = None,
+        *,
+        differentiable: bool = False,
     ) -> None:
         if learning_rate is None:
             raise ValueError("learning_rate is not set")
@@ -107,6 +109,7 @@ def __init__(
             weight_decay=weight_decay,
             grad_clip=grad_clip,
             name=name,
+            differentiable=differentiable,
         )
         self.type = "sgd"
         self._multi_precision = multi_precision

test/legacy_test/test_optimizer.py

@@ -182,6 +182,24 @@ def test_weight_decay_int(self):
         adam.step()
         adam.zero_grad(False)
 
+    def test_differentiable(self):
+        paddle.seed(100)
+        numpy.random.seed(100)
+        paddle.disable_static()
+        x = paddle.arange(26, dtype="float32").reshape([2, 13])
+        x.stop_gradient = False
+        linear = paddle.nn.Linear(13, 5)
+        optimizer = paddle.optimizer.SGD(
+            learning_rate=0.01,
+            parameters=linear.parameters(),
+            differentiable=True,
+        )
+        optimizer.zero_grad()
+        output = linear(x)
+        loss = paddle.mean(output)
+        loss.backward()
+        optimizer.step()
+
     def test_step_without_closure(self):
         paddle.seed(100)
         numpy.random.seed(100)