mindspore 2.7.0__cp310-cp310-win_amd64.whl → 2.7.0rc1__cp310-cp310-win_amd64.whl

mindspore/nn/cell.py

@@ -39,11 +39,10 @@ from typing import (
 
 import weakref
 import mindspore as ms
-import mindspore.ops as ops
 from mindspore._checkparam import args_type_check, check_hook_fn
-from mindspore.common.dynamic_shape._auto_dynamic import is_auto_dynamic, convert_inputs_to_dynamic
+from mindspore.common._auto_dynamic import is_auto_dynamic, convert_inputs_to_dynamic
 from mindspore import log as logger
-from mindspore.common.hook_handle import HookHandle, _update_hook_version
+from mindspore.common.hook_handle import HookHandle
 from mindspore import context
 from mindspore._c_expression import init_pipeline, update_func_graph_hyper_params, Cell_, FuncGraph, MixedPrecisionType
 from mindspore import _checkparam as Validator
@@ -93,8 +92,9 @@ def register_cell_buffer_registration_hook(hook: Callable[..., None],):
         A handle that can be used to remove the added hook by calling
         `handle.remove()`.
     """
-    handle = HookHandle(_global_buffer_registration_hooks)
-    _global_buffer_registration_hooks[handle.handle_id] = hook
+    from mindspore.utils.hooks import _RemovableHandle
+    handle = _RemovableHandle(_global_buffer_registration_hooks)
+    _global_buffer_registration_hooks[handle.id] = hook
     return handle
 
 
@@ -155,8 +155,7 @@ class Cell(Cell_):
     IGNORE_LIST = ['_scope', '_cell_init_args', '_auto_prefix', '_cells', '_params', '_create_time',
                    '_func_graph_flags', '_parameter_layout_dict', '_params_list', '_phase', '_bprop_debug',
                    '_forward_pre_hook', '_forward_hook', '_backward_pre_hook', '_backward_hook',
-                   '_cell_backward_pre_hook', '_cell_backward_hook', '_param_prefix',
-                   'requires_grad', 'cell_type', '_in_strategy', '_out_strategy']
+                   '_cell_backward_pre_hook', '_cell_backward_hook', '_param_prefix', 'requires_grad', 'cell_type']
     total_instance_count = 0
     _buffers: Dict[str, Optional[Tensor]]
     global_cells = weakref.WeakKeyDictionary()
@@ -207,7 +206,6 @@ class Cell(Cell_):
         super().__setattr__("mixed_precision_type", None)
         super().__setattr__("_lazy_construct_sig", None)
         super().__setattr__("_jit_graph_name", '')
-        super().__setattr__("_compiled", False)
         init_pipeline()
 
         # call gc to release GE session resources used by non-used cell objects
@@ -241,8 +239,6 @@ class Cell(Cell_):
         super().__setattr__("_amp_level", "")
         super().__setattr__("_init_flag", False)
         super().__setattr__("_shard_fn", None)
-        super().__setattr__("_in_strategy", None)
-        super().__setattr__("_out_strategy", None)
         super().__setattr__("has_bprop", False)
         if hasattr(self, "bprop"):
             super().__setattr__("has_bprop", True)
@@ -430,13 +426,6 @@ class Cell(Cell_):
         """
         return self._bprop_debug
 
-    @property
-    def compiled(self):
-        """
-        Get whether `Cell` is compiled in graph mode.
-        """
-        return self._compiled
-
     @bprop_debug.setter
     def bprop_debug(self, value):
         """
@@ -557,23 +546,10 @@ class Cell(Cell_):
 
     @property
     def pipeline_segment(self):
-        """
-        `pipeline_segment` represents the pipeline segment of current Cell.
-        """
         return self._pipeline_segment
 
     @pipeline_segment.setter
     def pipeline_segment(self, value):
-        """
-        Set the `pipeline_segment` of a Cell. Only effective in zero_bubble_v scheduler.
-
-        Args:
-            value (int): The pipeline segment of a parameter.
-
-        Raises:
-            TypeError: If `value` is not int type or is a bool type.
-            ValueError: If `value` is not a positive integer.
-        """
         if not isinstance(value, int) or isinstance(value, bool):
             raise TypeError("For 'context.set_auto_parallel_context', the argument 'pipeline_stages' "
                             "must be int type, but got type : {}".format(type(value)))
@@ -1051,13 +1027,12 @@ class Cell(Cell_):
         if self._forward_pre_hook:
             args, kwargs = self._run_forward_pre_hook(args, kwargs)
 
-        if self._backward_hook:
-            args = self._cell_backward_hook(args)
-
         if self._shard_fn is not None:
             output = self._shard_fn(*args, **kwargs)
         elif _pynative_executor.requires_grad():
-            if self._recompute_cell is not None:
+            if self._backward_hook:
+                output = self._backward_hook_construct(*args, **kwargs)
+            elif self._recompute_cell is not None:
                 output = self._recompute_cell(*args, **kwargs)
             elif self.has_bprop:
                 output = self._call_custom_bprop(*args, **kwargs)
@@ -1069,11 +1044,8 @@ class Cell(Cell_):
         if self._forward_hook:
             output = self._run_forward_hook(args, kwargs, output)
 
-        if self._backward_hook:
-            output = self._cell_backward_hook(output)
-
-        if self._backward_pre_hook:
-            output = self._cell_backward_pre_hook(output)
+        if self._backward_pre_hook and _pynative_executor.requires_grad():
+            output = self._run_backward_pre_hook(output)
 
         return output
 
@@ -1108,6 +1080,23 @@ class Cell(Cell_):
                             f"{default_args} default argument, total {positional_args + default_args}, "
                             f"but got {len(args)}.")
 
+    # pylint: disable=E0203
+    def _hook_fn_registered(self):
+        '''Hook function in graph mode'''
+        # Check super().__init__() in graph mode.
+        try:
+            if self._forward_pre_hook or self._forward_hook or self._backward_pre_hook or self._backward_hook:
+                return True
+        except AttributeError as e:
+            raise AttributeError(f"The '{type(self).__name__}' object does not inherit attribute from 'cell'. "
+                                 f"Please use 'super().__init__()'.") from e
+        if not self._is_recursion_hook:
+            self._is_recursion_hook = True
+            for cell in self.cells():
+                if cell._hook_fn_registered():
+                    return True
+        return False
+
     def _get_prims_recursively(self):
         all_prims = list()
         for _, value in self._primitives.items():
@@ -1133,6 +1122,9 @@ class Cell(Cell_):
             >>> net = nn.Dense(3, 4)
             >>> net.set_data_parallel()
         """
+        if context._get_mode() == context.PYNATIVE_MODE:
+            raise ValueError("set_data_parallel: does not support PyNative mode.")
+
         all_prims = self._get_prims_recursively()
         for prim in all_prims:
             prim.add_prim_attr("strategy_gen_mode", "data_parallel")
@@ -1211,6 +1203,8 @@ class Cell(Cell_):
             ...             out = self.blocks[i](out)
             ...         return out
         """
+        if context._get_mode() == context.PYNATIVE_MODE:
+            raise ValueError("The Cell offload does not support PyNative mode now.")
         if isinstance(backward_prefetch, str):
             Validator.check_string(backward_prefetch, ['Auto'], 'backward_prefetch', self.cls_name)
         else:
@@ -1218,10 +1212,11 @@ class Cell(Cell_):
         for prim in self._get_prims_recursively():
             prim._offload(backward_prefetch=backward_prefetch)
 
-    def shard(self, in_strategy, out_strategy=None, parameter_plan=None):
+    def shard(self, in_strategy, out_strategy=None, parameter_plan=None, device="Ascend", level=0):
         """
         Defining the input and output layouts of this cell and the parallel strategies of remaining ops will be
-        generated by sharding propagation. In Graph mode, use this method to specify distribution strategy for a Cell,
+        generated by sharding propagation. In PyNative mode, use this method to specify a Cell for distributed
+        execution in graph mode. In Graph mode, use this method to specify distribution strategy for a Cell,
         strategy for others will be set by sharding propagation.
         in_strategy and out_strategy define the input and output layout respectively.
         in_strategy/out_strategy should be a tuple, each element of which corresponds to the desired layout of
@@ -1233,14 +1228,11 @@ class Cell(Cell_):
               In other parallel modes, strategies set here will be ignored.
             - If the input contain Parameter, its strategy should be set in `in_strategy`.
 
-        .. warning::
-            The method is currently not supported in PyNative mode.
-
         Args:
             in_strategy (tuple): Define the layout of inputs, each element of the tuple should be a tuple. Tuple
                                  defines the layout of the corresponding input.
             out_strategy (Union[None, tuple]): Define the layout of outputs similar with in_strategy.
-                                               Default: ``None`` .
+                                               It is not in use right now. Default: ``None`` .
             parameter_plan (Union[dict, None]): Define the layout for the specified parameters. Each element in dict
                                                 defines the layout of the parameter like "param_name: layout".
                                                 The key is a parameter name of type 'str'.
@@ -1248,6 +1240,14 @@ class Cell(Cell_):
                                                 If the parameter name is incorrect or the corresponding parameter
                                                 has been set, the parameter setting will be ignored.
                                                 Default: ``None`` .
+            device (str): Select a certain device target. It is not in use right now.
+                             Support [ ``"CPU"`` , ``"GPU"`` , ``"Ascend"`` ]. Default: ``"Ascend"`` .
+            level (int): Option for parallel strategy infer algorithm, namely the object function, maximize computation
+                         over communication ratio, maximize speed performance, minimize memory usage etc. It is not in
+                         use right now. Support [ ``"0"`` , ``"1"`` , ``"2"`` ]. Default: ``0`` .
+
+        Returns:
+            Function, return the cell construct function that will be executed under auto parallel process.
 
         Examples:
             >>> import mindspore.nn as nn
@@ -1265,34 +1265,19 @@ class Cell(Cell_):
             ...   def __init__(self):
             ...     self.block1 = Block()
             ...     self.block2 = Block()
-            ...     self.block2.shard(in_strategy=((2, 1),), parameter_plan={'self.block2.dense1.weight': (4, 1)})
+            ...     self.block2_shard = self.block2.shard(in_strategy=((2, 1),),
+            ...                                           parameter_plan={'self.block2.shard.dense1.weight': (4, 1)})
             ...   def construct(self, x):
             ...     x = self.block1(x)
-            ...     x = self.block2(x)
+            ...     x = self.block2_shard(x)
             ...     return x
         """
         if ms.communication.management.get_group_size() == 1:
-            return
-
+            return self
         shard_fn = Shard()
-        self._shard_fn = shard_fn(self, in_strategy, out_strategy, parameter_plan)
-
-        if self._in_strategy is not None:  # pylint: disable=E0203
-            msg = (
-                "For '%s', 'Shard' has been configured more than once. "
-                "The existing in_strategy is %s and the existing out_strategy is %s. "
-                "The new in_strategy %s and out_strategy %s may not take effect. "
-                "It is recommended to configure 'Shard' only once."
-            ) % (
-                self._cell_tag,
-                self._in_strategy,  # pylint: disable=E0203
-                self._out_strategy,  # pylint: disable=E0203
-                shard_fn.in_strategy,
-                shard_fn.out_strategy,
-            )
-            logger.warning(msg)
-        self._in_strategy = shard_fn.in_strategy
-        self._out_strategy = shard_fn.out_strategy
+        fn = shard_fn(self, in_strategy, out_strategy, parameter_plan, device, level)
+        self._shard_fn = fn
+        return fn
 
     def _init_check(self):
         for param in self.get_parameters(expand=False):
@@ -1301,13 +1286,9 @@ class Cell(Cell_):
         self._init_flag = True
 
     def _self_check(self):
-        try:
-            if not self._is_check_and_refresh:  # pylint: disable=E0203
-                self.check_names_and_refresh_name()
-                self._is_check_and_refresh = True
-        except AttributeError as e:
-            raise AttributeError(f"The '{type(self).__name__}' object does not inherit attribute from 'cell'. "
-                                 f"Please use 'super().__init__()'.") from e
+        if not self._is_check_and_refresh:
+            self.check_names_and_refresh_name()
+            self._is_check_and_refresh = True
 
     def _predict(self, *args, **kwargs):
         '''Graph executor for predict'''
@@ -1328,7 +1309,6 @@ class Cell(Cell_):
     def __call__(self, *args, **kwargs):
         # Run in Graph mode.
         if context._get_mode() == context.GRAPH_MODE and os.getenv("MS_JIT") != '0':
-            self._compiled = True
             if kwargs:
                 bound_arguments = self._construct_sig.bind(*args, **kwargs)
                 bound_arguments.apply_defaults()
@@ -1339,8 +1319,11 @@ class Cell(Cell_):
             if predict_compiled:
                 return res
             self._check_construct_args(*args)
+
+            if self._hook_fn_registered():
+                logger.warning(f"For 'Cell', it's not support hook function in graph mode. If you want to use hook "
+                               f"function, please use context.set_context to set pynative mode.")
             self._self_check()
-            self.__compile_cell_hook__ = True
             out = self.compile_and_run(*args, **kwargs)
             return out
 
@@ -1438,7 +1421,16 @@ class Cell(Cell_):
                 exist_names.add(item.name)
                 self.insert_param_to_cell(item.name, item, check_name_contain_dot=False)
 
-            object.__setattr__(self, name, value)
+            if context._get_mode() == context.PYNATIVE_MODE:
+                if name in self.__dict__:
+                    del self.__dict__[name]
+                params = self.__dict__.get('_params')
+                if name in params:
+                    del params[name]
+                params_list = self.__dict__.get('_params_list')
+                params_list[name] = value
+            else:
+                object.__setattr__(self, name, value)
 
     def _set_attr_for_parameter_in_list_or_tuple(self, name, value):
         """Set attr for parameter in list or tuple."""
@@ -1617,6 +1609,8 @@ class Cell(Cell_):
                 _pynative_executor.set_dynamic_input(self, *self._dynamic_shape_inputs)
             else:
                 self._check_construct_args(*inputs)
+                # TODO(tronzhang): It may error for no actually args here. So just set in fullmode,
+                #                  which means that incremental mode is lacking dynamic input.
         else:
             self._dynamic_shape_inputs = _process_dyn_args(self.construct, kwargs)
 
@@ -2604,7 +2598,6 @@ class Cell(Cell_):
             raise ValueError(f"Negative 'fusion_size' {fusion_size} is invalid.")
         Tensor._flatten_tensors(self.trainable_params(), fusion_size)  # pylint: disable=W0212
 
-    @jit_forbidden_register
     def register_forward_pre_hook(self, hook_fn, with_kwargs=False):
         """
         Register forward pre hook function for Cell object.
@@ -2624,6 +2617,7 @@ class Cell(Cell_):
           `with_kwargs` is ``True`` .
 
         Note:
+            - The feature does not take effect in graph mode or in PyNative mode with functions decorated by jit.
             - The `hook_fn` can modify the forward inputs by returning new inputs. If `with_kwargs` is ``Flase`` , a
               single value (whick will be wrapped into a tuple unless already a tuple) or a tuple of args should be
               returned. If `with_kwargs` is ``True`` , both `args` and `kwargs` should be returned.
@@ -2674,15 +2668,15 @@ class Cell(Cell_):
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32,
             value= [ 2.00000000e+00]))
         """
+        if context._get_mode() == context.GRAPH_MODE:
+            return HookHandle()
         check_hook_fn(hook_fn)
         handle = HookHandle(self._forward_pre_hook, extra_dict=self._forward_pre_hook_with_kwargs)
         self._forward_pre_hook[handle.handle_id] = hook_fn
         if with_kwargs:
             self._forward_pre_hook_with_kwargs[handle.handle_id] = True
-        _update_hook_version()
         return handle
 
-    @jit_forbidden_register
     def _run_forward_pre_hook(self, args, kwargs):
         """
         Running forward pre hook function registered on Cell object.
