mindspore 2.3.0__cp310-cp310-win_amd64.whl → 2.4.1__cp310-cp310-win_amd64.whl

mindspore/nn/cell.py

@@ -32,7 +32,7 @@ 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
 from mindspore.common import dtype as mstype
-from mindspore.common.api import _cell_graph_executor, _pynative_executor, _get_args_for_run, cells_compile_cache
+from mindspore.common.api import _cell_graph_executor, _pynative_executor, _get_args_for_run, cells_compile_cache, _no_grad
 from mindspore.common.api import _generate_branch_control_input, _convert_python_data, _get_args_for_run_predict
 from mindspore.common.api import _process_dyn_args, _generate_dyn_compile_args
 from mindspore.common.parameter import Parameter, ParameterTuple
@@ -45,7 +45,6 @@ from mindspore._check_jit_forbidden_api import jit_forbidden_register
 from mindspore.common._decorator import deprecated
 from mindspore.common._register_for_recompute import recompute_registry
 
-
 class Cell(Cell_):
     """
     The basic building block of neural networks in MindSpore. The model or neural network layer should inherit this
@@ -101,9 +100,9 @@ class Cell(Cell_):
     """
 
     IGNORE_LIST = ['_scope', '_cell_init_args', '_auto_prefix', '_cells', '_params', '_create_time',
-                   '_func_graph_flags', '_parameter_layout_dict', '_params_list', '_phase',
-                   '_forward_pre_hook', '_forward_hook', '_enable_forward_pre_hook', '_enable_forward_hook',
-                   '_bprop_debug', '_enable_backward_hook', '_cell_backward_hook', '_is_run', '_param_prefix',
+                   '_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', '_is_run', '_param_prefix',
                    '_attr_synced', 'pynative', 'requires_grad', 'cell_type']
     total_instance_count = 0
 
@@ -135,7 +134,8 @@ class Cell(Cell_):
         self._id = 1
         self.exist_names = set("")
         self.exist_objs = set()
-        self.recompute_cell = None
+        self._recompute_cell = None
+        self.mixed_precision_type = None
         self.sig = inspect.signature(self.construct)
         init_pipeline()
 
@@ -146,13 +146,16 @@ class Cell(Cell_):
         if flags:
             self.add_flags(**flags)
         self._bprop_debug = False
+
+        # hook
         self._forward_pre_hook = OrderedDict()
         self._forward_hook = OrderedDict()
-        self._enable_forward_pre_hook = False
-        self._enable_forward_hook = False
-        self._enable_backward_hook = False
+        self._backward_pre_hook = OrderedDict()
+        self._cell_backward_pre_hook = None
+        self._backward_hook = OrderedDict()
         self._cell_backward_hook = None
         self._is_recursion_hook = False
+
         self.cell_type = None
         self.cast = Cast()
         self._has_config_recompute = False
@@ -166,6 +169,10 @@ class Cell(Cell_):
         self._is_check_and_refresh = False
         self._amp_level = ""
         self._init_flag = False
+        self._shard_fn = None
+        self.has_bprop = False
+        if hasattr(self, "bprop"):
+            self.has_bprop = True
 
     def __getstate__(self):
         base = Cell_.__getstate__(self)
@@ -223,8 +230,9 @@ class Cell(Cell_):
         Get whether cell custom bprop debug is enabled.
 
         Tutorial Examples:
-            - `Cell and Parameter - Custom Cell Reverse
-              <https://mindspore.cn/tutorials/en/master/advanced/modules/layer.html#custom-cell-reverse>`_
+            - `Custom Neural Network Layers - Custom Cell Reverse
+              <https://mindspore.cn/docs/en/master/model_train/custom_program/network_custom.html
+              #custom-cell-reverse>`_
         """
         return self._bprop_debug
 
@@ -374,6 +382,10 @@ class Cell(Cell_):
     def jit_config_dict(self):
         return self._jit_config_dict
 
+    @property
+    def enable_backward_hook(self):
+        return self._enable_backward_hook
+
     def get_func_graph_proto(self):
         """Return graph binary proto."""
         exec_id = ".".join([self.phase, str(self.create_time), str(id(self))])
@@ -401,8 +413,6 @@ class Cell(Cell_):
             cells_compile_cache.pop(id(self), None)
         if hasattr(self, "compile_cache") and self.compile_cache:
             _cell_graph_executor.del_net_res(self, self.compile_cache)
-        if isinstance(self, GraphCell):
-            _cell_graph_executor.dec_graph_cell_count()
         Cell.total_instance_count -= 1
 
     def __delattr__(self, name):
@@ -475,21 +485,28 @@ class Cell(Cell_):
         output = self._run_construct(cast_inputs, kwargs)
         return output
 
-    def _run_construct(self, cast_inputs, kwargs):
+    def _run_construct(self, *inputs, **kwargs):
         """Run the construct function"""
-        if self._enable_forward_pre_hook:
-            cast_inputs = self._run_forward_pre_hook(cast_inputs)
-        if self._enable_backward_hook:
-            output = self._backward_hook_construct(*cast_inputs, **kwargs)
-        elif hasattr(self, "_shard_fn"):
-            output = self._shard_fn(*cast_inputs, **kwargs)
+        if self._forward_pre_hook:
+            inputs = self._run_forward_pre_hook(inputs)
+
+        if self._backward_hook:
+            output = self._backward_hook_construct(*inputs, **kwargs)
+        elif self._shard_fn is not None:
+            output = self._shard_fn(*inputs, **kwargs)
+        elif self._recompute_cell is not None:
+            output = self._recompute_cell(*inputs, **kwargs)
+        elif self.has_bprop and _pynative_executor.requires_grad():
+            output = self._call_custom_bprop(*inputs, **kwargs)
         else:
-            if self.recompute_cell is not None:
-                output = self.recompute_cell(*cast_inputs, **kwargs)
-            else:
-                output = self.construct(*cast_inputs, **kwargs)
-        if self._enable_forward_hook:
-            output = self._run_forward_hook(cast_inputs, output)
+            output = self.construct(*inputs, **kwargs)
+
+        if self._forward_hook:
+            output = self._run_forward_hook(inputs, output)
+
+        if self._backward_pre_hook:
+            output = self._run_backward_pre_hook(output)
+
         return output
 
     def _check_construct_args(self, *args):
@@ -527,7 +544,7 @@ class Cell(Cell_):
         '''Hook function in graph mode'''
         # Check super().__init__() in graph mode.
         try:
-            if self._enable_forward_pre_hook or self._enable_forward_hook or self._enable_backward_hook:
+            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'. "
@@ -579,8 +596,7 @@ class Cell(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
-        this input/output, and None represents data_parallel,
-        which can refer to the description of `mindspore.ops.Primitive.shard`.
+        this input/output, which can refer to the description of `mindspore.ops.Primitive.shard`.
         The parallel strategies of remaining operators are derived from the strategy specified by the input and output.
 
         Note:
@@ -589,8 +605,8 @@ class Cell(Cell_):
             If the input contain Parameter, its strategy should be set in `in_strategy`.
 
         Args:
-            in_strategy (tuple): Define the layout of inputs, each element of the tuple should be a tuple or None. Tuple
-                             defines the layout of the corresponding input and None represents a data parallel strategy.
+            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.
                                                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
@@ -625,7 +641,7 @@ class Cell(Cell_):
             ...   def __init__(self):
             ...     self.block1 = Block()
             ...     self.block2 = Block()
-            ...     self.block2_shard = self.block2.shard(in_strategy=((2, 1),), out_strategy=(None,),
+            ...     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)
@@ -638,7 +654,7 @@ class Cell(Cell_):
 
         shard_fn = Shard()
         fn = shard_fn(self, in_strategy, out_strategy, parameter_plan, device, level)
-        object.__setattr__(self, "_shard_fn", fn)
+        self._shard_fn = fn
         return fn
 
     def auto_cast_inputs(self, inputs):
@@ -666,6 +682,7 @@ class Cell(Cell_):
         for param in self.get_parameters(expand=False):
             if param.has_init:
                 param.init_data()
+        self._init_flag = True
 
     def _self_check(self):
         if not self._is_check_and_refresh:
@@ -684,7 +701,7 @@ class Cell(Cell_):
 
     def __call__(self, *args, **kwargs):
         # Run in Graph mode.
-        if os.getenv("MS_JIT") != '0' and context._get_mode() == context.GRAPH_MODE:
+        if context._get_mode() == context.GRAPH_MODE and os.getenv("MS_JIT") != '0':
             if kwargs:
                 bound_arguments = self.sig.bind(*args, **kwargs)
                 bound_arguments.apply_defaults()
@@ -704,22 +721,69 @@ class Cell(Cell_):
             return out
 
         # Run in PyNative mode.
-        self._self_check()
-        if not self._init_flag:
+        if not (self._init_flag or self._is_check_and_refresh):
             self._init_check()
-            self._init_flag = True
+            self._self_check()
+
+        if not (self.requires_grad or self._dynamic_shape_inputs or self.mixed_precision_type):
+            if not (self._forward_pre_hook or self._forward_hook or self._backward_pre_hook or self._backward_hook or
+                    self._shard_fn or self._recompute_cell or (self.has_bprop and _pynative_executor.requires_grad())):
+                return self.construct(*args, **kwargs)
+
+            return self._run_construct(*args, **kwargs)
+
+        return self._complex_call(*args, **kwargs)
 
+    def _complex_call(self, *args, **kwargs):
+        """
+        PyNative call with requires_grad or hooks
+        """
+        self._call_pre_process(*args, **kwargs)
+
+        if not (self._forward_pre_hook or self._forward_hook or self._backward_pre_hook or self._backward_hook or
+                self._shard_fn or self._recompute_cell or self.has_bprop):
+            output = self.construct(*args, **kwargs)
+        else:
+            output = self._run_construct(*args, **kwargs)
+
+        self._call_post_process(output, *args, **kwargs)
+
+        return output
+
+    def _call_pre_process(self, *args, **kwargs):
+        """
+        Process cell info before call construct
+        """
         if self.requires_grad:
             _pynative_executor.set_grad_flag(True)
-
-        try:
             _pynative_executor.new_graph(self, *args, **kwargs)
-            output = self._run_construct(args, kwargs)
+        elif self._dynamic_shape_inputs is not None:
+            _pynative_executor.set_cell_use_dynamic_shape_process(True)
+
+        # Set mixed precision
+        if self.mixed_precision_type is not None:
+            _pynative_executor.set_mixed_precision_type(self.mixed_precision_type)
+
+    def _call_post_process(self, output, *args, **kwargs):
+        """
+        Process cell info after call construct
+        """
+        if self.requires_grad:
             _pynative_executor.end_graph(self, output, *args, **kwargs)
-        except Exception as err:
-            _pynative_executor.clear_res()
-            raise err
+        elif self._dynamic_shape_inputs is not None:
+            _pynative_executor.set_cell_use_dynamic_shape_process(False)
+
+        # mixed precision reset
+        if self.mixed_precision_type is not None:
+            _pynative_executor.set_mixed_precision_type(MixedPrecisionType.NOTSET, False)
 
+    def _call_custom_bprop(self, *args, **kwargs):
+        """
+        Call custom bprop for cell bprop.
+        """
+        with _no_grad():
+            output = self.construct(*args, **kwargs)
+        _pynative_executor.call_custom_bprop(self, output, *args, **kwargs)
         return output
 
     def _add_attr(self, name, value):
@@ -961,9 +1025,12 @@ class Cell(Cell_):
 
         if not kwargs:
             self._dynamic_shape_inputs = inputs
-            self._check_construct_args(*inputs)
             if context._get_mode() == context.PYNATIVE_MODE:
                 _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)
 
@@ -1682,10 +1749,13 @@ class Cell(Cell_):
     def _add_mixed_precision_flag(self, **flags):
         """Add mixed precision flag to current cell"""
         if "fp16" in flags and flags.get("fp16", False):
+            self.mixed_precision_type = MixedPrecisionType.FP16
             Cell_.set_mixed_precision_type(self, MixedPrecisionType.FP16)
         if "fp32" in flags and flags.get("fp32", False):
+            self.mixed_precision_type = MixedPrecisionType.FP32
             Cell_.set_mixed_precision_type(self, MixedPrecisionType.FP32)
         if "bf16" in flags and flags.get("bf16", False):
+            self.mixed_precision_type = MixedPrecisionType.BF16
             Cell_.set_mixed_precision_type(self, MixedPrecisionType.BF16)
 
     def apply(self, fn):
@@ -1750,9 +1820,6 @@ class Cell(Cell_):
         if not hasattr(self, "_func_graph_flags"):
             self._func_graph_flags = {}
         self._func_graph_flags.update({**flags})
-        if context._get_mode() == context.PYNATIVE_MODE and self._func_graph_flags.get("output_no_recompute"):
-            raise TypeError("Recompute is not supported in PyNative mode currently, you can use "
-                            "'context.set_context(mode=context.GRAPH_MODE)' or @jit to set graph mode.")
         self.__dict__.update({**flags})
         self._add_mixed_precision_flag(**flags)
         return self
@@ -2050,15 +2117,12 @@ 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()
         if not check_hook_fn("register_forward_pre_hook", hook_fn):
             return HookHandle()
-        self._enable_forward_pre_hook = True
-        _pynative_executor.set_hook_changed(self)
-        if not hasattr(self, '_forward_pre_hook_key'):
-            self._forward_pre_hook_key = -1
-        self._forward_pre_hook_key += 1
-        self._forward_pre_hook[self._forward_pre_hook_key] = hook_fn
-        handle = HookHandle(self, self._forward_pre_hook_key, "_forward_pre_hook")
+        handle = HookHandle(self._forward_pre_hook)
+        self._forward_pre_hook[handle.handle_id] = hook_fn
         return handle
 
     def _run_forward_pre_hook(self, inputs):
@@ -2074,14 +2138,23 @@ class Cell(Cell_):
         Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
         """
