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

mindspore/train/model.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.
@@ -18,6 +18,7 @@ from __future__ import absolute_import
 from collections.abc import Iterable
 from functools import wraps
 
+import sys
 import os
 import math
 import copy
@@ -34,8 +35,10 @@ from mindspore.common.tensor import Tensor
 from mindspore.train.metrics import get_metrics, get_metric_fn
 from mindspore._checkparam import check_input_data, check_output_data
 from mindspore import _checkparam as Validator
-from mindspore.train.callback import _InternalCallbackParam, RunContext, _CallbackManager, Callback, TimeMonitor
+from mindspore.train.callback import _InternalCallbackParam, RunContext, _CallbackManager, Callback, TimeMonitor,\
+    FlopsUtilizationCollector, TFTRegister
 from mindspore.train.callback import __all__ as internal_cb_names
+from mindspore.train.callback._cluster_monitor import ClusterMonitor
 from mindspore import context
 from mindspore.parallel._utils import _get_parallel_mode, _get_device_num, _get_parameter_broadcast, \
     _device_number_check, _parameter_broadcast_check, _parallel_predict_check, \
@@ -48,7 +51,7 @@ from mindspore.boost import AutoBoost
 from mindspore.context import ParallelMode
 from mindspore.parallel._recovery_context import _set_recovery_context, _get_recovery_context
 from mindspore.train.dataset_helper import DatasetHelper, connect_network_with_dataset
-from mindspore.common.api import _pynative_executor
+from mindspore.common.api import _pynative_executor, ARG_SPECIFIED, TOTAL_ARG_LEN
 from mindspore.dataset.core.config import get_debug_mode
 from mindspore.dataset.engine.datasets import _set_training_dataset, _reset_training_dataset
 from mindspore.train import amp
@@ -75,6 +78,7 @@ class _FrameworkProfilerCallback(Callback):
     """
     Profiler callback of framework for training.
     """
+
     def step_begin(self, run_context):
         _framework_profiler_step_start()
 
@@ -115,6 +119,236 @@ def _save_final_ckpt(func):
             func(self, *args, **kwargs)
     return wrapper
 
