mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.4.0__cp39-cp39-win_amd64.whl

mindspore/nn/cell.py

@@ -1,4 +1,4 @@
-# Copyright 2020-2023 Huawei Technologies Co., Ltd
+# Copyright 2020-2024 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,10 +20,9 @@ import inspect
 import os
 import time
 from collections import OrderedDict
-from types import FunctionType, MethodType
 import numpy
 
-from mindspore._checkparam import args_type_check
+from mindspore._checkparam import args_type_check, check_hook_fn
 from mindspore.common._auto_dynamic import is_auto_dynamic, convert_inputs_to_dynamic
 from mindspore import log as logger
 from mindspore.common.parameter import PARAMETER_NAME_DEFAULT
@@ -33,8 +32,9 @@ 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 _generate_branch_control_input
+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
 from mindspore.common.tensor import Tensor
 from mindspore.ops.operations import Cast
@@ -43,16 +43,7 @@ from mindspore.ops.operations import _inner_ops as inner
 from mindspore.parallel.shard import Shard
 from mindspore._check_jit_forbidden_api import jit_forbidden_register
 from mindspore.common._decorator import deprecated
-from mindspore._c_expression import PackExpander
-from mindspore.ops._tracefunc import _convert_tensor, _SetMixedPrecision, PackFunc
-
-
-def _check_args(args):
-    """Check the input args's type"""
-    for item in args:
-        if isinstance(item, Tensor) and item.has_init:
-            item.init_data()
-
+from mindspore.common._register_for_recompute import recompute_registry
 
 class Cell(Cell_):
     """
@@ -89,7 +80,7 @@ class Cell(Cell_):
 
     Examples:
         >>> import mindspore.nn as nn
-        >>> import mindspore.ops as ops
+        >>> from mindspore import ops
         >>> class MyCell(nn.Cell):
         ...     def __init__(self, forward_net):
         ...         super(MyCell, self).__init__(auto_prefix=False)
@@ -109,17 +100,19 @@ class Cell(Cell_):
     """
 
     IGNORE_LIST = ['_scope', '_cell_init_args', '_auto_prefix', '_cells', '_params', '_create_time',
-                   '_func_graph_flags', '_parameter_layout_dict', '_params_list', '_tensor_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
 
     def __init__(self, auto_prefix=True, flags=None):
         Cell_.__init__(self, self._cell_tag)
+        Cell.total_instance_count += 1
+        self.instance_count = Cell.total_instance_count
         self._params = OrderedDict()
         self._cells = OrderedDict()
         self._params_list = OrderedDict()
-        self._tensor_list = OrderedDict()
         self._primitives = OrderedDict()
         self.training = False
         self.requires_grad = False
@@ -135,11 +128,15 @@ class Cell(Cell_):
         self._create_time = int(time.time() * 1e9)
         self.arguments_key = ""
         self.compile_cache = set()
+        self.phase_cache = dict()
         cells_compile_cache[id(self)] = self.compile_cache
         self.parameter_broadcast_done = False
         self._id = 1
         self.exist_names = set("")
         self.exist_objs = set()
+        self._recompute_cell = None
+        self.mixed_precision_type = None
+        self.sig = inspect.signature(self.construct)
         init_pipeline()
 
         # call gc to release GE session resources used by non-used cell objects
@@ -149,24 +146,33 @@ 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
         self._user_parameters = []
         self._dynamic_shape_inputs = None
+        self._compile_args = None
         self.saved_dynamic_shape = None
         self._jit_config_dict = dict()
         self.grad_ops_label = False
         self.ge_sync_data = False
         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)
@@ -224,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/r2.2/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
 
@@ -317,10 +324,23 @@ class Cell(Cell_):
 
     @property
     def pipeline_stage(self):
+        """
+        `pipeline_stage` represents the pipeline stage of current Cell.
+        """
         return self._pipeline_stage
 
     @pipeline_stage.setter
     def pipeline_stage(self, value):