+        forward_pre_hook_inputs = inputs
         for fn in self._forward_pre_hook.values():
-            ret = fn(self, inputs)
+            ret = fn(self, forward_pre_hook_inputs)
             if ret is not None:
                 if not isinstance(ret, tuple):
-                    inputs = (ret,)
+                    forward_pre_hook_inputs = (ret,)
                 else:
-                    inputs = ret
-        return inputs
+                    forward_pre_hook_inputs = ret
+
+        if isinstance(inputs, tuple):
+            if not isinstance(forward_pre_hook_inputs, tuple):
+                forward_pre_hook_inputs = (forward_pre_hook_inputs,)
+            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
 
     def register_forward_hook(self, hook_fn):
         """
@@ -2142,15 +2215,12 @@ 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()
         if not check_hook_fn("register_forward_hook", hook_fn):
             return HookHandle()
-        self._enable_forward_hook = True
-        _pynative_executor.set_hook_changed(self)
-        if not hasattr(self, '_forward_hook_key'):
-            self._forward_hook_key = -1
-        self._forward_hook_key += 1
-        self._forward_hook[self._forward_hook_key] = hook_fn
-        handle = HookHandle(self, self._forward_hook_key, "_forward_hook")
+        handle = HookHandle(self._forward_hook)
+        self._forward_hook[handle.handle_id] = hook_fn
         return handle
 
     def _run_forward_hook(self, inputs, output):
@@ -2167,11 +2237,110 @@ class Cell(Cell_):
         Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
         """