+def _handle_tft(func):
+    """
+    Decorator function, which starts uce handle process when an exception occurs during training.
+    """
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        obj = None
+        if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), TFTRegister):
+            obj = kwargs.get('callbacks')
+        if kwargs.get('callbacks') and isinstance(kwargs.get('callbacks'), list):
+            for item in kwargs.get('callbacks'):
+                if isinstance(item, TFTRegister):
+                    obj = item
+        if obj:
+            tft = obj.tft
+            tft_env = os.getenv("MS_ENABLE_TFT", "")
+            uce_env = "UCE:1" in tft_env
+            while True:
+                try:
+                    return func(self, *args, **kwargs)
+                except RuntimeError as e:
+                    logger.info("uce wrapper caught RuntimeError")
+                    if not uce_env:
+                        logger.info("uce wrapper caught RuntimeError uce not enable")
+                        tft.tft_report_error(tft.ReportState.RS_UNKNOWN.value)
+                        raise e
+                    e_str = str(e)
+                    logger.info("uce wrapper caught RuntimeError e_str:{}".format(e_str))
+                    if "UCEError" in e_str:
+                        logger.info("uce wrapper report UCEError")
+                        tft.tft_report_error(tft.ReportState.RS_UCE.value)
+                    elif "ForceStopError" in e_str:
+                        logger.info("uce wrapper caught RuntimeError ForceStopError")
+                        force_stop_err = tft.ReportState.RS_NORMAL.value
+                        tft.tft_report_error(force_stop_err)
+                    else:
+                        logger.info("uce wrapper caught RuntimeError rankid: {} OTHER ERROR")
+                        tft.tft_report_error(tft.ReportState.RS_UNKNOWN.value)
+                        raise e
+                    ret = tft.tft_wait_next_action()
+                    if ret == tft.Action.EXIT.value:
+                        raise e
+                    repair_step = tft.tft_get_repair_step()
+                    logger.info("uce wrapper caught repair finish REPAIR STEP: {} batch_num: \
+{}".format(repair_step, self.batch_num))
+                    initial_epoch = int(repair_step/self.batch_num)
+                    initial_step = repair_step % self.batch_num
+                    kwargs["initial_epoch"] = initial_epoch
+
+                    train_dataset = args[1]
+                    dataset_sink_mode = args[3] if len(args) > 3 else kwargs.get('dataset_sink_mode', True)
+                    sink_size = args[4] if len(args) > 4 else kwargs.get('sink_size', -1)
+
+                    cb_initial_step = 0
+                    if dataset_sink_mode:
+                        train_dataset.set_init_step(initial_epoch)
+                        dataset_size = train_dataset.get_dataset_size()
+                        if sink_size != -1:
+                            cb_initial_step = initial_epoch * sink_size + initial_step
+                        else:
+                            cb_initial_step = initial_epoch * dataset_size + initial_step
+                    else:
+                        train_dataset.set_init_step(initial_step)
+                        cb_initial_step = initial_step
+
+                    kwargs["initial_step"] = cb_initial_step
+
+                    logger.info("uce wrapper repair complete  \
+initial_epoch: {}, cb_initial_step: {} ".format(initial_epoch, cb_initial_step))
+                    continue
+                except BaseException as e:
+                    logger.info("uce wrapper caught BaseException error")
+                    tft.tft_report_error(tft.ReportState.RS_UNKNOWN.value)
+                    raise e
+        else:
+            return func(self, *args, **kwargs)
+    return wrapper
+
+
+def _check_tft():
+    """Check if TFT is supported"""
+    tft_env = os.getenv("MS_ENABLE_TFT")
+    device_target = context.get_context("device_target")
+    if tft_env and device_target == "Ascend":
+        from mindspore._c_expression import MSContext
+        ascend_target = MSContext.get_instance().get_ascend_soc_version()
+        if ascend_target == 'ascend910':
+            raise ValueError("TFT is not supported when using ascend910")
+        ms_mode = context.get_context("mode")
+        if ms_mode != mindspore.GRAPH_MODE:
+            raise ValueError("TFT is only supported in GRAPH_MODE")
+        jit_level = context.get_context("jit_level")
+        if jit_level == "O2" and "UCE:1" in tft_env:
+            raise ValueError("TFT is not supported when using jit_level == O2")
+
+
+def _append_ccae(callbacks):
+    """Add cluster monitoring when CCAE is enabled."""
+    perf_config = os.getenv("PERF_DUMP_CONFIG")
+    if perf_config is None:
+        return callbacks
+    pairs = perf_config.split(',')
+    perf_config_dict = {}
+    for pair in pairs:
+        key, value = pair.split(':')
+        if value.lower() == 'true':
+            perf_config_dict[key] = True
+        elif value.lower() == 'false':
+            perf_config_dict[key] = False
+        elif value.isdigit():
+            perf_config_dict[key] = int(value)
+        else:
+            perf_config_dict[key] = value
+    if perf_config_dict.get("enable", False):
+        if callbacks is None:
+            callbacks = ClusterMonitor()
+        elif isinstance(callbacks, list):
+            callbacks.append(ClusterMonitor())
+        else:
+            callbacks = [callbacks, ClusterMonitor()]
+    return callbacks
+
+
+def _get_arg_infos(inputs):
+    """Get compile argument information from inputs.
+
+    Args:
+        inputs (Union[list, tuple, dict]): Argument got from cell which is set by `set_inputs`.
+
+    Raises:
+        RuntimeError: inputs is not a list, tuple or dict.
+        RuntimeError: inputs is a dict without necessary keys and values.
+
+    Returns:
+        _type_: _description_
+    """
+    if isinstance(inputs, (list, tuple)):
+        arg_specified = [[idx, arg] for idx, arg in enumerate(inputs)]
+        arg_len = len(inputs)
+    elif isinstance(inputs, dict):
+        arg_specified = inputs.get(ARG_SPECIFIED, None)
+        arg_len = inputs.get(TOTAL_ARG_LEN, None)
+        if arg_specified is None or arg_len is None:
+            raise RuntimeError(
+                "The incremental inputs should be processed(with \"%s\" and \"%s\"), but got %s." %
+                (ARG_SPECIFIED, TOTAL_ARG_LEN, str(inputs)))
+    else:
+        raise RuntimeError("inputs should be a list/tuple or a dict, but got %s!" % str(inputs))
+
+    return arg_len, arg_specified
+
+
+def _merge_inputs(inputs1, inputs2):
+    """Merge two processed inputs to a new inputs for latter setting cell's inputs."""
+    is_fullmode1 = isinstance(inputs1, (list, tuple))
+    is_fullmode2 = isinstance(inputs2, (list, tuple))
+
+    if is_fullmode1 and is_fullmode2:
+        return [*inputs1, *inputs2]
+
+    arg_len1, arg_specified1 = _get_arg_infos(inputs1)
+    arg_len2, arg_specified2 = _get_arg_infos(inputs2)
+
+    res_arg_len = arg_len1 + arg_len2
+    res_arg_specified = []
+    res_arg_specified.extend(arg_specified1)
+    # The second inputs should add offset before merging.
+    for idx, arg in arg_specified2:
+        res_arg_specified.append([idx + arg_len1, arg])
+
+    return {ARG_SPECIFIED: res_arg_specified, TOTAL_ARG_LEN: res_arg_len}
+
+
+def _process_loss_inputs(loss_inputs):
+    """Process loss's inputs whose first input should be dropped for train or eval.
+
+    Args:
+        loss_inputs (Union[list, tuple, dict]): Arguments save by `set_inputs` or `jit`.
+
+    Raises:
+        RuntimeError: inputs is not a list, tuple or dict.
+        RuntimeError: inputs is a dict without necessary keys and values.
+
+    Returns:
+        list, tuple or dict: Arguments for latter setting.
+    """
+    # For train or eval, the first input of loss is the inner-tensor, so drop it.
+    res = None
+    if isinstance(loss_inputs, (list, tuple)):
+        res = [*loss_inputs]
+        res.pop(0)
+    elif isinstance(loss_inputs, dict):
+        loss_arg_specified = loss_inputs.get(ARG_SPECIFIED, None)
+        loss_arg_len = loss_inputs.get(TOTAL_ARG_LEN, None)
+        if loss_arg_specified is None or loss_arg_len is None:
+            raise RuntimeError(
+                "The loss incremental inputs should be processed(with \"%s\" and \"%s\"), but got %s." %
+                (ARG_SPECIFIED, TOTAL_ARG_LEN, str(loss_inputs)))
+        res_loss_arg_specified = []
+        for idx, arg in loss_arg_specified:
+            if idx == 0:
+                continue
+            res_loss_arg_specified.append([idx, arg])
+        res = {ARG_SPECIFIED: res_loss_arg_specified, TOTAL_ARG_LEN: loss_arg_len - 1}
+    else:
+        raise RuntimeError("loss_inputs should be a list/tuple or a dict, but got %s!" % str(loss_inputs))
+
+    return res
+
+
+def _set_with_processed_inputs(network, inputs):
+    """Save set inputs for computation graph with processed inputs.
+
+    Args:
+        network (nn.Cell): Target cell.
+        inputs (Union[list, tuple, dict]): Inputs argument got from other cell.
+
+    Raises:
+        RuntimeError: network is not a nn.Cell.
+        RuntimeError: inputs is not a list, tuple or dict.
+    """
+    Validator.check_value_type('network', network, nn.Cell)
+    if isinstance(inputs, (list, tuple)):
+        network.set_inputs(*inputs)
+    elif isinstance(inputs, dict):
+        network.set_inputs(**inputs)
+    else:
+        raise RuntimeError(
+            "Reset inputs from a process inputs, should be a list/tuple or a dict, but got %s!" % str(inputs))
+
 
 class Model:
     """
@@ -133,8 +367,8 @@ class Model:
 
     Args:
         network (Cell): A training or testing network.
-        loss_fn (Cell): Objective function. If `loss_fn` is None, the `network` should contain the calculation of loss
-                        and parallel if needed. Default: ``None`` .
+        loss_fn (Cell): Objective function. If `loss_fn` is None, the `network` should contain the calculation of loss.
+                        Default: ``None`` .
         optimizer (Cell): Optimizer for updating the weights. If `optimizer` is None, the `network` needs to
                           do backpropagation and update weights. Default: ``None`` .
         metrics (Union[dict, set]): A Dictionary or a set of metrics for model evaluation.
@@ -151,21 +385,11 @@ class Model:
         amp_level (str): Option for argument `level` in :func:`mindspore.amp.build_train_network`, level for mixed
             precision training. Supports ["O0", "O1", "O2", "O3", "auto"]. Default: ``"O0"`` .
 
-            - "O0": Do not change.
-            - "O1": Cast the operators in white_list to float16, the remaining operators are kept in float32.
-              The operators in the whitelist: [Conv1d, Conv2d, Conv3d, Conv1dTranspose, Conv2dTranspose,
-              Conv3dTranspose, Dense, LSTMCell, RNNCell, GRUCell, MatMul, BatchMatMul, PReLU, ReLU, Ger].
-            - "O2": Cast network to float16, keep BatchNorm run in float32, using dynamic loss scale.
-            - "O3": Cast network to float16, the BatchNorm is also cast to float16, loss scale will not be used.
-            - "auto": Set level to recommended level in different devices. Set level to "O2" on GPU, set
-              level to "O3" on Ascend. The recommended level is chosen by the expert experience, not applicable to all
-              scenarios. User should specify the level for special network.
-
-            "O2" is recommended on GPU, "O3" is recommended on Ascend.
+            For details on `amp_level` , refer to :func:`mindspore.amp.auto_mixed_precision`.
+
             The BatchNorm strategy can be changed by `keep_batchnorm_fp32` settings in `kwargs`. `keep_batchnorm_fp32`
             must be a bool. The loss scale strategy can be changed by `loss_scale_manager` setting in `kwargs`.
             `loss_scale_manager` should be a subclass of :class:`mindspore.amp.LossScaleManager`.
-            The more detailed explanation of `amp_level` setting can be found at `mindspore.amp.build_train_network`.
 
         boost_level (str): Option for argument `level` in `mindspore.boost`, level for boost mode
             training. Supports ["O0", "O1", "O2"]. Default: ``"O0"`` .
@@ -190,7 +414,7 @@ class Model:
         >>> from mindspore.train import Model
         >>>
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
         >>> optim = nn.Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9)
@@ -199,7 +423,7 @@ class Model:
         >>> model.predict_network
         >>> model.eval_network
         >>> # Create the dataset taking MNIST as an example. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
         >>> dataset = create_dataset()
         >>> model.train(2, dataset)
     """
