[Dy2St] Fix Tensor hook error on multiple load_attr (#61505)

SigureMo · web-flow · commit d7744aeaddbd · 2024-02-04T11:06:10.000+08:00
* [Dy2St] Fix Tensor hook error on nested load

* add other block support

* remove useless map

* remove print
diff --git a/python/paddle/jit/dy2static/transformers/tensorhook_transformer.py b/python/paddle/jit/dy2static/transformers/tensorhook_transformer.py
@@ -12,17 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import collections
 
 from paddle.utils import gast
 
+from ..utils import ast_to_source_code
 from .base import BaseTransformer
 
 
+def get_loads(node: gast.AST):
+    for child in gast.walk(node):
+        if isinstance(
+            child, (gast.Name, gast.Attribute, gast.Subscript)
+        ) and isinstance(child.ctx, gast.Load):
+            yield child
+
+
 class RegisterHookTransformer(BaseTransformer):
     def __init__(self, root):
-        self.register_hook_pos_map = collections.defaultdict(list)
-        self.assignment_pos_map = collections.defaultdict(list)
         self.root = root
 
     def transform(self):
@@ -31,58 +37,75 @@ def transform(self):
         """
         self.visit(self.root)
 
-    def visit_FunctionDef(self, func_def):
-        # The inner function that has register_hook will not be processed
-        check_register_hook = next(
-            (
-                node
-                for node in gast.walk(func_def)
-                if isinstance(node, gast.Attribute)
-                and node.attr == 'register_hook'
-            ),
-            None,
-        )
-        if check_register_hook is None:
-            return func_def
-
-        register_hook_pos_map = self.register_hook_pos_map
-        assignment_pos_map = self.assignment_pos_map
-
-        for i in range(len(func_def.body) - 1, -1, -1):
-            body = func_def.body[i]
-            # Check if the code body contains the register_hook
-            if isinstance(body, gast.Expr):
-                for node in gast.walk(body):
-                    if (
-                        isinstance(node, gast.Attribute)
-                        and node.attr == 'register_hook'
-                    ):
-                        # parameter name for register_hook
-                        param_name = node.value.id
-                        register_hook_pos_map[param_name].append(i)
-            elif isinstance(body, gast.Assign):
-                for target in body.targets:
-                    assignment_pos_map[target.id].append(i)
-
-        # Confirm the order
-        order_map = {}
-        for k, idx_list in register_hook_pos_map.items():
-            for idx in idx_list:
-                if k not in assignment_pos_map:
-                    order_map[idx] = 1
-                else:
-                    for assignment_idx in assignment_pos_map[k]:
-                        if idx > assignment_idx:
-                            order_map[idx] = assignment_idx + 1
-                            break
-        code_order = [*range(len(func_def.body))]
-        for k, v in sorted(order_map.items(), key=lambda x: x[1], reverse=True):
-            if k == v:
+    def reorder_block_statements(self, stmts):
+        regisiter_hook_nodes = [
+            n
+            for n in stmts
+            for stmt in gast.walk(n)
+            if isinstance(stmt, gast.Attribute) and stmt.attr == "register_hook"
+        ]
+        # Analyze the register_hook nodes name dependency
+        dependents = {}
+        for n in regisiter_hook_nodes:
+            if n not in stmts:
                 continue
-            code_order.remove(k)
-            code_order.insert(v, k)
+            for load_node in get_loads(n):
+                load_name = ast_to_source_code(load_node)
+                if load_name not in dependents:
+                    dependents[load_name] = []
+                dependents[load_name].append(n)
+
+        # Reorder the register_hook nodes, insert it before the dependent nodes
+        idx = 0
+        reordered_stmts = list(stmts)
+        while idx < len(reordered_stmts):
+            stmt = reordered_stmts[idx]
+            loads = get_loads(stmt)
+            for load_node in loads:
+                load_name = ast_to_source_code(load_node)
+                if load_name in dependents:
+                    dep_nodes = dependents[load_name]
+                    for dep_node in dep_nodes:
+                        dep_idx = reordered_stmts.index(dep_node)
+                        if dep_idx <= idx:
+                            continue
+                        reordered_stmts.remove(dep_node)
+                        reordered_stmts.insert(idx, dep_node)
+                        idx += 1
+            idx += 1
+        return reordered_stmts
+
+    def visit_FunctionDef(self, node: gast.FunctionDef):
+        node.body = self.reorder_block_statements(node.body)
+        self.generic_visit(node)
+        return node
+
+    def visit_For(self, node: gast.For):
+        node.body = self.reorder_block_statements(node.body)
+        node.orelse = self.reorder_block_statements(node.orelse)
+        self.generic_visit(node)
+        return node
+
+    def visit_While(self, node: gast.While):
+        node.body = self.reorder_block_statements(node.body)
+        node.orelse = self.reorder_block_statements(node.orelse)
+        self.generic_visit(node)
+        return node
+
+    def visit_If(self, node: gast.If):
+        node.body = self.reorder_block_statements(node.body)
+        node.orelse = self.reorder_block_statements(node.orelse)
+        self.generic_visit(node)
+        return node
+
+    def visit_With(self, node: gast.With):
+        node.body = self.reorder_block_statements(node.body)
+        self.generic_visit(node)
+        return node
 
-        # rearrange the code according to the specified order
-        new_body = [func_def.body[i] for i in code_order]
-        func_def.body = new_body
-        return func_def
+    def visit_Try(self, node: gast.Try):
+        node.body = self.reorder_block_statements(node.body)
+        node.orelse = self.reorder_block_statements(node.orelse)
+        node.finalbody = self.reorder_block_statements(node.finalbody)
+        self.generic_visit(node)
+        return node
diff --git a/python/paddle/jit/dy2static/utils.py b/python/paddle/jit/dy2static/utils.py
@@ -363,6 +363,14 @@ def recover_globals_attribute(src_obj, dst_obj):
         if not (k.startswith('__') and k.endswith('__')):
             dst_globals[k] = v
 
+    # Inject source function closure into destination function globals
+    # Because the destination function is a standalone function, the original
+    # closure of the source function is compiled as LOAD_GLOBAL in the
+    # destination function.
+    src_closure = inspect.getclosurevars(src_obj)
+    for k, v in src_closure.nonlocals.items():
+        dst_globals[k] = v
+
 
 def func_to_source_code(function, dedent=True):
     """
diff --git a/test/dygraph_to_static/test_tensor_hook.py b/test/dygraph_to_static/test_tensor_hook.py
@@ -51,6 +51,73 @@ def h(g):
         loss.backward()
         np.testing.assert_allclose(x.grad.numpy(), x_jit.grad.numpy())
 
+    def test_hook_in_sub_block(self):
+        def f(x):
+            def hook1(grad):
+                return 2 * grad
+
+            def hook2(grad):
+                return 3 * grad
+
+            if x > 1:
+                y = x + 4
+                z = y**2
+                y.register_hook(hook1)
+            else:
+                y = x - 4
+                z = y**3
+                y.register_hook(hook2)
+            return z
+
+        x = paddle.to_tensor([2.0])
+        x.stop_gradient = False
+        loss = f(x)
+        loss.backward()
+
+        x_jit = paddle.to_tensor([2.0])
+        x_jit.stop_gradient = False
+        jit_f = to_static(f)
+        loss = jit_f(x_jit)
+        loss.backward()
+        np.testing.assert_allclose(x.grad.numpy(), x_jit.grad.numpy())
+
+    def test_hook_sub_attr(self):
+        IMAGE_SIZE = 784
+        CLASS_NUM = 10
+
+        def hook(grad):
+            return grad * 2
+
+        class Layer(nn.Layer):
+            def __init__(self):
+                super().__init__()
+                self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
+
+            def forward(self, x):
+                # breakpoint()
+                self._linear.weight.register_hook(hook)
+                y = self._linear(x)
+                return y
+
+        paddle.seed(0)
+        data = np.random.random([IMAGE_SIZE]).astype('float32')
+        x = paddle.to_tensor(data)
+        x.stop_gradient = False
+        layer = Layer()
+        loss = layer(x)
+        loss.backward()
+
+        paddle.seed(0)
+        x_jit = paddle.to_tensor(data)
+        x_jit.stop_gradient = False
+        jit_layer = to_static(Layer())
+        loss = jit_layer(x_jit)
+        loss.backward()
+        np.testing.assert_allclose(
+            layer._linear.weight.grad.numpy(),
+            jit_layer._linear.weight.grad.numpy(),
+        )
+
     def test_hook_for_reassignment_parameter(self):
         def f(x):
             def h(g):