+        forward_hook_output = output
         for fn in self._forward_hook.values():
-            ret = fn(self, inputs, output)
+            ret = fn(self, inputs, forward_hook_output)
             if ret is not None:
-                output = ret
-        return output
+                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
+
+    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:
+              hook_fn(cell, grad_output) -> New grad_output gradient or None.
+            - The 'hook_fn' is executed in the python environment. In order to prevent running failed when switching to
+              graph mode, it is not recommended to write it in the `construct` function of Cell object.
+            - In the pynative
+              mode, if the `register_backward_pre_hook` function is called in the `construct` function of the Cell
+              object, a hook function will be added at each run time of Cell object.
+
+        Args:
+            hook_fn (function): Python function. Backward pre hook function.
+
+        Returns:
+            A handle corresponding to the `hook_fn` . The handle can be used to remove the added `hook_fn` by calling
+            `handle.remove()` .
+
+        Raises:
+            TypeError: If the `hook_fn` is not a function of python.
+
+        Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+        Examples:
+            >>> import numpy as np
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor, nn, ops
+            >>> ms.set_context(mode=ms.PYNATIVE_MODE)
+            >>> def backward_pre_hook_fn(cell, grad_output):
+            ...     print("backward input: ", grad_output)
+            ...
+            >>> class Net(nn.Cell):
+            ...     def __init__(self):
+            ...         super(Net, self).__init__()
+            ...         self.relu = nn.ReLU()
+            ...         self.handle = self.relu.register_backward_pre_hook(backward_pre_hook_fn)
+            ...
+            ...     def construct(self, x):
+            ...         x = x + x
+            ...         x = self.relu(x)
+            ...         return x
+            >>> grad = ops.GradOperation(get_all=True)
+            >>> net = Net()
+            >>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)))
+            backward input: (Tensor(shape=[1], dtype=Float32, value= [ 1.00000000e+00]),)
+            >>> print(output)
+            (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]),)
+        """
+        if context._get_mode() == context.GRAPH_MODE:
+            return HookHandle()
+        if not check_hook_fn("register_backward_pre_hook", hook_fn):
+            return HookHandle()
+        handle = HookHandle(self._backward_pre_hook)
+        self._backward_pre_hook[handle.handle_id] = hook_fn
+        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()
+        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``
+        """
+        ret = self._cell_backward_pre_hook(outputs)
+        if isinstance(outputs, tuple):
+            if not isinstance(ret, tuple):
+                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 register_backward_hook(self, hook_fn):
         """
@@ -2180,11 +2349,11 @@ class Cell(Cell_):
         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_id` is the information of registered Cell object, including name and ID. `grad_input` is the
-              gradient passed to the Cell. `grad_output` is the gradient computed and passed to the next Cell or
-              primitive, which may be modified by returning a new output gradient.
+              `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
+              passed to the Cell.
             - The 'hook_fn' should have the following signature:
-              hook_fn(cell_id, grad_input, grad_output) -> New output gradient or none.
+              hook_fn(cell, grad_input, grad_output) -> New grad_input gradient or none.
             - The 'hook_fn' is executed in the python environment. In order to prevent running failed when switching to
               graph mode, it is not recommended to write it in the `construct` function of Cell object. In the pynative
               mode, if the `register_backward_hook` function is called in the `construct` function of the Cell object,
@@ -2208,9 +2377,9 @@ class Cell(Cell_):
             >>> import mindspore as ms
             >>> from mindspore import Tensor, nn, ops
             >>> ms.set_context(mode=ms.PYNATIVE_MODE)
-            >>> def backward_hook_fn(cell_id, grad_input, grad_output):
-            ...     print("backward input: ", grad_input)
-            ...     print("backward output: ", grad_output)
+            >>> def backward_hook_fn(cell, grad_input, grad_output):
+            ...     print("backward input: ", grad_output)
+            ...     print("backward output: ", grad_input)
             ...
             >>> class Net(nn.Cell):
             ...     def __init__(self):
@@ -2230,16 +2399,17 @@ class Cell(Cell_):
             >>> print(output)
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]),)
         """