@@ -240,7 +464,7 @@ class Model:
         self._mindspore_lite = None
         self._lite_infer = True  # if backend lite infer fails, set False
         self._mindspore_lite_model_group_id = id(self) & 0xFFFF
-
+        self.batch_num = -1
 
     def _check_for_graph_cell(self, kwargs):
         """Check for graph cell"""
@@ -312,13 +536,10 @@ class Model:
             raise ValueError("The argument 'optimizer' can not be None when set 'loss_scale_manager'.")
 
         net_inputs = network.get_inputs()
-        loss_inputs = [None]
         if self._loss_fn:
-            if self._loss_fn.get_inputs():
-                loss_inputs = [*self._loss_fn.get_inputs()]
-            loss_inputs.pop(0)
-            if net_inputs:
-                net_inputs = [*net_inputs, *loss_inputs]
+            if self._loss_fn.get_inputs() and net_inputs:
+                loss_inputs = _process_loss_inputs(self._loss_fn.get_inputs())
+                net_inputs = _merge_inputs(net_inputs, loss_inputs)
         if self._optimizer:
             amp_config = {}
             if self._loss_scale_manager_set:
@@ -336,7 +557,7 @@ class Model:
         # If need to check if loss_fn is not None, but optimizer is None
 
         if net_inputs is not None:
