From e58ac1215821cafc70f17937e059f3e251f1f1a9 Mon Sep 17 00:00:00 2001
From: Zhanlue Yang <jim19930609@gmail.com>
Date: Tue, 23 Nov 2021 10:36:56 +0800
Subject: [PATCH] =?UTF-8?q?Bug=20fix=20for=20snapshotting=20VariableWrappe?=
 =?UTF-8?q?r=20with=20initialized=20tensor=20but=20e=E2=80=A6=20(#37410)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Bug fix for snapshoting VariableWrapper with initialized tensor but empty allocation

* Added unittest for inplace&clear_gradient
---
 paddle/fluid/imperative/dygraph_grad_maker.h  | 26 +++++++---
 .../test_inplace_and_clear_gradient.py        | 50 +++++++++++++++++++
 2 files changed, 69 insertions(+), 7 deletions(-)
 create mode 100644 python/paddle/fluid/tests/unittests/test_inplace_and_clear_gradient.py

diff --git a/paddle/fluid/imperative/dygraph_grad_maker.h b/paddle/fluid/imperative/dygraph_grad_maker.h
index f1eb8aa62c..f5a4b0fa31 100644
--- a/paddle/fluid/imperative/dygraph_grad_maker.h
+++ b/paddle/fluid/imperative/dygraph_grad_maker.h
@@ -347,15 +347,27 @@ class TracedGradOp {
     //  Use original var_wrapper if its inplace_version is not
     //  changed. Otherwise, it will affect the accuracy of the model
     //  results and affect double grad.
-    if (!var_wrapper->MutableVar()->IsInitialized() ||
-        var_wrapper->InplaceVersionSnapshot() ==
-            var_wrapper->MutableVar()->CurrentInplaceVersion()) {
+    if (!var_wrapper->MutableVar()->IsInitialized()) {
       return var_wrapper;
-    } else {
-      VariableWrapper new_var_wrapper = *var_wrapper.get();
-      new_var_wrapper.ResetInplaceVersion();
-      return std::make_shared<VariableWrapper>(new_var_wrapper);
+    } else if (var_wrapper->InplaceVersionSnapshot() ==
+               var_wrapper->MutableVar()->CurrentInplaceVersion()) {
+      return var_wrapper;
+    } else if (var_wrapper->MutableVar()->IsType<framework::LoDTensor>() ||
+               var_wrapper->MutableVar()->IsType<framework::SelectedRows>()) {
+      auto* tensor =
+          var_wrapper->MutableVar()->IsType<framework::LoDTensor>()
+              ? var_wrapper->MutableVar()->GetMutable<framework::LoDTensor>()
+              : var_wrapper->MutableVar()
+                    ->GetMutable<framework::SelectedRows>()
+                    ->mutable_value();
+      if (!tensor->IsInitialized()) {
+        return var_wrapper;
+      }
     }
+
+    VariableWrapper new_var_wrapper = *var_wrapper.get();
+    new_var_wrapper.ResetInplaceVersion();
+    return std::make_shared<VariableWrapper>(new_var_wrapper);
   }
 
  private:
diff --git a/python/paddle/fluid/tests/unittests/test_inplace_and_clear_gradient.py b/python/paddle/fluid/tests/unittests/test_inplace_and_clear_gradient.py
new file mode 100644
index 0000000000..b685900ead
--- /dev/null
+++ b/python/paddle/fluid/tests/unittests/test_inplace_and_clear_gradient.py
@@ -0,0 +1,50 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import paddle
+import paddle.fluid as fluid
+from paddle import _C_ops
+import unittest
+
+paddle.disable_static()
+
+
+def clear_grad(w, a):
+    @paddle.no_grad()
+    def warp(*_):
+        assert w.grad is not None
+        _C_ops.scale_(w.grad, 'scale', 0.5)
+        w.clear_gradient(False)
+
+    return warp
+
+
+class TestInplaceAndClearGradient(unittest.TestCase):
+    def test(self):
+        paddle.set_device('cpu')
+
+        input_data = np.ones([2, 2]).astype('float32')
+        w = paddle.to_tensor(input_data, 'float32', stop_gradient=False)
+
+        _clear_grad = clear_grad(w, a="1")
+        w._register_backward_hook(_clear_grad)
+
+        for i in range(10):
+            out = _C_ops.scale(w, 'scale', 0.1)
+            out.backward()
+
+
+if __name__ == '__main__':
+    unittest.main()
-- 
GitLab