+        if context._get_mode() == context.GRAPH_MODE:
+            return HookHandle()
         if not check_hook_fn("register_backward_hook", hook_fn):
             return HookHandle()
+        handle = HookHandle(self._backward_hook)
+        self._backward_hook[handle.handle_id] = hook_fn
         if self._cell_backward_hook is None:
-            self._enable_backward_hook = True
-            self._cell_backward_hook = inner.CellBackwardHook(self.cls_name + "(" + str(id(self)) + ")")
-            backward_hook_key = self._cell_backward_hook.register_backward_hook(hook_fn)
-            handle = HookHandle(self, backward_hook_key, "_cell_backward_hook")
-        else:
-            backward_hook_key = self._cell_backward_hook.register_backward_hook(hook_fn)
-            handle = HookHandle(self, backward_hook_key, "_cell_backward_hook")
+            # 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()
         return handle
 
     def _backward_hook_construct(self, *inputs, **kwargs):
@@ -2256,21 +2426,31 @@ class Cell(Cell_):
         Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
         """
-        if len(inputs) > 1:
-            inputs = self._cell_backward_hook(inputs)
-        else:
-            inputs = self._cell_backward_hook(*inputs)
-            inputs = (inputs,)
-        if self.recompute_cell is not None:
-            if isinstance(inputs, tuple):
-                outputs = self.recompute_cell(*inputs, **kwargs)
+        # 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.recompute_cell(inputs, **kwargs)
+                outputs = self._call_custom_bprop(outputs, **kwargs)
         else:
-            if isinstance(inputs, tuple):
-                outputs = self.construct(*inputs, **kwargs)
+            if is_need_unwrap:
+                outputs = self.construct(*outputs, **kwargs)
             else:
-                outputs = self.construct(inputs, **kwargs)
+                outputs = self.construct(outputs, **kwargs)
+
         outputs = self._cell_backward_hook(outputs)
         return outputs
 
@@ -2401,7 +2581,8 @@ class Cell(Cell_):
                 Default: ``False`` .
         """
         if context.get_context("mode") == context.PYNATIVE_MODE:
-            self.recompute_cell = recompute_registry.get()(self.construct)
+            self._recompute_cell = recompute_registry.get()(self.construct)
+            self._recompute()
             return
         self._recompute()
         if 'mp_comm_recompute' in kwargs.keys():
@@ -2579,7 +2760,6 @@ class GraphCell(Cell):
         params_dict = update_func_graph_hyper_params(self.graph, params_init)
         for name, param in params_dict.items():
             self._params[name] = param
-        _cell_graph_executor.inc_graph_cell_count()
 
     def construct(self, *inputs):
         return self.graph(*inputs)