+        """
+        Set the `pipeline_stage` of a Cell.
+
+        Args:
+            value (int): The pipeline stage 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 'Cell', the property 'pipeline_stage' "
                             "must be int type, but got type : {}".format(type(value)))
@@ -362,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))])
@@ -376,10 +400,6 @@ class Cell(Cell_):
             cells = self.__dict__['_cells']
             if name in cells:
                 return cells[name]
-        if '_tensor_list' in self.__dict__:
-            tensor_list = self.__dict__['_tensor_list']
-            if name in tensor_list:
-                return tensor_list[name]
         if '_params_list' in self.__dict__:
             params_list = self.__dict__['_params_list']
             if name in params_list:
@@ -391,12 +411,9 @@ class Cell(Cell_):
             # while deepcopy a cell instance, the copied cell instance can't be added to cells_compile_cache
             # here using pop(id(self), None) to avoid KeyError exception
             cells_compile_cache.pop(id(self), None)
-        try:
-            if self.compile_cache:
-                _cell_graph_executor.del_net_res(self, self.compile_cache)
-        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 hasattr(self, "compile_cache") and self.compile_cache:
+            _cell_graph_executor.del_net_res(self, self.compile_cache)
+        Cell.total_instance_count -= 1
 
     def __delattr__(self, name):
         if name in self._params:
@@ -405,8 +422,6 @@ class Cell(Cell_):
             del self._cells[name]
         elif '_params_list' in self.__dict__ and name in self._params_list:
             del self._params_list[name]
-        elif '_tensor_list' in self.__dict__ and name in self._tensor_list:
-            del self._tensor_list[name]
         else:
             object.__delattr__(self, name)
         self._attr_synced = False
@@ -420,7 +435,7 @@ class Cell(Cell_):
             elif isinstance(item, float):
                 res.append(self.cast(item, dst_type))
             elif hasattr(item, "dtype") and item.dtype in \
-                {mstype.float16, mstype.float32, mstype.float64, mstype.bfloat16} and item.dtype != dst_type:
+                    {mstype.float16, mstype.float32, mstype.float64, mstype.bfloat16} and item.dtype != dst_type:
                 res.append(self.cast(item, dst_type))
             else:
                 res.append(item)
@@ -470,18 +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:
-            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):
@@ -519,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'. "
@@ -566,22 +591,22 @@ class Cell(Cell_):
     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 PyNative mode, use this method
-        to specify a Cell for distributed execution in graph mode.
+        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
-        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:
-            Only effective in PYNATIVE_MODE and in either ParallelMode.AUTO_PARALLEL with
-            search_mode in auto_parallel_context set as sharding_propagation.
+            If Cell.shard is called, the parallel mode in `set_auto_parallel_context` (parallel_mode) will be set to
+            "auto_parallel" and the search mode (search_mode) to "sharding_propagation".
             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
@@ -598,7 +623,7 @@ class Cell(Cell_):
                          use right now. Support [ ``"0"`` , ``"1"`` , ``"2"`` ]. Default: ``0`` .
 
         Returns:
-            Cell, the cell itself.
+            Function, return the cell construct function that will be executed under auto parallel process.
 
         Examples:
             >>> import mindspore.nn as nn
@@ -616,22 +641,21 @@ class Cell(Cell_):
             ...   def __init__(self):
             ...     self.block1 = Block()
             ...     self.block2 = Block()
-            ...     self.block2.shard(in_strategy=((2, 1),), out_strategy=(None,),
-            ...                       parameter_plan={'self.block2.shard.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 context.get_context("mode") != context.PYNATIVE_MODE or \
-                context.get_auto_parallel_context("parallel_mode") not in ["auto_parallel"]:
-            raise AssertionError(f"Cell shard only supports auto parallel under PyNative mode. "
-                                 f"Please check if you call Cell.shard in the script.")
+        if context.get_auto_parallel_context("parallel_mode") not in ["auto_parallel", "semi_auto_parallel"]:
+            raise AssertionError(f"Cell shard only supports auto parallel or semi_auto_parallel "
+                                 f"Please check the parallel mode in parallel context.")
 
         shard_fn = Shard()
         fn = shard_fn(self, in_strategy, out_strategy, parameter_plan, device, level)
-        object.__setattr__(self, "_shard_fn", fn)
-        return self
+        self._shard_fn = fn
+        return fn
 
     def auto_cast_inputs(self, inputs):
         """
@@ -654,65 +678,113 @@ class Cell(Cell_):
 
         return cast_inputs
 
-    def __call__(self, *args, **kwargs):
-        if self.__class__.construct is Cell.construct:
-            raise AttributeError("For 'Cell', the method 'construct' is not defined.")
-
-        if kwargs:
-            bound_arguments = inspect.signature(self.construct).bind(*args, **kwargs)
-            bound_arguments.apply_defaults()
-            args = bound_arguments.args
-            kwargs = bound_arguments.kwargs
+    def _init_check(self):
+        for param in self.get_parameters(expand=False):
+            if param.has_init:
+                param.init_data()
+        self._init_flag = True
 
-        if PackFunc.is_tracing():
-            return self._run_tracefunc(*args, **kwargs)
-
-        if hasattr(self, '_is_check_and_refresh') and not self._is_check_and_refresh:
+    def _self_check(self):
+        if not self._is_check_and_refresh:
             self.check_names_and_refresh_name()
             self._is_check_and_refresh = True
 
