Support if/else in dygraph_to_static (#22540)

* support nested if/else

* support to derivate returns the parameter list automatically

* polish tranform function of slice

* fix modify x.numpy()[i] slice function

* support to transform ast.node into callable function

* fix get_name_ids bug and add more unittest test=develop

* fix requirements.txt test=develop

* remove useless import statement test=develop

* Fixed version compatibility issues in param of function test=develop

* use decorater to test ast_to_func test=develop

* add textwrap.dedent for source_code test=develop

* polish code comment

* fix compatibility with python2 and python3 test=develop

* fix gast version error test=develop

* fix gast repo test=develop

* polish transfer_from_node_type code test=develop

* add nested_if_else unittest test=develop

* split IfElseTransformer test=develop

* specify gast version test=develop

* fix ast_to_func root type test=develop

python/paddle/fluid/dygraph/dygraph_to_static/ast_transformer.py

@@ -1,4 +1,4 @@
-#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#   Copyright (c) 2020 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.
@@ -15,11 +15,78 @@
 from __future__ import print_function
 
 import gast
+# gast is a generic AST to represent Python2 and Python3's Abstract Syntax Tree(AST).
+# It provides a compatibility layer between the AST of various Python versions,
+# as produced by ast.parse from the standard ast module.
+# See details in https://github.com/serge-sans-paille/gast/
+from .ast_utils import is_control_flow_if, create_cond_node, transform_if_else
 
 from .static_analysis import AstNodeWrapper, StaticAnalysisVisitor
 
 __all__ = ['DygraphToStaticAst']
 
+DECORATOR_NAME = 'dygraph_to_static_output'
+
+
+class IfElseTransformer(gast.NodeTransformer):
+    """
+    Transform if/else statement of Dygraph into Static Graph.
+    """
+
+    def __init__(self, wrapper_root):
+        assert isinstance(
+            wrapper_root, AstNodeWrapper
+        ), "Type of input node should be AstNodeWrapper, but received %s ." % type(
+            wrapper_root)
+        self.wrapper_root = wrapper_root
+        self.root = wrapper_root.node
+        self.new_func_nodes = []
+
+    def ast_visit(self):
+        """
+        Main function to transform AST.
+        """
+        self.visit(self.root)
+        self.after_visit(self.root)
+
+    def visit_If(self, node):
+        assert isinstance(node, gast.If)
+        self.generic_visit(node)
+        if is_control_flow_if(node.test):
+            pred_node = node.test
+            true_func_node, false_func_node, return_name_ids = transform_if_else(
+                node, self.root)
+            self.new_func_nodes += [true_func_node, false_func_node]
+            # create layers.cond
+            new_node = create_cond_node(return_name_ids, pred_node,
+                                        true_func_node, false_func_node)
+            return new_node
+        else:
+            return node
+
+    def visit_Call(self, node):
+        # Remove `numpy()` statement, like `Tensor.numpy()[i]` -> `Tensor[i]`
+        # Todo: should be removed. it may be considered as basic api transformation.
+        if isinstance(node.func, gast.Attribute):
+            attribute = node.func
+            if attribute.attr == 'numpy':
+                node = attribute.value
+        return node
+
+    def after_visit(self, node):
+        """
+        This function will add some postprocessing operations with node.
+        It can be used to add the created `true_fn/false_fn` in front of
+        the node.body before they are called in cond layer.
+        """
+        assert hasattr(node, 'body')
+        # add new ast.funcDef of `if/else`
+        if self.new_func_nodes:
+            node.body = self.new_func_nodes + node.body
+
+    def get_new_func_nodes(self):
+        return self.new_func_nodes
+
 
 class DygraphToStaticAst(gast.NodeTransformer):
     """
@@ -27,12 +94,37 @@ class DygraphToStaticAst(gast.NodeTransformer):
     """
 
     def get_static_ast(self, root):
-        # save root for some analysis may need global AST 
+        # save root for some analysis may need global AST
         self.root = root
         self.static_analysis_root = StaticAnalysisVisitor(
             root).get_node_wrapper_root()
+        self.decorate_func_name = None
         self.transfer_from_node_type(self.static_analysis_root)
         return self.static_analysis_root
 
     def transfer_from_node_type(self, node):
-        print("Not implemented")
+        # Generic transformation
+        self.visit(node.node)
+        # Transform all if/else statement of Dygraph into Static Graph.
+        IfElseTransformer(node).ast_visit()
+
+    def visit_FunctionDef(self, node):
+        if self.decorate_func_name is None:
+            self.decorate_func_name = node.name
+        self.generic_visit(node)
+        # Remove the decorated name of dygraph_to_static
+        if hasattr(node, 'decorator_list'):
+            decorator_list = [
+                d for d in node.decorator_list if d.id != DECORATOR_NAME
+            ]
+            node.decorator_list = decorator_list
+        return node
+
+    def get_module_name(self):
+        """
+        Return the main function name which will be used as module name
+        in ast_to_func.
+        """
+        # Should consider BaseAPITransformer which add new module name in Yamei's PR.
+        assert self.decorate_func_name, "decorate_func_name shall not be None."
+        return self.decorate_func_name

python/paddle/fluid/dygraph/jit.py

@@ -16,10 +16,12 @@ __all__ = ['TracedLayer', 'dygraph_to_static_output']
 
 import gast
 import inspect
+import textwrap
 
 from ..wrapped_decorator import wrap_decorator
 from .base import program_desc_tracing_guard, switch_to_static_graph
 from .dygraph_to_static import DygraphToStaticAst
+from .dygraph_to_static.ast_utils import ast_to_func
 from .layers import Layer
 from paddle.fluid import core
 from paddle.fluid.framework import Program, Block, Variable, _dygraph_tracer, dygraph_only, _dygraph_guard, _current_expected_place, in_dygraph_mode
@@ -54,14 +56,15 @@ def _dygraph_to_static_output_(dygraph_func):
     def __impl__(*args, **kwargs):
         # Get AST from dygraph function
         dygraph_code = inspect.getsource(dygraph_func)
+        dygraph_code = textwrap.dedent(dygraph_code)
         root = gast.parse(dygraph_code)
+        # Transform AST
+        dygraph_to_static = DygraphToStaticAst()
+        root_wrapper = dygraph_to_static.get_static_ast(root)
+        func_name = dygraph_to_static.get_module_name()
 
-        root = DygraphToStaticAst().get_static_ast(root)
+        static_func, file_name = ast_to_func(root_wrapper.node, func_name)
 
-        # TODO static_func should a callable from AST, like
-        # static_func = ast_to_func(root)
-        # currently just use dygraph_func
-        static_func = dygraph_func
         return static_func(*args, **kwargs)
 
     return __impl__

python/paddle/fluid/tests/unittests/test_ast_util.py

@@ -0,0 +1,165 @@
+#   Copyright (c) 2020 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.
+
+from __future__ import print_function
+
+import unittest
+import textwrap
+import gast
+import inspect
+import numpy as np
+import paddle.fluid as fluid
+from paddle.fluid.dygraph.dygraph_to_static.ast_utils import get_name_ids, ast_to_func
+
+
+class TestGetNameIds(unittest.TestCase):
+    """
+    Test for parsing the ast.Name list from the ast.Nodes
+    """
+
+    def setUp(self):
+        self.source = """
+          def test_fn(x):
+            return x+1
+        """
+        self.all_name_ids = {'x': [gast.Param()]}
+
+    def test_get_name_ids(self):
+        source = textwrap.dedent(self.source)
+        root = gast.parse(source)
+        all_name_ids = get_name_ids([root])
+        self.assertDictEqual(
+            self.transfer_dict(self.all_name_ids),
+            self.transfer_dict(all_name_ids))
+
+    def transfer_dict(self, name_ids_dict):
+        new_dict = {}
+        for name, ctxs in name_ids_dict.items():
+            new_dict[name] = [type(ctx) for ctx in ctxs]
+        return new_dict
+
+
+class TestGetNameIds2(TestGetNameIds):
+    def setUp(self):
+        self.source = """
+          def test_fn(x, y):
+            a = 1
+            x = y + a
+            if x > y:
+               z = x * x
+               z = z + a
+            else:
+               z = y * y
+            return z
+        """
+        self.all_name_ids = {
+            'x': [
+                gast.Param(), gast.Store(), gast.Load(), gast.Load(),
+                gast.Load()
+            ],
+            'a': [gast.Store(), gast.Load(), gast.Load()],
+            'y':
+            [gast.Param(), gast.Load(), gast.Load(), gast.Load(), gast.Load()],
+            'z': [gast.Store(), gast.Load(), gast.Store(), gast.Store()]
+        }
+
+
+class TestGetNameIds3(TestGetNameIds):
+    def setUp(self):
+        self.source = """
+          def test_fn(x, y):
+            z = 1
+            if x > y:
+               z = x * x
+               z = z + y
+            return z
+        """
+        self.all_name_ids = {
+            'x': [gast.Param(), gast.Load(), gast.Load(), gast.Load()],
+            'y': [gast.Param(), gast.Load(), gast.Load()],
+            'z': [gast.Store(), gast.Store(), gast.Load(), gast.Store()]
+        }
+
+
+def dyfunc_with_if_else(x_v):
+    if fluid.layers.mean(x_v).numpy()[0] > 5:
+        x_v = x_v - 1
+    else:
+        x_v = x_v + 1
+    return x_v
+
+
+def dyfunc_with_if_else2(x):
+    i, j = 0, 0
+    if fluid.layers.reduce_mean(x).numpy()[0] > x.numpy()[i][j]:
+        y = fluid.layers.relu(x)
+    else:
+        x_pow = fluid.layers.pow(x, 2)
+        y = fluid.layers.tanh(x_pow)
+    return y
+
+
+class TestAST2Func(unittest.TestCase):
+    """
+    TestCase for the transformation from ast.AST into python callable function.
+    """
+
+    def _ast2func(self, func):
+        source = inspect.getsource(func)
+        source = textwrap.dedent(source)
+        ast_root = gast.parse(source)
+        transformed_func, _ = ast_to_func(ast_root, func.__name__)
+        return transformed_func
+
+    def test_ast2func(self):
+        def func(x, y):
+            return x + y
+
+        x, y = 10, 20
+        self.assertEqual(func(x, y), self._ast2func(func)(x, y))
+
+    def test_ast2func_dygraph(self):
+        func = dyfunc_with_if_else
+        x_data = np.random.random([10, 16]).astype('float32')
+        with fluid.dygraph.guard():
+            x_v = fluid.dygraph.to_variable(x_data)
+            true_ret = func(x_v).numpy()
+            test_ret = self._ast2func(func)(x_v).numpy()
+            self.assertTrue((true_ret == test_ret).all())
+
+    def test_ast2func_static(self):
+        def func(x):
+            y = fluid.layers.relu(x)
+            loss = fluid.layers.mean(y)
+            return loss
+
+        x_data = np.random.random([10, 16]).astype('float32')
+        main_program = fluid.Program()
+        with fluid.program_guard(main_program):
+            x_v = fluid.layers.assign(x_data)
+            true_ret = func(x_v)
+            test_ret = self._ast2func(func)(x_v)
+            exe = fluid.Executor(fluid.CPUPlace())
+            ret = exe.run(main_program, fetch_list=[true_ret, test_ret])
+            self.assertTrue((ret[0] == ret[1]).all())
+
+    def test_ast2func_error(self):
+        with self.assertRaises(Exception) as e:
+            self.assertRaises(TypeError, ast_to_func("x = a + b", 'foo'))
+        self.assertTrue("Type of ast_root should be gast.AST or ast.AST" in
+                        str(e.exception))
+
+
+if __name__ == '__main__':
+    unittest.main()

python/paddle/fluid/tests/unittests/test_dygraph_to_static_basic.py

@@ -16,8 +16,6 @@ from __future__ import print_function
 
 import numpy as np
 import paddle.fluid as fluid
-import paddle.fluid.layers as layers
-import paddle.fluid.core as core
 import unittest
 
 from paddle.fluid.dygraph.jit import dygraph_to_static_output
@@ -25,37 +23,85 @@ from paddle.fluid.dygraph.jit import dygraph_to_static_output
 np.random.seed(1)
 
 
-def dyfunc(a, b):
+def dyfunc_with_if_else(x_v):
-    with fluid.dygraph.guard():
+    if fluid.layers.mean(x_v).numpy()[0] > 5:
-        x = fluid.dygraph.to_variable(a)
+        x_v = x_v - 1
-        y = fluid.dygraph.to_variable(b)
+    else:
-        x.stop_gradient = False
+        x_v = x_v + 1
-        y.stop_gradient = False
+    return x_v
 
-        inputs = {'X': [x], 'Y': [y]}
-        loss = core.ops.elementwise_mul(inputs)['Out'][0]
 
-        loss.backward()
+def dyfunc_with_if_else2(x):
-        x_grad = x.gradient()
+    i, j = 0, 0
-        y_grad = y.gradient()
+    if fluid.layers.reduce_mean(x).numpy()[0] > x.numpy()[i][j]:
-        return x_grad, y_grad
+        y = fluid.layers.relu(x)
+    else:
+        x_pow = fluid.layers.pow(x, 2)
+        y = fluid.layers.tanh(x_pow)
+    return y
 
 
-@dygraph_to_static_output
+def nested_if_else(x_v):
-def dyfunc_to_static(a, b):
+    batch_size = x_v.shape[0]
-    return dyfunc(a, b)
+    feat_size = x_v.shape[-1]
+    bias = fluid.layers.fill_constant([feat_size], dtype='float32', value=1)
+    if fluid.layers.mean(x_v).numpy()[0] < 0:
+        y = x_v + bias
+        w = fluid.layers.fill_constant([feat_size], dtype='float32', value=10)
+        if y.numpy()[0] < 10:
+            tmp = y * w
+            y = fluid.layers.relu(tmp)
+            if fluid.layers.mean(y).numpy()[0] < batch_size:
+                y = fluid.layers.abs(y)
+            else:
+                tmp = fluid.layers.fill_constant(
+                    [feat_size], dtype='float32', value=-1)
+                y = y - tmp
+    else:
+        y = x_v - bias
+    return y
 
 
-class TestBasicModel(unittest.TestCase):
+class TestDygraphIfElse(unittest.TestCase):
-    def test_dygraph_static_same_output(self):
+    """
-        a = np.random.uniform(
+    TestCase for the transformation from control flow `if/else`
-            low=0.1, high=1, size=(3, 4, 5)).astype(np.float32)
+    dependent on tensor in Dygraph into Static `fluid.layers.cond`.
-        b = np.random.uniform(
+    """
-            low=0.1, high=1, size=(3, 4, 5)).astype(np.float32)
+
-        dy_output = dyfunc(a, b)
+    def setUp(self):
-        static_output = dyfunc_to_static(a, b)
+        self.x = np.random.random([10, 16]).astype('float32')
-        self.assertTrue(np.array_equal(dy_output[0], static_output[0]))
+        self.dyfunc = dyfunc_with_if_else
-        self.assertTrue(np.array_equal(dy_output[1], static_output[1]))
+
+    def _run_static(self):
+        main_program = fluid.Program()
+        with fluid.program_guard(main_program):
+            x_v = fluid.layers.assign(self.x)
+            # Transform into static graph
+            out = dygraph_to_static_output(self.dyfunc)(x_v)
+            exe = fluid.Executor(fluid.CPUPlace())
+            ret = exe.run(main_program, fetch_list=out)
+            return ret
+
+    def _run_dygraph(self):
+        with fluid.dygraph.guard():
+            x_v = fluid.dygraph.to_variable(self.x)
+            ret = self.dyfunc(x_v)
+            return ret.numpy()
+
+    def test_ast_to_func(self):
+        self.assertTrue((self._run_dygraph() == self._run_static()).all())
+
+
+class TestDygraphIfElse2(TestDygraphIfElse):
+    def setUp(self):
+        self.x = np.random.random([10, 16]).astype('float32')
+        self.dyfunc = dyfunc_with_if_else2
+
+
+class TestDygraphIfElse3(TestDygraphIfElse):
+    def setUp(self):
+        self.x = np.random.random([10, 16]).astype('float32')
+        self.dyfunc = nested_if_else
 
 
 if __name__ == '__main__':

python/requirements.txt

@@ -2,6 +2,7 @@ requests>=2.20.0
 numpy>=1.12, <=1.16.4 ; python_version<"3.5"
 numpy>=1.12 ; python_version>="3.5"
 protobuf>=3.1.0
+gast>=0.3.3
 matplotlib<=2.2.4 ; python_version<"3.6"
 scipy>=0.19.0, <=1.2.1 ; python_version<"3.5"
 nltk>=3.2.2, <=3.4 ; python_version<"3.5"
@@ -17,5 +18,4 @@ pyyaml
 decorator
 prettytable
 objgraph
-gast
 astor