-            network.set_inputs(*net_inputs)
+            _set_with_processed_inputs(network, net_inputs)
         return network
 
     def _build_eval_network(self, metrics, eval_network, eval_indexes):
@@ -362,17 +583,13 @@ class Model:
                                  f" optional, and then you can set `eval_network` or `loss_fn`. For the latter case,"
                                  f" framework will automatically build an evaluation network with `network` and"
                                  f" `loss_fn`.")
-
             net_inputs = self._network.get_inputs()
-            loss_inputs = [None]
-            if self._loss_fn.get_inputs():
-                loss_inputs = [*self._loss_fn.get_inputs()]
-            loss_inputs.pop(0)
-            if net_inputs:
-                net_inputs = [*net_inputs, *loss_inputs]
+            if self._loss_fn.get_inputs() and net_inputs:
+                loss_inputs = _process_loss_inputs(self._loss_fn.get_inputs())
+                net_inputs = _merge_inputs(net_inputs, loss_inputs)
             self._eval_network = nn.WithEvalCell(self._network, self._loss_fn, self._amp_level in ["O2", "O3", "auto"])
             if net_inputs is not None:
-                self._eval_network.set_inputs(*net_inputs)
+                _set_with_processed_inputs(self._eval_network, net_inputs)
             self._eval_indexes = [0, 1, 2]
 
     def _build_predict_network(self):
@@ -437,15 +654,16 @@ class Model:
             dataset.__loop_size__ = 1
 
         if dataset_helper is None:
+            logger.info("Begin to create DatasetHelper.")
             dataset_helper = DatasetHelper(dataset, dataset_sink_mode, sink_size, epoch_num)
 
         if dataset_sink_mode:
+            logger.info("Begin to connect network with dataset.")
             network = connect_network_with_dataset(network, dataset_helper)
 
         if _get_recovery_context("enable_recovery") and is_train:
             _set_training_dataset(dataset_helper)
 
-
         network.set_train(is_train)
         network.phase = phase
         self._backbone_is_train = is_train
@@ -459,8 +677,46 @@ class Model:
         if self._backbone_is_train != is_train:
             network.set_train(is_train)
             self._backbone_is_train = is_train
+        # Mode train and eval are the same net, network will be set_grad in _build_train_network.
+        # But if mode just want to  do predict or eval, must set network set_grad False
+        if not is_train:
+            network.set_grad(False)
         return network
 