+    def _predict(self, *args, **kwargs):
+        if not hasattr(self, "phase"):
+            return False, None
+        if (self.phase == "prefill" or self.phase == 'increment') and self.phase in self.phase_cache:
+            new_args = _get_args_for_run_predict(self, args, kwargs, self._compile_args)
+            res = _cell_graph_executor._graph_executor(tuple(new_args), self.phase_cache[self.phase])
+            res = _convert_python_data(res)
+            return True, res
+        return False, None
+
+    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()
+                args = bound_arguments.args
+                kwargs = bound_arguments.kwargs
+
+            predict_compiled, res = self._predict(*args, **kwargs)
+            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()
             out = self.compile_and_run(*args, **kwargs)
             return out
 
         # Run in PyNative mode.
-        if _pynative_executor.is_first_cell():
-            _pynative_executor._optimizer = getattr(self, "optimizer", None)
-            _pynative_executor._top_cell = self
-            # There many Casts in parameter_broadcast. Enable build faster.
-            self._do_parameter_broadcast()
+        if not (self._init_flag or self._is_check_and_refresh):
+            self._init_check()
+            self._self_check()
 
-        _check_args(args)
-        self._check_cell_flags_in_pynative()
+        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)
 
-        if self.requires_grad and _pynative_executor.enable_grad():
-            _pynative_executor.set_grad_flag(True)
+            return self._run_construct(*args, **kwargs)
 
-        if self._dynamic_shape_inputs is not None:
-            self._check_compile_dynamic_shape(self._dynamic_shape_inputs, args)
+        return self._complex_call(*args, **kwargs)
 
-        try:
+    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)
             _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)
 
-        if isinstance(output, Parameter):
-            output = output.data
-        return output
+        # mixed precision reset
+        if self.mixed_precision_type is not None:
+            _pynative_executor.set_mixed_precision_type(MixedPrecisionType.NOTSET, False)
 
-    def _check_cell_flags_in_pynative(self):
-        """Check the flags added to cell in pynative mode"""
-        if hasattr(self, "_func_graph_flags") 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.")
+    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):
         if name and name[:2] != '__' and name not in Cell.IGNORE_LIST:
@@ -830,15 +902,6 @@ class Cell(Cell_):
         else:
             self.insert_param_to_cell(name, None)
 
-    def _set_attr_for_tensor(self, name, value):
-        if context._get_mode() == context.PYNATIVE_MODE:
-            tensor_list = self.__dict__.get('_tensor_list')
-            if name in self.__dict__:
-                del self.__dict__[name]
-            tensor_list[name] = value
-        else:
-            object.__setattr__(self, name, value)
-
     def __setattr__(self, name, value):
         cells = self.__dict__.get('_cells')
         params = self.__dict__.get('_params')
@@ -856,8 +919,6 @@ class Cell(Cell_):
             if value is not None:
                 raise TypeError(f"For 'Cell', the type of {name} must be cell, but got {type(value).__name__}.")
             self._cells[name] = None
-        elif isinstance(value, Tensor):
-            self._set_attr_for_tensor(name, value)
         else:
             if isinstance(value, Primitive):
                 value.set_prim_instance_name(name)
@@ -910,14 +971,25 @@ class Cell(Cell_):
         """
         logger.warning("'set_parallel_input_with_inputs' function is deprecated.")
 
-    def set_inputs(self, *inputs):
+    def set_inputs(self, *inputs, **kwargs):
         """
         Save set inputs for computation graph. The number of inputs should be the same with that of the datasets. When
         using Model for dynamic shape, please make sure that all networks and loss functions passed to the Model are
-        configured with set_inputs. The inputs can be Tensor of either dynamic or static shape.
+        configured with set_inputs. The shape of input Tensor can be either dynamic or static.
+
+        .. note::
+            There are two mode:
+
+            - Full mode: arguments will be used as all compile inputs for graph-compiling.
+            - Incremental mode: arguments will set to some of the Cell inputs, which will be substituted into the input
+              at the corresponding position for graph-compiling.
+
+            Only one of inputs or kwargs can be set. Inputs for full mode and kwargs for incremental mode.
 
         Args:
-            inputs (tuple): Inputs of the Cell object.
+            inputs (tuple): Full mode arguments.
+            kwargs (dict): Incremental mode arguments. The acceptable key is the name of parameter defined
+                in `self.construct`.
 
         .. warning::
             This is an experimental API that is subject to change or deletion.
@@ -937,16 +1009,30 @@ class Cell(Cell_):
             >>> net = ReluNet()
             >>> input_dyn = Tensor(shape=[3, None], dtype=ms.float32)
             >>> net.set_inputs(input_dyn)