@@ -2706,35 +2700,6 @@ class Cell(Cell_):
                     args = ret
         return args, kwargs
 
-    def _jit_forward_pre_hook(self, inputs):
-        """
-        Compile forward pre hook function registered on Cell object.
-
-        Args:
-            inputs: The input objects of cell object.
-
-        Returns:
-            - **outputs** - New input objects or none.
-
-        Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-        """
-        forward_pre_hook_inputs = inputs
-        for fn in self._forward_pre_hook.values():
-            ret = fn(self, forward_pre_hook_inputs)
-            if ret is not None:
-                if not isinstance(ret, tuple):
-                    forward_pre_hook_inputs = (ret,)
-                else:
-                    forward_pre_hook_inputs = ret
-
-        if len(forward_pre_hook_inputs) != len(inputs):
-            raise TypeError(
-                "The forward pre hook return value size is {} not equal to input size {}".format(
-                    len(forward_pre_hook_inputs), len(inputs)))
-        return forward_pre_hook_inputs
-
-    @jit_forbidden_register
     def register_forward_hook(self, hook_fn, with_kwargs=False):
         """
         Register forward hook function for Cell object.
@@ -2755,6 +2720,7 @@ class Cell(Cell_):
         - `output`: Output generated by the `construct` function.
 
         Note:
+            - The feature does not take effect in graph mode or in PyNative mode with functions decorated by jit.
             - The `hook_fn` can modify the forward outputs by returning new outputs.
             - In order to prevent running failed when switching to graph mode, it is not recommended to call it in the
               `construct` function of Cell object.
@@ -2807,44 +2773,15 @@ class Cell(Cell_):
         """
         if self.has_bprop:
             return HookHandle()