+    def _check_need_ckpt(self, callbacks):
+        """Check callback list contain ckpt"""
+        need_ckpt = False
+        save_ckpt_steps = 1
+        last_triggered_step = 0
+        for cb in callbacks:
+            if isinstance(cb, ModelCheckpoint):
+                need_ckpt = True
+                cfg_size = cb._get_save_checkpoint_steps
+                save_ckpt_steps = save_ckpt_steps if (cfg_size is None or cfg_size >= sys.maxsize) else cfg_size
+                last_triggered_step = cb._get_last_trigger_step
+                break
+        return need_ckpt, save_ckpt_steps, last_triggered_step
+
+    def _store_training_step_info(self, cb_params):
+        """
+        cache train step info
+        :param cb_params: callback params
+        :return: none
+        """
+        if os.environ.get("MS_ENABLE_CKPT_D2H_ASYNC") != "1":
+            return
+        if (context.get_context("mode") == context.GRAPH_MODE) and (context.get_context("device_target") == "Ascend"):
+            cb_params.need_ckpt, cb_params.save_checkpoint_steps, \
+            cb_params.last_triggered_step = self._check_need_ckpt(cb_params.list_callback)
+            logger.info(f"need_ckpt:{cb_params.need_ckpt},"
+                        f"save_checkpoint_steps:{cb_params.save_checkpoint_steps},"
+                        f"cur_step_num:{cb_params.cur_step_num},"
+                        f"last_triggered_step:{cb_params.last_triggered_step}")
+            context.set_context(ascend_config={"need_ckpt": cb_params.need_ckpt,
+                                               "save_checkpoint_steps": cb_params.save_checkpoint_steps,
+                                               "cur_step_num": cb_params.cur_step_num,
+                                               "last_triggered_step": cb_params.last_triggered_step})
+
     def _warmup_dataset(self, epoch, train_dataset, sink_size=-1):
         """
         Trigger dataset pipeline running before graph compiling.
@@ -486,6 +742,22 @@ class Model:
         train_dataset._dataset_helper = dataset_helper
         train_dataset._warmup_epoch = epoch
 
+    def _waiting_for_dataset_warmup_ready(self, train_dataset):
+        """
+        Wait for the dataset to warmup until there is a batch of data available for training on the device side.
+
+        Args:
+            train_dataset (Dataset): A training dataset iterator. If `train_dataset` is defined, training graphs will be
+                                     initialized. Default: ``None``.
+        """
+        mbuf_size = train_dataset.__transfer_dataset__.get_mbuf_queue_size()
+        while mbuf_size == 0:
+            time.sleep(10)
+            mbuf_size = train_dataset.__transfer_dataset__.get_mbuf_queue_size()
+            if mbuf_size != 0:
+                break
+            logger.warning(f"Waiting for the dataset warmup, current device queue size: {mbuf_size}")
+
     def _init(self, train_dataset=None, valid_dataset=None, sink_size=-1, epoch=1):
         """
         Initialize compute graphs and data graphs with the sink mode.
@@ -507,6 +779,7 @@ class Model:
         if not train_dataset and not valid_dataset:
             raise ValueError("The argument 'train_dataset' and 'valid_dataset' can not both be None or empty.")
 
+        logger.info("Begin to check device number in model.build() procedure.")
         _device_number_check(self._parallel_mode, self._device_number)
 
         if train_dataset:
@@ -514,30 +787,34 @@ class Model:
                 raise TypeError("The type of 'train_dataset' must be `Dataset`, "
                                 "but got {}.".format(type(train_dataset)))
 
+            logger.info("Begin to check parameter broadcast in model.build() procedure.")
             _parameter_broadcast_check(self._parallel_mode, self._parameter_broadcast)
             if self._parameter_broadcast:
                 self._train_network.set_broadcast_flag()
 
+            logger.info("Begin to exec preprocess in model.build() procedure.")
             train_dataset.__no_send__ = True
             train_dataset_helper, train_network = self._exec_preprocess(is_train=True,
                                                                         dataset=train_dataset,
                                                                         dataset_sink_mode=True,
                                                                         sink_size=sink_size)
+            logger.info("Begin to warmup dataset in model.build() procedure.")
             self._warmup_dataset(epoch, train_dataset, sink_size)
-            if train_dataset.get_init_step() > 0:
-                mbuf_size = train_dataset.__transfer_dataset__.get_mbuf_queue_size()
-                while mbuf_size == 0:
-                    time.sleep(10)
-                    mbuf_size = train_dataset.__transfer_dataset__.get_mbuf_queue_size()
-                    if mbuf_size != 0:
-                        break
-                    logger.warning(f"Failover mode, waiting for dataset recover to specify step, "
-                                   f"current device queue size: {mbuf_size}")
+
+            # Since dataset pipeline has been triggered, delete flag
+            delattr(train_dataset, "__no_send__")
+
+            # Waiting for the dataset warmup ready
+            logger.info("Begin waiting for dataset warmup in model.build() procedure.")
+            self._waiting_for_dataset_warmup_ready(train_dataset)
+            logger.info("The dataset warmup was successful in model.build() procedure.")
 
             if context.get_auto_parallel_context("pipeline_stages") > 1 and valid_dataset:
                 train_network.add_flags_recursive(is_first_iteration=True)
             for inputs in train_dataset_helper:
+                logger.info("Begin to compile train network in model.build() procedure.")
                 train_network.compile(*inputs)
+                self._train_network.parameter_layout_dict = train_network.parameter_layout_dict
                 break
 
         if valid_dataset:
@@ -555,6 +832,7 @@ class Model:
             if context.get_auto_parallel_context("pipeline_stages") > 1:
                 eval_network.add_flags_recursive(is_first_iteration=False)
             for inputs in valid_dataset_helper:
+                logger.info("Begin to compile eval network in model.build() procedure.")
                 eval_network.compile(*inputs)
                 break
 
@@ -569,9 +847,10 @@ class Model:
 
         return [callbacks]
 
+    @_handle_tft
     @_save_final_ckpt
     def _train(self, epoch, train_dataset, callbacks=None, dataset_sink_mode=True, sink_size=-1, initial_epoch=0,
-               valid_dataset=None, valid_frequency=1, valid_dataset_sink_mode=True):
+               valid_dataset=None, valid_frequency=1, valid_dataset_sink_mode=True, initial_step=0):
         """
         Training.
 
@@ -595,12 +874,14 @@ class Model:
             self._train_network.set_broadcast_flag()
 
         cb_params = _InternalCallbackParam()
+        cb_params.cur_step_num = initial_step
         cb_params.train_network = self._train_network
         cb_params.epoch_num = epoch - initial_epoch
         if dataset_sink_mode and sink_size > 0:
             cb_params.batch_num = sink_size
         else:
             cb_params.batch_num = train_dataset.get_dataset_size()
+        self.batch_num = cb_params.batch_num
         cb_params.mode = "train"
         cb_params.loss_fn = self._loss_fn
         cb_params.optimizer = self._optimizer
@@ -610,6 +891,10 @@ class Model:
         cb_params.list_callback = self._transform_callbacks(callbacks)
         valid_infos = (valid_dataset, valid_frequency, valid_dataset_sink_mode)
         cb_params.list_callback.insert(0, _FrameworkProfilerCallback())
+        if os.environ.get("ENABLE_FLOPS_UTILIZATION_COLLECTOR") == "1" and \
+            FlopsUtilizationCollector not in cb_params.list_callback:
+            cb_params.list_callback.insert(0, FlopsUtilizationCollector(
+                cb_params.batch_num, full_flops=False))
         if context.get_context("mode") == context.PYNATIVE_MODE:
             cb_params.list_callback.insert(0, _StepSync())
         callbacks = cb_params.list_callback
@@ -625,11 +910,13 @@ class Model:
         with _CallbackManager(callbacks) as list_callback:
             self._check_reuse_dataset(train_dataset)
             if not dataset_sink_mode:
-                self._train_process(epoch, train_dataset, list_callback, cb_params, initial_epoch, valid_infos)
+                self._train_process(epoch, train_dataset, list_callback, cb_params, initial_epoch,
+                                    valid_infos)
             elif context.get_context("device_target") == "CPU":
                 logger.info("The CPU cannot support dataset sink mode currently."
                             "So the training process will be performed with dataset not sink.")
-                self._train_process(epoch, train_dataset, list_callback, cb_params, initial_epoch, valid_infos)
+                self._train_process(epoch, train_dataset, list_callback, cb_params, initial_epoch,
+                                    valid_infos)
             else:
                 self._train_dataset_sink_process(epoch, train_dataset, list_callback,
                                                  cb_params, sink_size, initial_epoch, valid_infos)
@@ -659,18 +946,17 @@ class Model:
         is_graph = (context.get_context("mode") == context.GRAPH_MODE)
         dataset_size = train_dataset.get_dataset_size()
         if dataset_size % sink_size != 0:
-            logger.warning("In dataset_sink mode (dataset_size % sink_size) should equal to 0, "
-                           "it is suggested to pad/drop data or adjust sink_size. "
-                           "But got 'dataset_size': {}, 'sink_size': {}.".format(dataset_size, sink_size))
+            logger.info("In dataset_sink mode (dataset_size % sink_size) should equal to 0, "
+                        "it is suggested to pad/drop data or adjust sink_size. "
+                        "But got 'dataset_size': {}, 'sink_size': {}.".format(dataset_size, sink_size))
         if sink_size == -1:
             dataset_sink_num = epoch
         else:
             dataset_sink_num = math.ceil(epoch * sink_size / dataset_size)
             train_dataset.__total_batch__ = epoch * sink_size
 
-        cb_params.cur_step_num = 0
+        cb_params.sink_size = sink_size
         cb_params.dataset_sink_mode = True
-
         run_context = RunContext(cb_params)
         list_callback.on_train_begin(run_context)
         # used to stop training for early stop, such as stopAtTIme or stopATStep
@@ -679,7 +965,6 @@ class Model:
             dataset_helper = train_dataset._dataset_helper
 
         self.epoch_iter = 0
-
         self._check_enable_recovery()
         # Used to check whether need perform recovery for process which is restarted.
         self._check_need_load_ckpt(cb_params, dataset_size, sink_size)
@@ -713,6 +998,7 @@ class Model:
                 else:
                     cb_params.cur_step_num += 1
                 self._current_step_num = int((cb_params.cur_step_num - 1) % cb_params.batch_num + 1)
+                self._store_training_step_info(cb_params)
                 cb_params.train_dataset_element = inputs
                 list_callback.on_train_step_begin(run_context)
                 train_network = self._check_network_mode(train_network, True)
@@ -814,7 +1100,6 @@ class Model:
             dataset_size (int): The number of batches in a dataset.
             sink_size (int): Control the amount of data in each sink. Default: -1.
         """