-            >>> input1 = Tensor(np.random.random([3, 10]), dtype=ms.float32)
-            >>> output = net(input1)
+            >>> input = Tensor(np.random.random([3, 10]), dtype=ms.float32)
+            >>> output = net(input)
+            >>>
+            >>> net2 = ReluNet()
+            >>> net2.set_inputs(x=input_dyn)
+            >>> output = net2(input)
         """
         if self.grad_ops_label:
             logger.warning(f'For Cell, set_inputs must be set before the gradient function of the network is '
                            f'generated.')
-        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)
+        if kwargs and inputs:
+            raise ValueError('For Cell, set_inputs should only set inputs or kwargs(inputs: %s, kwargs: %s)!'
+                             % (inputs, kwargs))
+
+        if not kwargs:
+            self._dynamic_shape_inputs = 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)
 
     def get_inputs(self):
         """
@@ -981,6 +1067,48 @@ class Cell(Cell_):
 
         return self._dynamic_shape_inputs
 
+    def _check_parameter_consistency(self, set_inputs, net_inputs):
+        """Check consistency for parameter."""
+        for index, (set_input, net_input) in enumerate(zip(set_inputs, net_inputs)):
+            if isinstance(set_input, Tensor):
+                if not isinstance(net_input, Tensor):
+                    raise TypeError(
+                        f"For 'set_inputs' and tuple(list) in 'set_inputs',the type of {index + 1}th input must "
+                        f"be Tensor, but got {type(net_input)}.")
+                if isinstance(set_input, Parameter) != isinstance(net_input, Parameter):
+                    raise TypeError(
+                        f"For 'set_inputs' and tuple(list) in 'set_inputs', the {index + 1}th input must be the same "
+                        f"as expected, but got expected: {type(set_input)} and input: {type(net_input)}.")
+            elif isinstance(set_input, (tuple, list)):
+                if not isinstance(net_input, (tuple, list)):
+                    raise TypeError(
+                        f"The {index + 1}th input type of 'set_inputs' or tuple(list) in "
+                        f"'set_inputs' must be tuple or list, but got {type(net_input)}.")
+                self._check_parameter_consistency(set_input, net_input)
+
+    def _get_compile_args(self, args):
+        """Get compile arguments."""
+        # this is used only for test
+        set_by_auto_dynamic = False
+        if is_auto_dynamic():
+            if self._dynamic_shape_inputs is None:
+                set_by_auto_dynamic = True
+            else:
+                if isinstance(self._dynamic_shape_inputs, (list, tuple)) and self._dynamic_shape_inputs[0] is None:
+                    set_by_auto_dynamic = True
+        if set_by_auto_dynamic:
+            self._dynamic_shape_inputs = convert_inputs_to_dynamic(*args)
+
+        if self._dynamic_shape_inputs is not None:
+            logger.debug("Compiled Graph with dynamic shape")
+            compile_args = _generate_dyn_compile_args(args, self._dynamic_shape_inputs)
+            _cell_graph_executor._graph_executor.check_argument_consistency(compile_args, args, "set_inputs")
+            self._check_parameter_consistency(compile_args, args)
+            Validator.check_symbolic_shape(compile_args, args)
+            self.saved_dynamic_shape = compile_args
+            return compile_args
+        return args
+
     def compile(self, *args, **kwargs):
         """
         Compile Cell as a computation graph, the input must be consistent with the input defined in construct.
@@ -989,19 +1117,9 @@ class Cell(Cell_):
             args (tuple): Args of the Cell object.
             kwargs (dict): Kwargs of the Cell object.
         """
-        # this is used only for test
-        if is_auto_dynamic() and (self._dynamic_shape_inputs is None or self._dynamic_shape_inputs[0] is None):
-            self._dynamic_shape_inputs = convert_inputs_to_dynamic(*args)
-
-        if self._dynamic_shape_inputs is None:
-            _cell_graph_executor.compile(self, phase=self.phase,
-                                         jit_config_dict=self._jit_config_dict, *args, **kwargs)
-        else:
-            self._check_compile_dynamic_shape(self._dynamic_shape_inputs, args)
-            self.saved_dynamic_shape = self._dynamic_shape_inputs
-            _cell_graph_executor.compile(self, *self._dynamic_shape_inputs, phase=self.phase,
-                                         jit_config_dict=self._jit_config_dict, **kwargs)
-            logger.debug("Compiled Graph with dynamic shape")
+        self._compile_args = self._get_compile_args(args)
+        _cell_graph_executor.compile(self, *self._compile_args, phase=self.phase,
+                                     jit_config_dict=self._jit_config_dict, **kwargs)
 
     def compile_and_run(self, *args, **kwargs):
         """
@@ -1019,7 +1137,7 @@ class Cell(Cell_):
         """
         self.compile(*args, **kwargs)
         self.add_flags(ge_sync_data=False)
-        new_args = _get_args_for_run(self, args, kwargs)
+        new_args = _get_args_for_run(self, args, kwargs, self._compile_args)
         return _cell_graph_executor(self, *new_args, phase=self.phase)
 
     def auto_parallel_compile_and_run(self):