+        if context._get_mode() == context.GRAPH_MODE:
+            return HookHandle()
         check_hook_fn(hook_fn)
         handle = HookHandle(self._forward_hook, extra_dict=self._forward_hook_with_kwargs)
         self._forward_hook[handle.handle_id] = hook_fn
         if with_kwargs:
             self._forward_hook_with_kwargs[handle.handle_id] = True
-        _update_hook_version()
         return handle
 
-    def _jit_forward_hook(self, inputs, output):
-        """
-        Compile forward hook function registered on Cell object.
-
-        Args:
-            inputs: The input objects of Cell object.
-            output: The output object of Cell object.
-
-        Returns:
-            - **output** - New output object or none.
-
-        Supported Platforms:
-        ``Ascend`` ``GPU`` ``CPU``
-        """
-        forward_hook_output = output
-        for fn in self._forward_hook.values():
-            ret = fn(self, inputs, forward_hook_output)
-            if ret is not None:
-                forward_hook_output = ret
-
-        if isinstance(output, tuple):
-            if not isinstance(forward_hook_output, tuple):
-                forward_hook_output = (forward_hook_output,)
-            if len(forward_hook_output) != len(output):
-                raise TypeError(
-                    "The forward hook return value size is {} not equal to output size {}".format(
-                        len(forward_hook_output), len(output)))
-        return forward_hook_output
-
-    @jit_forbidden_register
     def _run_forward_hook(self, args, kwargs, output):
         """
         Running forward hook function registered on Cell object.
@@ -2858,12 +2795,12 @@ class Cell(Cell_):
                 output = ret
         return output
 
-    @jit_forbidden_register
     def register_backward_pre_hook(self, hook_fn):
         """
         Register the backward pre hook function.
 
         Note:
+            - The `register_backward_pre_hook(hook_fn)` does not work in graph mode or functions decorated with 'jit'.
             - The 'hook_fn' must be defined as the following code.
               `cell` is the Cell object. `grad_output` is the gradient passed to the Cell.
             - The 'hook_fn' should have the following signature:
@@ -2912,17 +2849,44 @@ class Cell(Cell_):
             >>> print(output)
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]),)
         """
+        if context._get_mode() == context.GRAPH_MODE:
+            return HookHandle()
         check_hook_fn(hook_fn)
-        handle = HookHandle(self._backward_pre_hook, extra_dict=None)
+        handle = HookHandle(self._backward_pre_hook)
         self._backward_pre_hook[handle.handle_id] = hook_fn
-        if self._cell_backward_pre_hook is None:  # pylint: disable=E0203
+        if self._cell_backward_pre_hook is None:
             # Generate a CellBackwardHook prim, and add function for it
             self._cell_backward_pre_hook = inner.CellBackwardHook(self.cls_name + "(" + str(id(self)) + ")",
                                                                   self, self._backward_pre_hook)
             self._cell_backward_pre_hook.register_backward_pre_hook()
-        _update_hook_version()
         return handle
 
+    def _run_backward_pre_hook(self, outputs):
+        """
+        Running backward pre hook function registered on Cell object.
+
+        Args:
+            outputs: The output objects of cell object.
+
+        Returns:
+            - **outputs** - New backward gradient or None.
+
+        Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+        """
+        if isinstance(outputs, tuple):
+            ret = self._cell_backward_pre_hook(*outputs)
+        else:
+            ret = self._cell_backward_pre_hook(outputs)
+        if isinstance(outputs, tuple):
+            if len(outputs) == 1:
+                ret = (ret,)
+            if len(ret) != len(outputs):
+                raise TypeError(
+                    "The backward pre hook return value size is {} not equal to output size {}".format(
+                        len(ret), len(outputs)))
+        return ret
+
     def get_extra_state(self) -> Any:
         """Return any extra state to include in the cell's state_dict.
 