-
         if not self.enable_recovery:
             self.need_load_ckpt = False
 
@@ -901,7 +1186,6 @@ class Model:
                                                   dataset=train_dataset,
                                                   dataset_sink_mode=False,
                                                   epoch_num=epoch)
-        cb_params.cur_step_num = 0
         cb_params.dataset_sink_mode = False
         run_context = RunContext(cb_params)
         list_callback.on_train_begin(run_context)
@@ -923,7 +1207,6 @@ class Model:
                                      "returned by 'train_dataset'".format(len_element))
                 cb_params.cur_step_num += 1
                 self._current_step_num = int((cb_params.cur_step_num - 1) % cb_params.batch_num + 1)
-
                 cb_params.train_dataset_element = next_element
                 list_callback.on_train_step_begin(run_context)
                 self._check_network_mode(self._train_network, True)
@@ -1020,10 +1303,10 @@ class Model:
             >>> from mindspore.train import Model
             >>>
             >>> # Create the dataset taking MNIST as an example. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
             >>> dataset = create_dataset()
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
             >>> loss_scale_manager = ms.FixedLossScaleManager(1024., False)
@@ -1032,9 +1315,10 @@ class Model:
             ...                  loss_scale_manager=loss_scale_manager)
             >>> model.train(2, dataset)
         """
+        _check_tft()
+        device_target = context.get_context("device_target")
         # prepare dataset for obfuscated model
         train_dataset = self._prepare_obf_dataset(train_dataset)
-        device_target = context.get_context("device_target")
         if _is_ps_mode() and not _cache_enable() and (device_target in ["Ascend", "CPU"]) and dataset_sink_mode:
             logger.info("For PS mode, reset datasink mode to False when using Ascend or CPU backend.")
             dataset_sink_mode = False
@@ -1074,9 +1358,9 @@ class Model:
 
         _device_number_check(self._parallel_mode, self._device_number)
 
+        callbacks = _append_ccae(callbacks)
         if callbacks:
             self._check_methods_for_custom_callbacks(callbacks, "train")
-
         self._train(epoch,
                     train_dataset,
                     callbacks=callbacks,
@@ -1110,7 +1394,7 @@ class Model:
             callbacks = [callbacks]
         for cb in callbacks:
             cb_name = cb.__class__.__name__
-            if  cb_name not in internal_cb_names:
+            if cb_name not in internal_cb_names:
                 cb_methods_names = set(cb.__class__.__dict__.keys())
                 invalid_methods_names = cb_methods_names & old_version_methods_names
                 if invalid_methods_names:
@@ -1173,11 +1457,11 @@ class Model:
             >>> from mindspore.train import Model
             >>>
             >>> # Create the dataset taking MNIST as an example. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
             >>> train_dataset = create_dataset("train")
             >>> valid_dataset = create_dataset("test")
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
             >>> optim = nn.Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9)
@@ -1186,7 +1470,7 @@ class Model:
 
         Tutorial Examples:
             - `Advanced Encapsulation: Model - Train and Save Model
-              <https://www.mindspore.cn/tutorials/en/r2.2/advanced/model.html#training-and-saving-model>`_
+              <https://www.mindspore.cn/docs/en/master/model_train/train_process/model.html#training-and-saving-model>`_
         """
         device_target = context.get_context("device_target")
         if _is_ps_mode() and not _cache_enable() and (device_target in ["Ascend", "CPU"]) and dataset_sink_mode:
@@ -1266,10 +1550,10 @@ class Model:
             >>> from mindspore.amp import FixedLossScaleManager
             >>>
             >>> # Create the dataset taking MNIST as an example. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
             >>> dataset = create_dataset()
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> loss = nn.SoftmaxCrossEntropyWithLogits()
             >>> loss_scale_manager = FixedLossScaleManager()