@@ -1033,6 +1151,7 @@ class Cell(Cell_):
 
     def exec_checkpoint_graph(self):
         """Executes GE saving checkpoint graph operation."""
+        logger.warning("'exec_checkpoint_graph' function is deprecated.")
         self.add_flags(ge_sync_data=True)
         _cell_graph_executor(self, phase='save')
 
@@ -1070,14 +1189,14 @@ class Cell(Cell_):
             Parameter(name=bias, shape=(3,), dtype=Int64, requires_grad=True)
         """
         if not param_name:
-            raise KeyError("For 'insert_param_to_cell', the argument 'param_name' should not be None.")
+            raise KeyError(f"For 'insert_param_to_cell', the argument 'param_name' should not be None.")
         if check_name_contain_dot and '.' in param_name:
-            raise KeyError("For 'insert_param_to_cell', the argument 'param_name' should not contain \".\"")
+            raise KeyError(f"For 'insert_param_to_cell', the argument 'param_name' should not contain'.' ")
         if '_params' not in self.__dict__:
-            raise AttributeError("For 'insert_param_to_cell', please call Cell.__init__() firstly.")
+            raise AttributeError(f"For 'insert_param_to_cell', please call Cell.__init__() firstly.")
         if hasattr(self, param_name) and param_name not in self._params:
-            raise KeyError("For 'insert_param_to_cell', the {} parameter already exists in the network. Cannot "
-                           "insert another parameter with the same name.".format(param_name))
+            raise KeyError(f"For 'insert_param_to_cell', the {param_name} parameter already exists in the network."
+                           f"Cannot insert another parameter with the same name.")
         if not isinstance(param, Parameter) and param is not None:
             raise TypeError(f"For 'insert_param_to_cell', the argument 'param' must be 'Parameter' if not None, "
                             f"but got {type(param)}.")
@@ -1139,11 +1258,11 @@ class Cell(Cell_):
             raise TypeError(f"For 'insert_child_to_cell', the type of parameter 'child_name' must be str, "
                             f"but got {type(child_name)}.")
         if not child_name or '.' in child_name:
-            raise KeyError("For 'insert_child_to_cell', the parameter 'child_name' can not be None and "
-                           "can not contain '.'")
+            raise KeyError(f"For 'insert_child_to_cell', the parameter 'child_name' can not be None and "
+                           "can not contain '.' ")
         if hasattr(self, child_name) and child_name not in self._cells:
-            raise KeyError("For 'insert_child_to_cell', the {} child cell already exists in the network. Cannot "
-                           "insert another child cell with the same name.".format(child_name))
+            raise KeyError(f"For 'insert_child_to_cell', the {child_name} child cell already exists in the network."
+                           f"Cannot insert another child cell with the same name.")
         if not isinstance(child_cell, Cell) and child_cell is not None:
             raise TypeError(f"For 'insert_child_to_cell', the argument 'child_cell' must be 'Cell' if not None, "
                             f"but got type {type(child_cell)}.")
@@ -1163,7 +1282,7 @@ class Cell(Cell_):
         Returns:
             Tensor, returns the computed result.
         """
-        return None
+        raise AttributeError("For 'Cell', the method 'construct' is not defined.")
 
     def remove_redundant_parameters(self):
         """
@@ -1361,7 +1480,7 @@ class Cell(Cell_):
 
         Tutorial Examples:
             - `Model Training - Optimizer
-              <https://mindspore.cn/tutorials/en/r2.2/beginner/train.html#optimizer>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/train.html#optimizer>`_
         """
         return list(filter(lambda x: x.requires_grad, self.get_parameters(expand=recurse)))
 
@@ -1472,7 +1591,7 @@ class Cell(Cell_):
 
         Tutorial Examples:
             - `Building a Network - Model Parameters
-              <https://mindspore.cn/tutorials/en/r2.2/beginner/model.html#model-parameters>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/model.html#model-parameters>`_
         """
         cells = []
         if expand:
@@ -1630,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):
@@ -1698,6 +1820,9 @@ 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
@@ -1808,7 +1933,7 @@ class Cell(Cell_):
         accelerate the algorithm in the algorithm library.
 
         If `boost_type` is not in the algorithm library, please view the algorithm in the algorithm library through
-        `algorithm library <https://gitee.com/mindspore/mindspore/tree/r2.2/mindspore/python/mindspore/boost>`_.
+        `algorithm library <https://gitee.com/mindspore/mindspore/tree/master/mindspore/python/mindspore/boost>`_.
 
         Note:
             Some acceleration algorithms may affect the accuracy of the network, please choose carefully.
@@ -1865,7 +1990,7 @@ class Cell(Cell_):
 
         Tutorial Examples:
             - `Model Training - Implementing Training and Evaluation
-              <https://mindspore.cn/tutorials/en/r2.2/beginner/train.html#training-and-evaluation>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/train.html#training-and-evaluation>`_
         """
         if mode:
             self._phase = 'train'