@@ -2975,8 +2939,9 @@ class Cell(Cell_):
             A handle that can be used to remove the added hook by calling
             `handle.remove()`.
         """
-        handle = HookHandle(self._state_dict_hooks)
-        self._state_dict_hooks[handle.handle_id] = hook
+        from mindspore.utils.hooks import _RemovableHandle
+        handle = _RemovableHandle(self._state_dict_hooks)
+        self._state_dict_hooks[handle.id] = hook
         return handle
 
     @jit_forbidden_register
@@ -3022,8 +2987,9 @@ class Cell(Cell_):
             >>> print("extra_param" in net_state_dict)
             True
         """
-        handle = HookHandle(self._state_dict_pre_hooks)
-        self._state_dict_pre_hooks[handle.handle_id] = hook
+        from mindspore.utils.hooks import _RemovableHandle
+        handle = _RemovableHandle(self._state_dict_pre_hooks)
+        self._state_dict_pre_hooks[handle.id] = hook
         return handle
 
     def _save_to_state_dict(self, destination, prefix, keep_vars):
@@ -3169,8 +3135,9 @@ class Cell(Cell_):
             A handle that can be used to remove the added hook by calling
             `handle.remove()`.
         """
-        handle = HookHandle(self._load_state_dict_pre_hooks)
-        self._load_state_dict_pre_hooks[handle.handle_id] = hook
+        from mindspore.utils.hooks import _RemovableHandle
+        handle = _RemovableHandle(self._load_state_dict_pre_hooks)
+        self._load_state_dict_pre_hooks[handle.id] = hook
         return handle
 
     @jit_forbidden_register
@@ -3202,8 +3169,9 @@ class Cell(Cell_):
             A handle that can be used to remove the added hook by calling
             `handle.remove()`.
         """
