PaddlePaddle
diff --git a/‎python/paddle/distributed/auto_parallel/static/helper.py
Copy file name to clipboard
+4-3Lines changed: 4 additions & 3 deletions b/‎python/paddle/distributed/auto_parallel/static/helper.py
Copy file name to clipboard
+4-3Lines changed: 4 additions & 3 deletions
diff --git a/‎python/paddle/jit/api.py
Copy file name to clipboard
+7-6Lines changed: 7 additions & 6 deletions b/‎python/paddle/jit/api.py
Copy file name to clipboard
+7-6Lines changed: 7 additions & 6 deletions
diff --git a/‎python/paddle/jit/dy2static/convert_call_func.py
Copy file name to clipboard
+3-12Lines changed: 3 additions & 12 deletions b/‎python/paddle/jit/dy2static/convert_call_func.py
Copy file name to clipboard
+3-12Lines changed: 3 additions & 12 deletions
diff --git a/‎python/paddle/jit/dy2static/function_spec.py
Copy file name to clipboard
+26-7Lines changed: 26 additions & 7 deletions b/‎python/paddle/jit/dy2static/function_spec.py
Copy file name to clipboard
+26-7Lines changed: 26 additions & 7 deletions
diff --git a/‎python/paddle/jit/dy2static/program_translator.py
Copy file name to clipboard
+41-33Lines changed: 41 additions & 33 deletions b/‎python/paddle/jit/dy2static/program_translator.py
Copy file name to clipboard
+41-33Lines changed: 41 additions & 33 deletions
@@ -64,9 +64,10 @@ def __init__(self, layer, loss_func, metrics):
 
         # Consider ProxyLayer as not Paddle inner function because it contains
         # user-defined layer.
-        as_not_paddle_func(
-            inspect.getmodule(ProxyLayer).__name__ + ".ProxyLayer"
-        )
+        for fn_name in ["_train", "_eval", "_predict"]:
+            as_not_paddle_func(
+                f"{inspect.getmodule(ProxyLayer).__name__}.ProxyLayer.{fn_name}"
+            )
 
     @paddle.jit.not_to_static
     def append_loss_to_shadow_output(self, mode):
 
@@ -133,7 +133,8 @@ def copy_decorator_attrs(original_func, decorated_obj):
 
     decorated_obj.__name__ = original_func.__name__
     decorated_obj._decorator_name = decorator_name
-    decorated_obj.__wrapped__ = original_func
+    if not inspect.ismethod(original_func):
+        decorated_obj.__wrapped__ = original_func
     decorated_obj.__doc__ = original_func.__doc__
     if hasattr(original_func, "__module__"):
         decorated_obj.__module__ = original_func.__module__
@@ -1271,7 +1272,7 @@ def save(
                 )
                 concrete_program = static_function.concrete_program
 
-                if static_function._class_instance is None:
+                if static_function.class_instance is None:
                     warnings.warn(
                         f'`jit.save` will only save the `Program`, not the parameters. If you have to save the parameters, please make sure that {layer} is a member function of `paddle.nn.Layer` and the saved parameters are in `state_dict`'
                     )
@@ -1281,9 +1282,9 @@ def save(
         if isinstance(inner_layer, Layer):
             dygraph_state_dict = inner_layer.to_static_state_dict()
         elif isinstance(attr_func, StaticFunction):
-            if static_func._class_instance:
+            if static_func.class_instance:
                 dygraph_state_dict = (
-                    static_func._class_instance.to_static_state_dict()
+                    static_func.class_instance.to_static_state_dict()
                 )
 
         if dygraph_state_dict:
@@ -1778,9 +1779,9 @@ def set_dynamic_shape(variable, shape_list):
 
 def get_ast_static_function(function):
     if isinstance(function, SymbolicStaticFunction):
-        if function._class_instance:
+        if function.class_instance:
             dygraph_function = types.MethodType(
-                function._dygraph_function, function._class_instance
+                function._dygraph_function, function.class_instance
             )
         else:
             dygraph_function = function._dygraph_function
 
@@ -42,7 +42,7 @@
     convert_to_static,
     unwrap_decorators,
 )