@@ -1945,11 +2070,11 @@ class Cell(Cell_):
         Note:
             - The `register_forward_pre_hook(hook_fn)` does not work in graph mode or functions decorated with 'jit'.
             - 'hook_fn' must be defined as the following code.
-              `cell_id` is the information of registered Cell object, including name and ID. `inputs` is the forward
+              `cell` is the object of registered Cell. `inputs` is the forward
               input objects passed to the Cell. The 'hook_fn' can modify the forward input objects by returning new
               forward input objects.
             - It should have the following signature:
-              hook_fn(cell_id, inputs) -> new input objects or none.
+              hook_fn(cell, inputs) -> new input objects or none.
             - 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_forward_pre_hook` function is
               called in the `construct` function of the Cell object, a hook function will be added at each run time of
@@ -1959,8 +2084,8 @@ class Cell(Cell_):
             hook_fn (function): Python function. Forward pre hook function.
 
         Returns:
-            Handle, it is an instance of `mindspore.common.hook_handle.HookHandle` and corresponding to the `hook_fn` .
-            The handle can be used to remove the added `hook_fn` by calling `handle.remove()` .
+            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.
@@ -1973,7 +2098,7 @@ class Cell(Cell_):
             >>> import mindspore as ms
             >>> from mindspore import Tensor, nn, ops
             >>> ms.set_context(mode=ms.PYNATIVE_MODE)
-            >>> def forward_pre_hook_fn(cell_id, inputs):
+            >>> def forward_pre_hook_fn(cell, inputs):
             ...     print("forward inputs: ", inputs)
             ...
             >>> class Net(nn.Cell):
@@ -1995,24 +2120,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_context("mode") != context.PYNATIVE_MODE:
-            logger.warning(f"'register_forward_pre_hook' function is only supported in pynative mode, you can use "
-                           f"context.set_context to set pynative mode.")
+        if context._get_mode() == context.GRAPH_MODE:
             return HookHandle()
-
-        if not isinstance(hook_fn, (FunctionType, MethodType)):
-            raise TypeError(f"When using 'register_forward_pre_hook(hook_fn)', the type of 'hook_fn' must be python "
-                            f"function, but got {type(hook_fn)}.")
-        if hook_fn.__code__.co_name == "staging_specialize":
-            raise TypeError(f"Decorating hook function {hook_fn.__name__} with '@jit' is not supported.")
-
-        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")
+        if not check_hook_fn("register_forward_pre_hook", hook_fn):
+            return HookHandle()
+        handle = HookHandle(self._forward_pre_hook)
+        self._forward_pre_hook[handle.handle_id] = hook_fn
         return handle
 
     def _run_forward_pre_hook(self, inputs):
@@ -2028,15 +2141,23 @@ class Cell(Cell_):
         Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
         """
-        cell_id = self.cls_name + "(" + str(id(self)) + ")"
+        forward_pre_hook_inputs = inputs
         for fn in self._forward_pre_hook.values():
-            ret = fn(cell_id, 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):
         """
@@ -2045,11 +2166,11 @@ class Cell(Cell_):
         Note:
             - The `register_forward_hook(hook_fn)` does not work in graph mode or functions decorated with 'jit'.
             - 'hook_fn' must be defined as the following code.
-              `cell_id` is the information of registered Cell object, including name and ID. `inputs` is the forward
+              `cell` is the object of registered Cell. `inputs` is the forward
               input objects passed to the Cell. `output` is the forward output object of the Cell. The 'hook_fn' can
               modify the forward output object by returning new forward output object.
             - It should have the following signature:
-              hook_fn(cell_id, inputs, output) -> new output object or none.
+              hook_fn(cell, inputs, output) -> new output object or none.
             - 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_forward_hook` function is
               called in the `construct` function of the Cell object, a hook function will be added at each run time of
@@ -2059,8 +2180,8 @@ class Cell(Cell_):
             hook_fn (function): Python function. Forward hook function.
 
         Returns:
-            Handle, it is an instance of `mindspore.common.hook_handle.HookHandle` and corresponding to the `hook_fn` .
-            The handle can be used to remove the added `hook_fn` by calling `handle.remove()` .
+            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.
@@ -2073,7 +2194,7 @@ class Cell(Cell_):
             >>> import mindspore as ms
             >>> from mindspore import Tensor, nn, ops
             >>> ms.set_context(mode=ms.PYNATIVE_MODE)
-            >>> def forward_hook_fn(cell_id, inputs, output):
+            >>> def forward_hook_fn(cell, inputs, output):
             ...     print("forward inputs: ", inputs)
             ...     print("forward output: ", output)
             ...