-        handle = HookHandle(self._load_state_dict_post_hooks)
-        self._load_state_dict_post_hooks[handle.handle_id] = hook
+        from mindspore.utils.hooks import _RemovableHandle
+        handle = _RemovableHandle(self._load_state_dict_post_hooks)
+        self._load_state_dict_post_hooks[handle.id] = hook
         return handle
 
     def _load_from_state_dict(
@@ -3439,12 +3407,12 @@ class Cell(Cell_):
             )
         return _IncompatibleKeys(missing_keys, unexpected_keys)
 
-    @jit_forbidden_register
     def register_backward_hook(self, hook_fn):
         """
         Register the backward hook function.
 
         Note:
+            - The `register_backward_hook(hook_fn)` does not work in graph mode or functions decorated with 'jit'.
             - The 'hook_fn' must be defined as the following code.
               `cell` is the registered Cell object. `grad_input` is the gradient computed and passed to
               the next Cell or primitive, which can be return a new gradient or None. `grad_output` is the gradient
@@ -3496,17 +3464,65 @@ class Cell(Cell_):
             >>> print(output)
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]),)
         """
+        if context._get_mode() == context.GRAPH_MODE:
+            return HookHandle()
         check_hook_fn(hook_fn)
-        handle = HookHandle(self._backward_hook, extra_dict=None)
+        handle = HookHandle(self._backward_hook)
         self._backward_hook[handle.handle_id] = hook_fn
-        if self._cell_backward_hook is None:  # pylint: disable=E0203
+        if self._cell_backward_hook is None:
             # Generate a CellBackwardHook prim, and add function for it
             self._cell_backward_hook = inner.CellBackwardHook(self.cls_name + "(" + str(id(self)) + ")",
                                                               self, self._backward_hook)
             self._cell_backward_hook.register_backward_hook()
-        _update_hook_version()
         return handle
 
