From 36080ae85d593d21d0c6bba18c3cc63a37b6f006 Mon Sep 17 00:00:00 2001
From: Aurelius84 <liujiezhangbupt@gmail.com>
Date: Wed, 14 Jul 2021 11:41:22 +0800
Subject: [PATCH] [Dy2Stat] Ignore to convert paddle.no_grad (#34136)

* fix paddle.no_grad

* fix paddle.no_grad
---
 .../dygraph_to_static/grad_transformer.py     |  5 +++
 .../unittests/dygraph_to_static/test_grad.py  | 38 ++++++++++++++++---
 2 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/python/paddle/fluid/dygraph/dygraph_to_static/grad_transformer.py b/python/paddle/fluid/dygraph/dygraph_to_static/grad_transformer.py
index f7a59063ae6..272d480c5b7 100644
--- a/python/paddle/fluid/dygraph/dygraph_to_static/grad_transformer.py
+++ b/python/paddle/fluid/dygraph/dygraph_to_static/grad_transformer.py
@@ -83,5 +83,10 @@ def is_grad_api_node(node):
     assert isinstance(node, gast.Call)
     api_name = utils.ast_to_source_code(node.func).strip()
     if utils.is_paddle_api(node):
+        if 'no_grad' in api_name:
+            warnings.warn(
+                "paddle.no_grad is only supported for inference model, and not supported for training under @to_static."
+            )
+            return False
         return api_name.endswith("grad")
     return False
diff --git a/python/paddle/fluid/tests/unittests/dygraph_to_static/test_grad.py b/python/paddle/fluid/tests/unittests/dygraph_to_static/test_grad.py
index ab87beb9e10..6c437962158 100644
--- a/python/paddle/fluid/tests/unittests/dygraph_to_static/test_grad.py
+++ b/python/paddle/fluid/tests/unittests/dygraph_to_static/test_grad.py
@@ -48,6 +48,22 @@ class GradLinearLayer(paddle.nn.Layer):
         return dx
 
 
+class NoGradLinearLayer(paddle.nn.Layer):
+    def __init__(self):
+        super(NoGradLinearLayer, self).__init__()
+        self.linear = paddle.nn.Linear(5, 5, bias_attr=False)
+
+    @paddle.jit.to_static
+    def forward(self, x):
+        x.stop_gradient = False
+
+        with paddle.no_grad():
+            y = self.linear(x)
+
+        out = y + x
+        return out
+
+
 class TestGrad(unittest.TestCase):
     def setUp(self):
         self.func = GradLayer()
@@ -72,15 +88,16 @@ class TestGradLinear(TestGrad):
         self.func = GradLinearLayer()
         self.x = paddle.ones(shape=[10, 2, 5], dtype='float32')
         self.x.stop_gradient = False
+        self.infer_model_path = "double_grad_infer_model"
+        self.train_model_path = "double_grad_train_model"
 
     def test_save_infer_program(self):
-        path = "double_grad_infer_model"
         input_spec = [
             paddle.static.InputSpec(
                 shape=[10, 2, 5], dtype='float32')
         ]
-        paddle.jit.save(self.func, path, input_spec=input_spec)
-        load_func = paddle.jit.load(path)
+        paddle.jit.save(self.func, self.infer_model_path, input_spec=input_spec)
+        load_func = paddle.jit.load(self.infer_model_path)
 
         origin_res = self.func(self.x).numpy()
         load_res = load_func(self.x).numpy()
@@ -96,16 +113,25 @@ class TestGradLinear(TestGrad):
             avg_loss = paddle.mean(paddle.abs(out - 1))
             avg_loss.backward()
             optimizer.minimize(avg_loss)
+            print(self.x.grad.mean())
             self.func.clear_gradients()
 
-        path = "double_grad_train_model"
-        paddle.jit.save(self.func, path)
-        load_func = paddle.jit.load(path)
+        paddle.jit.save(self.func, self.train_model_path)
+        load_func = paddle.jit.load(self.train_model_path)
 
         origin_res = self.func(self.x).numpy()
         load_res = load_func(self.x).numpy()
         self.assertTrue(np.allclose(origin_res, load_res))
 
 
+class TestNoGradLinear(TestGradLinear):
+    def setUp(self):
+        self.func = NoGradLinearLayer()
+        self.x = paddle.ones(shape=[10, 2, 5], dtype='float32')
+        self.x.stop_gradient = False
+        self.infer_model_path = "no_grad_infer_model"
+        self.train_model_path = "no_grad_train_model"
+
+
 if __name__ == '__main__':
     unittest.main()
-- 
GitLab