@@ -2097,24 +2218,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_context("mode") != context.PYNATIVE_MODE:
-            logger.warning(f"'register_forward_hook' function is only supported in pynative mode, you can use "
-                           f"context.set_context to set pynative mode.")
+        if context._get_mode() == context.GRAPH_MODE:
             return HookHandle()
-
-        if not isinstance(hook_fn, (FunctionType, MethodType)):
-            raise TypeError(f"When using 'register_forward_hook(hook_fn)', the type of 'hook_fn' must be python "
-                            f"function, but got {type(hook_fn)}.")
-        if hook_fn.__code__.co_name == "staging_specialize":
-            raise TypeError(f"Decorating hook function {hook_fn.__name__} with '@jit' is not supported.")
-
-        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")
+        if not check_hook_fn("register_forward_hook", hook_fn):
+            return HookHandle()
+        handle = HookHandle(self._forward_hook)
+        self._forward_hook[handle.handle_id] = hook_fn
         return handle
 
     def _run_forward_hook(self, inputs, output):
@@ -2131,12 +2240,110 @@ class Cell(Cell_):
         Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
         """
-        cell_id = self.cls_name + "(" + str(id(self)) + ")"
+        forward_hook_output = output
         for fn in self._forward_hook.values():
-            ret = fn(cell_id, 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):
         """
@@ -2145,11 +2352,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,
@@ -2159,8 +2366,8 @@ class Cell(Cell_):
             hook_fn (function): Python function. Backward hook function.
 
         Returns:
-            Handle, it is an instance of `mindspore.common.hook_handle.HookHandle` and corresponding to the `hook_fn` .
-            The handle can be used to remove the added `hook_fn` by calling `handle.remove()` .
+            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.
@@ -2173,9 +2380,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):
@@ -2195,22 +2402,17 @@ class Cell(Cell_):
             >>> print(output)
             (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]),)
         """
-        if context.get_context("mode") != context.PYNATIVE_MODE:
-            logger.warning(f"'register_backward_hook' function is only supported in pynative mode, you can use "
-                           f"context.set_context to set pynative mode.")
+        if context._get_mode() == context.GRAPH_MODE:
             return HookHandle()
-
-        if not isinstance(hook_fn, (FunctionType, MethodType)):
-            raise TypeError(f"When using 'register_backward_hook(hook_fn)', the type of 'hook_fn' must be python "
-                            f"function, but got {type(hook_fn)}.")
+        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):
@@ -2227,15 +2429,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 isinstance(inputs, tuple):
-            outputs = self.construct(*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._call_custom_bprop(outputs, **kwargs)
         else:
-            outputs = self.construct(inputs, **kwargs)
+            if is_need_unwrap:
+                outputs = self.construct(*outputs, **kwargs)
+            else:
+                outputs = self.construct(outputs, **kwargs)
+
         outputs = self._cell_backward_hook(outputs)
         return outputs
 
@@ -2365,6 +2583,9 @@ class Cell(Cell_):
                 introduced by optimizer shard are recomputed in auto parallel or semi auto parallel mode.
                 Default: ``False`` .
         """
+        if context.get_context("mode") == context.PYNATIVE_MODE:
+            self._recompute_cell = recompute_registry.get()(self.construct)
+            return
         self._recompute()
         if 'mp_comm_recompute' in kwargs.keys():
             self._mp_comm_recompute(kwargs.get('mp_comm_recompute', False))
@@ -2384,16 +2605,13 @@ class Cell(Cell_):
                                  "the key kwargs must be 'mp_comm_recompute', "
                                  "'parallel_optimizer_comm_recompute', 'recompute_slice_activation'" % key)
 