@@ -1283,7 +1567,9 @@ class Model:
         if hasattr(self._train_network, '_is_check_and_refresh') and not self._train_network._is_check_and_refresh:
             self._train_network.check_names_and_refresh_name()
             self._train_network._is_check_and_refresh = True
+        logger.info("Begin to init dataset in model.build() procedure.")
         self._init(train_dataset, valid_dataset, sink_size, epoch)
+        logger.info("The model.build() which contains dataset warmup and network compile is success.")
 
     def _eval_in_fit(self, valid_dataset, callbacks=None, dataset_sink_mode=True, cb_params=None):
         """
@@ -1442,10 +1728,10 @@ class Model:
             >>> from mindspore.train import Model
             >>>
             >>> # Create the dataset taking MNIST as an example. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
             >>> dataset = create_dataset()
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
             >>> model = Model(net, loss_fn=loss, optimizer=None, metrics={'acc'})
@@ -1453,7 +1739,7 @@ class Model:
 
         Tutorial Examples:
             - `Advanced Encapsulation: Model - Train and Save Model
-              <https://www.mindspore.cn/tutorials/en/r2.2/advanced/model.html#training-and-saving-model>`_
+              <https://www.mindspore.cn/docs/en/master/model_train/train_process/model.html#training-and-saving-model>`_
         """
         valid_dataset = self._prepare_obf_dataset(valid_dataset)
         dataset_sink_mode = Validator.check_bool(dataset_sink_mode)
@@ -1473,6 +1759,10 @@ class Model:
         cb_params.mode = "eval"
         cb_params.cur_step_num = 0
         cb_params.list_callback = self._transform_callbacks(callbacks)
+        if os.environ.get("ENABLE_FLOPS_UTILIZATION_COLLECTOR") == "1" and \
+            FlopsUtilizationCollector not in cb_params.list_callback:
+            cb_params.list_callback.insert(0, FlopsUtilizationCollector(
+                cb_params.batch_num, full_flops=False))
         cb_params.network = self._network
 
         self._clear_metrics()
@@ -1686,7 +1976,7 @@ class Model:
                               "execution_plan" : {"op_name3" : "data_type:float16", "op_name4" : "data_type:float32"}}
 
                 Note that both the "configPath" is configured in the config_dict and the config_item,
-                    in this case, the path_b in the config_dict takes precedence.
+                in this case, the path_b in the config_dict takes precedence.
 
         Returns:
             Tensor, array(s) of predictions.
@@ -1699,7 +1989,7 @@ class Model:
             >>>
             >>> input_data = Tensor(np.random.randint(0, 255, [1, 1, 32, 32]), mindspore.float32)
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
             >>> model = Model(LeNet5())
             >>> result = model.predict(input_data)
         """
@@ -1719,8 +2009,25 @@ class Model:
                 self._lite_infer = False
                 logger.warning(f"Lite inference failed, {e.__str__()}, fallback to original inference!")
 
+        def _check_input_data():
+            """Input data check."""
+            for item in predict_data:
+                if item is None:
+                    continue
+                if isinstance(item, Tensor):
+                    if item.size == 0:
+                        msg = "The input data can not be empty."
+                        logger.critical(msg)
+                        raise ValueError(msg)
+                    continue
+                if not isinstance(item, (int, float, str)):
+                    data_class_str = "Tensor, None, int, float, str"
+                    raise TypeError(f'The types of input data must be in the Union({data_class_str}, ' \
+                                    f'tuple[{data_class_str}], list[{data_class_str}], dict[{data_class_str}]), ' \
+                                    f'but got type {item if item is None else type(item).__name__}.')
+
         self._check_network_mode(self._predict_network, False)
-        check_input_data(*predict_data, data_class=(int, float, str, None, Tensor))
+        _check_input_data()
         _parallel_predict_check()
         result = self._predict_network(*predict_data)
 
@@ -1807,10 +2114,10 @@ class Model:
             >>> ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.SEMI_AUTO_PARALLEL)
             >>>
             >>> # Create the dataset taking MNIST as an example. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/mnist.py
             >>> dataset = create_dataset()
             >>> # Define the network structure of LeNet5. Refer to
-            >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+            >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
             >>> net = LeNet5()
             >>> loss = nn.SoftmaxCrossEntropyWithLogits()
             >>> loss_scale_manager = ms.FixedLossScaleManager()
@@ -1832,7 +2139,6 @@ class Model:
         train_dataset.__model_hash__ = hash(self)
         return train_network.parameter_layout_dict
 
-
     def infer_predict_layout(self, *predict_data, skip_backend_compile=False):
         """
         Generate parameter layout for the predict network in 'AUTO_PARALLEL' or 'SEMI_AUTO_PARALLEL' mode.