+    def _backward_hook_construct(self, *inputs, **kwargs):
+        """
+        Backward hook construct method to replace original construct method.
+
+        Args:
+            inputs: The input objects of Cell object.
+            kwargs (dict): Dictionary of variable keyword parameters.
+
+        Returns:
+            - **outputs** - The output objects of Cell object.
+
+        Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+        """
+        # cell_backward_hook has CellBackwardHook op, so keep input args as they are.
+        outputs = self._cell_backward_hook(*inputs)
+        # If the inputs have more than two args, the outputs will also have more than two args and will be wrapped into
+        # a tuple, so need to do unwrapping. If inputs is empty, we also need to unwrap it.
+        # Because when output of runop method is one, it will not wrap a tuple, we need not unwrap it.
+        is_need_unwrap = False
+        if isinstance(outputs, tuple) and len(inputs) != 1:
+            is_need_unwrap = True
+
+        if self._recompute_cell is not None:
+            if is_need_unwrap:
+                outputs = self._recompute_cell(*outputs, **kwargs)
+            else:
+                outputs = self._recompute_cell(outputs, **kwargs)
+        elif self.has_bprop:
+            if is_need_unwrap:
+                outputs = self._call_custom_bprop(*outputs, **kwargs)
+            else:
+                outputs = self._call_custom_bprop(outputs, **kwargs)
+        else:
+            if is_need_unwrap:
+                outputs = self.construct(*outputs, **kwargs)
+            else:
+                outputs = self.construct(outputs, **kwargs)
+        if isinstance(outputs, tuple):
+            new_outputs = self._cell_backward_hook(*outputs)
+        else:
+            new_outputs = self._cell_backward_hook(outputs)
+        # if outputs is (X,) and new_outpus is X
+        if isinstance(outputs, tuple) and len(outputs) == 1:
+            new_outputs = (new_outputs,)
+        return new_outputs
+
     def set_param_ps(self, recurse=True, init_in_server=False):
         """
         Set whether the trainable parameters are updated by parameter server and whether the
@@ -3585,7 +3601,7 @@ class Cell(Cell_):
         """
         Validator.check_bool(mode)
         Validator.check_bool(output_recompute)
-        if not self._has_config_recompute:  # pylint: disable=E0203
+        if not self._has_config_recompute:
             self._has_config_recompute = True
         else:
             logger.info("The recompute interface can be configured only once."
@@ -3628,7 +3644,7 @@ class Cell(Cell_):
                 introduced by optimizer shard are recomputed in auto parallel or semi auto parallel mode.
                 Default: ``False`` .
         """
-        if context._get_mode() == context.PYNATIVE_MODE:
+        if context.get_context("mode") == context.PYNATIVE_MODE:
             self._recompute_cell = recompute_registry.get()(self.construct)
         self._recompute()
         if 'mp_comm_recompute' in kwargs.keys():
@@ -3689,64 +3705,6 @@ class Cell(Cell_):
         """
         self._jit_graph_name = key
 
-    def _jit_backward_pre_hook(self, grad_output):
-        new_grad_output = grad_output
-        if not isinstance(grad_output, tuple):
-            new_grad_output = (grad_output,)
-
-        for fn in self._backward_pre_hook.values():
-            ret = fn(self, new_grad_output)
-            if ret is not None:
-                if not isinstance(ret, tuple):
-                    output = (ret,)
-                else:
-                    output = ret
-            else:
-                output = ops.Depend()(new_grad_output, ret)
-            new_grad_output = output
-
-        if not isinstance(grad_output, tuple):
-            if len(new_grad_output) == 1:
-                return new_grad_output[0]
-            raise TypeError(
-                "The backward pre hook return value size is {} not equal to input size 1".format(
-                    len(new_grad_output)))
-
-        if len(new_grad_output) != len(grad_output):
-            raise TypeError(
-                "The backward pre hook return value size is {} not equal to input size {}".format(
-                    len(new_grad_output), len(grad_output)))
-
-        return new_grad_output
-
-    def _jit_backward_hook(self, grad_input, grad_output):
-        backward_hook_input = grad_input
-        backward_hook_output = grad_output
-        if not isinstance(grad_input, tuple):
-            backward_hook_input = (grad_input,)
-        if not isinstance(grad_output, tuple):
-            backward_hook_output = (grad_output,)
-
-        for fn in self._backward_hook.values():
-            ret = fn(self, backward_hook_input, backward_hook_output)
-            if ret is not None:
-                if not isinstance(ret, tuple):
-                    output = (ret,)
-                else:
-                    output = ret
-            else:
-                output = ops.Depend()(backward_hook_input, ret)
-
-            backward_hook_input = output
-
-        if not isinstance(grad_input, tuple):
-            return backward_hook_input[0]
-
-        if len(backward_hook_input) != len(grad_input):
-            raise TypeError(
-                "The backward hook return value size is {} not equal to input size {}".format(
-                    len(backward_hook_input), len(grad_input)))
-        return backward_hook_input
 
 class GraphCell(Cell):
     """