+    @deprecated("2.3", "infer_param_pipeline_stage")
     def infer_param_pipeline_stage(self):
         """
         Infer pipeline stages of all parameters in the cell.
 
         Note:
-            - If a parameter does not belong to any cell which has been set pipeline_stage,
-              the parameter should use add_pipeline_stage to add it's pipeline_stage information.
-            - If a parameter P has been used by two operators in different stages "stageA" and "stageB",
-              the parameter P should use P.add_pipeline_stage(stageA) and P.add_pipeline_stage(stageB)
-              to add it's stage information before using infer_param_pipeline_stage.
+            - The interface is deprecated from version 2.3 and will be removed in a future version.
 
         Returns:
             The params belong to current stage in pipeline parallel.
@@ -2448,86 +2666,6 @@ class Cell(Cell_):
         for op in all_ops:
             op.place(role, rank_id)
 
-    def _check_dynamic_tensor(self, set_input, net_input, index):
-        """
-        Check if tensor is correctly set for dynamic shape.
-
-        Args:
-            set_input (Tensor): Tensor set for dynamic shape.
-            net_input (Tensor): Input tensor of the Cell object.
-            index (int): Tensor index for set inputs.
-        """
-        if not isinstance(net_input, Tensor):
-            raise TypeError(
-                f"For 'set_inputs' and tuple(list) in 'set_inputs',the type of {index + 1}th input must be Tensor, "
-                f"but got {type(net_input)}.")
-        is_param_set_input = isinstance(set_input, Parameter)
-        is_param_net_input = isinstance(net_input, Parameter)
-        if (is_param_set_input and not is_param_net_input) or (is_param_net_input and not is_param_set_input):
-            raise TypeError(
-                f"For 'set_inputs' and tuple(list) in 'set_inputs', the {index + 1}th input must be the same "
-                f"as network's input, but got 'set_inputs': {type(set_input)} and network's input: {type(net_input)}.")
-        if set_input.dtype != net_input.dtype:
-            raise TypeError(
-                f"For 'set_inputs' and tuple(list) in 'set_inputs',the dtype of {index + 1}th input must be the same "
-                f"as network's input, but got 'set_inputs': {set_input.dtype} and network's input: {net_input.dtype}.")
-        if -2 not in set_input.shape:
-            if net_input.dim() != 0 and set_input.dim() != net_input.dim():
-                raise ValueError(
-                    f"For 'set_inputs' and tuple(list) in 'set_inputs',the dims of {index + 1}th input must be the "
-                    f"same as network's input, but got 'set_inputs': {set_input.dim()} and network's input: "
-                    f"{net_input.dim()}.")
-            if not all([ele1 in (-1, ele2) for ele1, ele2 in zip(set_input.shape, net_input.shape)]):
-                raise ValueError(
-                    f"For 'set_inputs' and tuple(list) in 'set_inputs',the shape of {index + 1}th input must be the "
-                    f"same as network's input, but got 'set_inputs': {set_input.shape} and network's input: "
-                    f"{net_input.shape}.")
-
-    def _check_compile_dynamic_shape(self, set_inputs, net_inputs):
-        """
-        Check if graph has been compiled with dynamic shape.
-
-        Args:
-            net_inputs (tuple): Inputs of the Cell object.
-        """
-        set_inputs_len = len(set_inputs)
-        net_inputs_len = len(net_inputs)
-        if set_inputs_len != net_inputs_len:
-            raise ValueError("The length of 'set_inputs' or tuple(list) in 'set_inputs' must be equal to network's "
-                             f"inputs, but got 'set_inputs': {set_inputs_len} and network's input: {net_inputs_len}.")
-        for index, (set_input, net_input) in enumerate(zip(set_inputs, net_inputs)):
-            if isinstance(set_input, Tensor):
-                self._check_dynamic_tensor(set_input, net_input, index)
-            elif isinstance(set_input, (tuple, list)):
-                if not isinstance(net_input, (tuple, list)):
-                    raise TypeError(
-                        f"The {index + 1}th input type of 'set_inputs' or tuple(list) in 'set_inputs' must be tuple or "
-                        f"list, but got {type(net_input)}.")
-                self._check_compile_dynamic_shape(set_input, net_input)
-            else:
-                if context._get_mode() == context.PYNATIVE_MODE and set_input is None:
-                    continue
-                if net_input != set_input:
-                    raise ValueError(
-                        f"The {index + 1}th input of 'set_inputs' or tuple(list) in 'set_inputs' must be the same with "
-                        f"network's input, but got set_inputs: {set_input} and network's input: {net_input}.")
-
-    def _run_tracefunc(self, *args, **kwargs):
-        """ Run Packed Cell in Pack."""
-        args = self._mixed_precision_cast(args)
-        need_subgraph = hasattr(self, "bprop") or hasattr(self, "_pipeline_stage") or self.get_flags()
-        if not PackFunc.current.is_pynative_mode and need_subgraph:
-            expander = PackExpander.get_instance()
-            args = expander.begin_subgraph(self, *args)
-            args = [_convert_tensor(a) for a in args]
-            output = self._run_construct(args, kwargs)
-            ret = expander.end_subgraph(self, output)
-            output = _convert_tensor(ret)
-        else:
-            with _SetMixedPrecision(self):
-                output = self._run_construct(args, kwargs)
-        return output
-
     def _mixed_precision_cast(self, inputs):
         mixed_type = self.get_mixed_precision_type()
         if mixed_type == MixedPrecisionType.NOTSET:
@@ -2633,7 +2771,7 @@ class GraphCell(Cell):
         self._add_attr("graph_load_from_mindir", self.graph)
         if not self.obf_random_seed:
             return self.compile_and_run(*args, **kwargs)
-        append_input = Tensor((numpy.ones((1, 1)) * self._branch_control_input).astype(numpy.int32))
+        append_input = Tensor((numpy.ones((1,)) * self._branch_control_input).astype(numpy.int32))
         return self.compile_and_run(*args, append_input, **kwargs)