-from .utils import is_builtin, is_paddle_func
+from .utils import WeakMethod, is_builtin, is_paddle_func
 
 if TYPE_CHECKING:
     from types import ModuleType
@@ -192,15 +192,6 @@ def is_unsupported(func):
             )
             return True
 
-    # NOTE: should be placed before `is_paddle_func`
-    # The api(s) should be considered as plain function and convert
-    # them into static layer code.
-    from paddle.nn import Sequential
-
-    PADDLE_NEED_CONVERT_APIS = [Sequential]
-    if type(func) in PADDLE_NEED_CONVERT_APIS:
-        return False
-
     if is_paddle_func(func):
         translator_logger.log(
             2,
@@ -367,11 +358,11 @@ def convert_call(func):
             try:
                 _, forward_func = unwrap_decorators(func.forward)
                 func._original_funcs['forward'] = forward_func.__func__
-                forward_func = convert_to_static(forward_func)
+                forward_func = convert_to_static(forward_func.__func__)
                 # Bound method will be convert into plain function after `convert_to_static`.
                 # So descriptor mechanism is used to bound `self` instance on function to
                 # keep it as bound method.
-                func.forward = forward_func.__get__(func)
+                func.forward = WeakMethod(forward_func, func)
             except (OSError, TypeError):
                 # NOTE: func.forward may have been decorated.
                 func_self = None if func_self else func_self
 
@@ -14,6 +14,7 @@
 
 import collections
 import inspect
+import weakref
 
 import numpy as np
 
@@ -46,7 +47,12 @@ class FunctionSpec:
     """
 
     def __init__(self, function, input_spec=None):
-        self._dygraph_function = function
+        if inspect.ismethod(function):
+            self._dygraph_function = function.__func__
+            self._class_instance = weakref.ref(function.__self__)
+        else:
+            self._dygraph_function = function
+            self._class_instance = None
         if input_spec is None:
             self._input_spec = None
             self._flat_input_spec = None
@@ -83,7 +89,7 @@ def unified_args_and_kwargs(self, args, kwargs):
             New arguments tuple containing default kwargs value.
         """
         if len(self._arg_names) < len(args):
-            error_msg = f"The decorated function `{self._dygraph_function.__name__}` requires {len(self._arg_names)} arguments: {self._arg_names}, but received {len(args)} with {args}."
+            error_msg = f"The decorated function `{self.dygraph_function.__name__}` requires {len(self._arg_names)} arguments: {self._arg_names}, but received {len(args)} with {args}."
             if args and inspect.isclass(args[0]):
                 error_msg += "\n\tMaybe the function has more than one decorator, we don't support this for now."
                 raise NotImplementedError(error_msg)
@@ -100,7 +106,7 @@ def unified_args_and_kwargs(self, args, kwargs):
             else:
                 if arg_name not in self._default_kwargs:
                     raise ValueError(
-                        f"`{self._dygraph_function.__name__}()` requires `{arg_name}` arguments, but not found in input `args`: {args} and `kwargs`: {kwargs}."
+                        f"`{self.dygraph_function.__name__}()` requires `{arg_name}` arguments, but not found in input `args`: {args} and `kwargs`: {kwargs}."
                     )
                 args.append(self._default_kwargs[arg_name])
 
@@ -128,7 +134,7 @@ def args_to_input_spec(self, args, kwargs):
             # So we don't support to deal this case while specifying `input_spec` currently.
             if kwargs:
                 raise ValueError(
-                    f"{self._dygraph_function.__name__} got unexpected keyword arguments: {kwargs}. Cannot trace the function when `input_spec` is specified."
+                    f"{self.dygraph_function.__name__} got unexpected keyword arguments: {kwargs}. Cannot trace the function when `input_spec` is specified."
                 )
 
             # Note: The length of `input_spec` can be greater than `args`,
@@ -279,14 +285,27 @@ def _verify_input_spec(self, input_spec):
 
     def __repr__(self):
         return "function: {}({}), input_spec: {}".format(
-            self._dygraph_function.__name__,
+            self.dygraph_function.__name__,
             ','.join(self._arg_names),
             self._input_spec,
         )
 
+    @property
+    def class_instance(self):
+        if self._class_instance is None:
+            return None
+        if self._class_instance() is None:
+            raise RuntimeError(
+                "The instance of class has been deleted, please re-create the instance."
+            )
+        return self._class_instance()
+
     @property
     def dygraph_function(self):
-        return self._dygraph_function
+        if self.class_instance is not None:
+            return self._dygraph_function.__get__(self.class_instance)
+        else:
+            return self._dygraph_function
 
     @property
     def args_name(self):
@@ -302,7 +321,7 @@ def flat_input_spec(self):
 
     @property
     def code(self):
-        return func_to_source_code(self._dygraph_function)
+        return func_to_source_code(self.dygraph_function)
 
 
 def get_parameters(layer_instance, include_sublayer=True):
 
@@ -364,7 +364,7 @@ def unwrap_decorators(func):
         if isinstance(cur, StaticFunction):
             decorators.append(cur)
             # Note: if `cur` is a method, keep it as bound method of class.
-            instance = cur._class_instance
+            instance = cur.class_instance
             if instance is not None:
                 cur = cur.dygraph_function.__get__(instance)
             else:
@@ -388,14 +388,14 @@ def __init__(self, function, input_spec=None, **kwargs):
 
         if inspect.ismethod(function):
             self._dygraph_function = function.__func__
-            self._class_instance = function.__self__
+            self._class_instance = weakref.ref(function.__self__)
 
-            if not hasattr(self._class_instance, '_original_funcs'):
+            if not hasattr(self.class_instance, '_original_funcs'):
                 raise TypeError(
                     "When using 'to_static' to convert method of a class, "
                     "please ensure the class inherits from nn.Layer"
                 )
-            self._class_instance._original_funcs[function.__name__] = (
+            self.class_instance._original_funcs[function.__name__] = (
                 self._dygraph_function
             )
         else:
@@ -417,8 +417,8 @@ def __init__(self, function, input_spec=None, **kwargs):
 
     def _get_debug_name(self) -> str:
         try:
-            if self._class_instance:
-                self._debug_name = self._class_instance.__class__.__name__
+            if self.class_instance:
+                self._debug_name = self.class_instance.__class__.__name__
             else:
                 self._debug_name = self._dygraph_function.__name__
         except Exception:
@@ -431,8 +431,8 @@ def is_property(self) -> bool:
 
     def train(self) -> None:
         if (
-            isinstance(self._class_instance, layers.Layer)
-            and self._class_instance.training is False
+            isinstance(self.class_instance, layers.Layer)
+            and self.class_instance.training is False
         ):
             raise RuntimeError(
                 f"Failed to switch train mode. {self.dygraph_function} is a Layer's method, "
@@ -442,8 +442,8 @@ def train(self) -> None:
 
     def eval(self) -> None:
         if (
-            isinstance(self._class_instance, layers.Layer)
-            and self._class_instance.training is True
+            isinstance(self.class_instance, layers.Layer)
+            and self.class_instance.training is True
         ):
             raise RuntimeError(
                 f"Failed to switch eval mode. {self.dygraph_function} is a Layer's method, "
@@ -488,7 +488,7 @@ def forward(self, x, y):
                 instance._original_funcs[self._dygraph_function.__name__] = (
                     self._dygraph_function
                 )
-            new_static_layer._class_instance = instance
+            new_static_layer._class_instance = weakref.ref(instance)
             self._descriptor_cache[instance] = new_static_layer
 
         return self._descriptor_cache[instance]
@@ -535,13 +535,13 @@ def __call__(self, *args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
         return self._perform_call(*args, **kwargs)
 
     def _is_train_mode(self) -> bool:
-        if self._class_instance is not None:
-            if not hasattr(self._class_instance, 'training'):
+        if self.class_instance is not None:
+            if not hasattr(self.class_instance, 'training'):
                 raise TypeError(
                     "When using 'to_static' to convert method of a class, "
                     "please ensure the class inherits from nn.Layer"
                 )
-            return self._class_instance.training
+            return self.class_instance.training
         else:
             return self._training
 
@@ -584,13 +584,23 @@ def get_traced_count(self):
     def code(self) -> str:
         raise NotImplementedError("Not implemented yet.")
 
+    @property
+    def class_instance(self):
+        if self._class_instance is None:
+            return None
+        if self._class_instance() is None:
+            raise RuntimeError(
+                "The instance of class has been deleted, please re-create the instance."
+            )
+        return self._class_instance()
+
     @property
     def dygraph_function(self) -> Callable[_InputT, _RetT]:
         """
         Returns the original decorated function.
         """
-        if self._class_instance is not None:
-            return self._dygraph_function.__get__(self._class_instance)
+        if self.class_instance is not None:
+            return self._dygraph_function.__get__(self.class_instance)
         else:
             return self._dygraph_function
 
@@ -645,23 +655,23 @@ def rollback_impl(class_instance):
             for sublayer in class_instance.sublayers(include_self=False):
                 rollback_impl(sublayer)
 
-        if self._class_instance is None:
+        if self.class_instance is None:
             return self._dygraph_function
 
         # only rollback sub-functions on path of top _dygraph_function
         func_name = self._dygraph_function.__name__
         assert (
-            func_name in self._class_instance._original_funcs
-        ), f"Not Found function '{func_name}' in class '{self._class_instance.__class__}'."
-        func = self._class_instance._original_funcs[func_name]
+            func_name in self.class_instance._original_funcs
+        ), f"Not Found function '{func_name}' in class '{self.class_instance.__class__}'."
+        func = self.class_instance._original_funcs[func_name]
         setattr(
-            self._class_instance, func_name, func.__get__(self._class_instance)
+            self.class_instance, func_name, func.__get__(self.class_instance)
         )
 
-        for sublayer in self._class_instance.sublayers(include_self=False):
+        for sublayer in self.class_instance.sublayers(include_self=False):
             rollback_impl(sublayer)
 
-        return getattr(self._class_instance, func_name)
+        return getattr(self.class_instance, func_name)
 
     def __deepcopy__(self, memo):
         """
@@ -695,17 +705,15 @@ def __deepcopy__(self, memo):
 
         Please attention that original 'net' will unwrap @to_static and rollback into simple Layer.
         """
-        if self._class_instance is not None:
-            net_name = type(self._class_instance).__name__
+        if self.class_instance is not None:
+            net_name = type(self.class_instance).__name__
             logging_utils.log(
                 level=-1,
                 msg=f"Not recommend to deepcopy '{net_name}' decorated with @to_static, it has side effect that will"
                 f" rollback into original state before @to_static. Please deepcopy '{net_name}' before applying @to_static.",
             )
             self.rollback()
-            return self._dygraph_function.__get__(
-                memo[id(self._class_instance)]
-            )
+            return self._dygraph_function.__get__(memo[id(self.class_instance)])
         else:
             return self._dygraph_function
 
@@ -771,8 +779,8 @@ def _perform_call(self, *args, **kwargs):
             training=self._is_train_mode(),
             backend=backend,
         )
-        if self._class_instance is not None:
-            args = (self._class_instance, *args)
+        if self.class_instance is not None:
+            args = (self.class_instance, *args)
         return traced_fun(*args, **kwargs)
 
     @property
@@ -831,8 +839,8 @@ def _perform_call(self, *args, **kwargs):
                 *args, **kwargs, is_train=self._is_train_mode()
             )
             # 2. synchronize self.training attribute.
-            if isinstance(self._class_instance, layers.Layer):
-                partial_program_layer.training = self._class_instance.training
+            if isinstance(self.class_instance, layers.Layer):
+                partial_program_layer.training = self.class_instance.training
             else:
                 partial_program_layer.training = self._training
 
@@ -899,7 +907,7 @@ def get_concrete_program(
             self._function_spec,
             input_args_with_spec,
             input_kwargs_with_spec,
-            self._class_instance,
+            self.class_instance,
             **self._kwargs,
             with_hook=with_hook,
             is_train=is_train,