[Dy2stat] Add a decorator paddle.jit.not_to_static to support that not to...

[Dy2stat] Add a decorator paddle.jit.not_to_static to support that not to convert a function in Dynamic-to-Static. (#29253)

Usage scenarios：A function could have run successfully in static mode,  you can use it to decorate a function in the following cases:
  1. An unknown error occurs in the dynamic-to-static conversion process of the function;
  2. In the internal implementation of the function, it has two branches: dynamic branch and static branch;
  3. Users don't want to convert the function in the process of dynamic to static.

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

@@ -14,8 +14,6 @@
 
 from __future__ import print_function
 
-__all__ = ['convert_call']
-
 import collections
 import copy
 import functools
@@ -35,6 +33,8 @@ from paddle.fluid.dygraph.dygraph_to_static.program_translator import convert_to
 from paddle.fluid.dygraph.dygraph_to_static.program_translator import unwrap_decorators
 from paddle.fluid.dygraph.layers import Layer
 
+__all__ = ["convert_call"]
+
 # TODO(liym27): A better way to do this.
 BUILTIN_LIKELY_MODULES = [
     collections, pdb, copy, inspect, re, six, numpy, logging
@@ -42,6 +42,22 @@ BUILTIN_LIKELY_MODULES = [
 
 translator_logger = TranslatorLogger()
 
+CONVERSION_OPTIONS = "An attribute for a function that indicates conversion flags of the function in dynamic-to-static."
+
+
+class ConversionOptions(object):
+    """
+    A container for conversion flags of a function in dynamic-to-static.
+
+    Attributes:
+        not_convert(bool): An attribute indicates that the function won't be converted in dynamic-to-static.
+
+    NOTE(liym27): More attributes and methods can be added in this class.
+    """
+
+    def __init__(self, not_convert=False):
+        self.not_convert = not_convert
+
 
 def is_builtin(func):
     if isinstance(func, types.BuiltinFunctionType):
@@ -133,6 +149,14 @@ def convert_call(func):
     # in this case, unwraps it into a raw method or function.
     _, func = unwrap_decorators(func)
 
+    options = getattr(func, CONVERSION_OPTIONS, None)
+    if options is not None and options.not_convert:
+        translator_logger.log(
+            2,
+            "{} is not converted when it is decorated by 'paddle.jit.not_to_static'.".
+            format(func))
+        return func
+
     if is_builtin_len(func):
         return convert_len
 

python/paddle/fluid/dygraph/jit.py

@@ -28,6 +28,7 @@ from paddle.fluid.data_feeder import check_type
 from paddle.fluid.layers.utils import flatten
 from paddle.fluid.dygraph.base import program_desc_tracing_guard, switch_to_static_graph
 from paddle.fluid.dygraph.dygraph_to_static import logging_utils
+from paddle.fluid.dygraph.dygraph_to_static.convert_call_func import ConversionOptions, CONVERSION_OPTIONS
 from paddle.fluid.dygraph.dygraph_to_static.logging_utils import set_code_level, set_verbosity
 from paddle.fluid.dygraph.dygraph_to_static.program_translator import ProgramTranslator, StaticFunction, unwrap_decorators
 from paddle.fluid.dygraph.io import TranslatedLayer, INFER_MODEL_SUFFIX, INFER_PARAMS_SUFFIX, INFER_PARAMS_INFO_SUFFIX
@@ -40,7 +41,7 @@ from paddle.fluid.wrapped_decorator import wrap_decorator
 
 __all__ = [
     'TracedLayer', 'declarative', 'dygraph_to_static_func', 'set_code_level',
-    'set_verbosity', 'save', 'load'
+    'set_verbosity', 'save', 'load', 'not_to_static'
 ]
 
 
@@ -225,6 +226,46 @@ def declarative(function=None, input_spec=None):
     return decorated
 
 
+def not_to_static(func=None):
+    """
+    A Decorator to suppresses the convertion of a function.
+
+    Args:
+        func(callable): The function to decorate.
+
+    Returns:
+        callable: A function which won't be converted in Dynamic-to-Static.
+
+    Examples:
+        .. code-block:: python
+
+            import paddle
+
+            @paddle.jit.not_to_static
+            def func_not_to_static(x):
+                res = x - 1
+                return res
+
+            @paddle.jit.to_static
+            def func(x):
+                if paddle.mean(x) < 0:
+                    out = func_not_to_static(x)
+                else:
+                    out = x + 1
+                return out
+
+            x = paddle.ones([1, 2], dtype='float32')
+            out = func(x)
+            print(out) # [[2. 2.]]
+    """
+    if func is None:
+        return not_to_static
+
+    options = ConversionOptions(not_convert=True)
+    setattr(func, CONVERSION_OPTIONS, options)
+    return func
+
+
 class _SaveLoadConfig(object):
     def __init__(self):
         self._output_spec = None

python/paddle/fluid/tests/unittests/dygraph_to_static/test_recursive_call.py → python/paddle/fluid/tests/unittests/dygraph_to_static/test_convert_call.py

@@ -19,18 +19,22 @@ import unittest
 import logging
 import numpy as np
 
+import paddle
 import paddle.fluid as fluid
 from paddle.fluid.dygraph import ProgramTranslator
-from paddle.fluid.dygraph import declarative
+from paddle.fluid.dygraph.dygraph_to_static.convert_call_func import CONVERSION_OPTIONS
+from test_program_translator import get_source_code
 
 program_translator = ProgramTranslator()
 
 SEED = 2020
 np.random.seed(SEED)
 
+# Situation 1 : test recursive call
+
 
 # Use a decorator to test exception
-@declarative
+@paddle.jit.to_static
 def dyfunc_with_if(x_v):
     if fluid.layers.mean(x_v).numpy()[0] > 5:
         x_v = x_v - 1
@@ -39,7 +43,7 @@ def dyfunc_with_if(x_v):
     return x_v
 
 
-@declarative
+@paddle.jit.to_static
 def nested_func(x_v):
     x_v = fluid.dygraph.to_variable(x_v)
 
@@ -50,7 +54,7 @@ def nested_func(x_v):
     return res
 
 
-@declarative
+@paddle.jit.to_static
 def dyfunc_with_third_library_logging(x_v):
     logging.info('test dyfunc_with_third_library_logging')
     if fluid.layers.mean(x_v).numpy()[0] > 5:
@@ -106,14 +110,14 @@ class MyConvLayer(fluid.dygraph.Layer):
             bias_attr=fluid.ParamAttr(
                 initializer=fluid.initializer.Constant(value=0.5)))
 
-    @declarative
+    @paddle.jit.to_static
     def forward(self, inputs):
         y = dyfunc_with_if(inputs)
         y = lambda_fun(y)
         y = self.dymethod(y)
         return y
 
-    @declarative
+    @paddle.jit.to_static
     def dymethod(self, x_v):
         x_v = fluid.layers.assign(x_v)
         return x_v
@@ -133,7 +137,7 @@ class MyLayer(fluid.dygraph.Layer):
             bias_attr=fluid.ParamAttr(
                 initializer=fluid.initializer.Constant(value=0.5)))
 
-    @declarative
+    @paddle.jit.to_static
     def forward(self, inputs):
         h = self.conv(inputs)
         out = self.fc(h)
@@ -143,15 +147,15 @@ class MyLayer(fluid.dygraph.Layer):
 class TestRecursiveCall2(unittest.TestCase):
     def setUp(self):
         self.input = np.random.random((1, 3, 3, 5)).astype('float32')
-        self.Layer = MyLayer
         self.place = fluid.CUDAPlace(0) if fluid.is_compiled_with_cuda(
         ) else fluid.CPUPlace()
+        self.set_func()
+
+    def set_func(self):
+        self.dygraph_func = MyLayer()
 
     def _run(self):
         with fluid.dygraph.guard():
-            self.dygraph_func = self.Layer()
-            fluid.default_startup_program.random_seed = SEED
-            fluid.default_main_program.random_seed = SEED
             data = fluid.dygraph.to_variable(self.input)
             res = self.dygraph_func(data)
 
@@ -175,14 +179,106 @@ class TestRecursiveCall2(unittest.TestCase):
 
 
 class TestThirdPartyLibrary(TestRecursiveCall2):
-    def _run(self):
-        with fluid.dygraph.guard():
+    def set_func(self):
         self.dygraph_func = dyfunc_with_third_library_logging
-            fluid.default_startup_program.random_seed = SEED
-            fluid.default_main_program.random_seed = SEED
-            data = fluid.dygraph.to_variable(self.input)
-            res = self.dygraph_func(data)
-            return res.numpy()
+
+
+# Situation 2 : test not_to_static
+
+
+def func_sum(x):
+    res = paddle.sum(x)
+    return res
+
+
+@paddle.jit.not_to_static
+def func_not_to_static(x):
+    res = func_sum(x)
+    return res
+
+
+@paddle.jit.to_static
+def func_convert_then_not_to_static(x):
+    y = func_not_to_static(x)
+    return y
+
+
+class TestClass(paddle.nn.Layer):
+    @paddle.jit.not_to_static
+    def called_member(self, x):
+        return paddle.sum(x)
+
+    @paddle.jit.to_static
+    def forward(self, x):
+        y = self.called_member(x)
+        return y
+
+
+class TestNotToConvert(TestRecursiveCall2):
+    def set_func(self):
+        self.dygraph_func = func_not_to_static
+
+    def test_conversion_options(self):
+        options = getattr(self.dygraph_func, CONVERSION_OPTIONS, None)
+        self.assertIsNotNone(options)
+        self.assertTrue(options.not_convert)
+
+
+class TestNotToConvert2(TestRecursiveCall2):
+    def set_func(self):
+        self.dygraph_func = func_convert_then_not_to_static
+
+
+class TestNotToConvert3(TestRecursiveCall2):
+    def set_func(self):
+        self.dygraph_func = TestClass()
+
+
+class TestDynamicToStaticCode(unittest.TestCase):
+    def setUp(self):
+        self.set_func()
+        self.set_answer_func()
+
+    def set_func(self):
+        self.func = func_not_to_static
+
+    def set_answer_func(self):
+        class StaticCode():
+            @paddle.jit.not_to_static
+            def func_not_to_static(x):
+                res = func_sum(x)
+                return res
+
+        self.answer_func = StaticCode.func_not_to_static
+
+    def _get_answer_code(self):
+        return get_source_code(self.answer_func)
+
+    def _get_transformed_code(self):
+        transformed_func = paddle.jit.dy2static.convert_call(self.func)
+        return get_source_code(transformed_func)
+
+    def test_code(self):
+        transformed_code = self._get_transformed_code()
+        answer_code = self._get_answer_code()
+        self.assertEqual(
+            answer_code,
+            transformed_code,
+            msg="\ntransformed_code : \n{}\nanswer_code : \n{}".format(
+                transformed_code, answer_code))
+
+
+class TestDynamicToStaticCode2(TestDynamicToStaticCode):
+    def set_func(self):
+        self.func = func_convert_then_not_to_static
+
+    def set_answer_func(self):
+        class StaticCode():
+            def func_convert_then_not_to_static(x):
+                y = paddle.jit.dy2static.convert_call(func_not_to_static)(x)
+                return y
+
+        self.answer_func = StaticCode.func_convert_then_not_to_static
 
 
 if __name__ == '__main__':

python/paddle/jit/__init__.py

@@ -20,6 +20,7 @@ from ..fluid.dygraph.jit import TracedLayer  #DEFINE_ALIAS
 from ..fluid.dygraph.jit import set_code_level  #DEFINE_ALIAS
 from ..fluid.dygraph.jit import set_verbosity  #DEFINE_ALIAS
 from ..fluid.dygraph.jit import declarative as to_static  #DEFINE_ALIAS
+from ..fluid.dygraph.jit import not_to_static  #DEFINE_ALIAS
 from ..fluid.dygraph import ProgramTranslator  #DEFINE_ALIAS
 from ..fluid.dygraph.io import TranslatedLayer  #DEFINE_ALIAS
 
@@ -27,5 +28,5 @@ from . import dy2static
 
 __all__ = [
     'save', 'load', 'TracedLayer', 'to_static', 'ProgramTranslator',
-    'TranslatedLayer', 'set_code_level', 'set_verbosity'
+    'TranslatedLayer', 'set_code_level', 'set_verbosity', 'not_to_static'
 ]