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

mindspore/train/model.py

@@ -36,7 +36,7 @@ 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,\
-    FlopsUtilizationCollector, MindIOTTPAdapter
+    FlopsUtilizationCollector, TFTRegister
 from mindspore.train.callback import __all__ as internal_cb_names
 from mindspore.train.callback._cluster_monitor import ClusterMonitor
 from mindspore import context
@@ -46,6 +46,7 @@ from mindspore.parallel._utils import _get_parallel_mode, _get_device_num, _get_
 from mindspore.parallel._ps_context import _is_role_worker, _is_role_pserver, _is_ps_mode, \
     _cache_enable, _enable_distributed_mindrt
 from mindspore.train.metrics import Loss
+from mindspore.train._utils import vlog_print
 from mindspore import nn
 from mindspore.boost import AutoBoost
 from mindspore.context import ParallelMode
@@ -119,6 +120,101 @@ 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."""
@@ -290,21 +386,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"`` .
@@ -379,6 +465,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"""
@@ -568,9 +655,13 @@ class Model:
             dataset.__loop_size__ = 1
 
         if dataset_helper is None:
+            vlog_print("1", "ME", __file__, sys._getframe().f_lineno, "Begin to create DatasetHelper.")
+            logger.info("Begin to create DatasetHelper.")
             dataset_helper = DatasetHelper(dataset, dataset_sink_mode, sink_size, epoch_num)
 
         if dataset_sink_mode:
+            vlog_print("1", "ME", __file__, sys._getframe().f_lineno, "Begin to connect network with dataset.")
+            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:
@@ -589,6 +680,10 @@ 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):
@@ -687,6 +782,8 @@ 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.")
 
+        vlog_print("1", "ME", __file__, sys._getframe().f_lineno, "Begin to check device number in model.build().")
+        logger.info("Begin to check device number in model.build() procedure.")
         _device_number_check(self._parallel_mode, self._device_number)
 
         if train_dataset:
@@ -694,27 +791,44 @@ class Model:
                 raise TypeError("The type of 'train_dataset' must be `Dataset`, "
                                 "but got {}.".format(type(train_dataset)))
 
+            vlog_print("1", "ME", __file__, sys._getframe().f_lineno,
+                       "Begin to check parameter broadcast in model.build().")
+            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()
 
+            vlog_print("1", "ME", __file__, sys._getframe().f_lineno, "Begin to exec preprocess in model.build().")
+            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)
+            vlog_print("1", "ME", __file__, sys._getframe().f_lineno, "Begin to warmup dataset in model.build().")
+            logger.info("Begin to warmup dataset in model.build() procedure.")
             self._warmup_dataset(epoch, train_dataset, sink_size)
 
             # Since dataset pipeline has been triggered, delete flag
             delattr(train_dataset, "__no_send__")
 
             # Waiting for the dataset warmup ready
+            vlog_print("1", "ME", __file__, sys._getframe().f_lineno,
+                       "Begin waiting for dataset warmup in model.build().")
+            logger.info("Begin waiting for dataset warmup in model.build() procedure.")
             self._waiting_for_dataset_warmup_ready(train_dataset)
+            vlog_print("1", "ME", __file__, sys._getframe().f_lineno,
+                       "The dataset warmup was successful in model.build().")
+            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:
+                vlog_print("1", "ME", __file__, sys._getframe().f_lineno,
+                           "Begin to compile train network in model.build().")
+                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:
@@ -732,6 +846,9 @@ 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:
+                vlog_print("1", "ME", __file__, sys._getframe().f_lineno,
+                           "Begin to compile eval network in model.build().")
+                logger.info("Begin to compile eval network in model.build() procedure.")
                 eval_network.compile(*inputs)
                 break
 
@@ -746,9 +863,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.
 
@@ -772,12 +890,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
@@ -801,16 +921,19 @@ class Model:
             epoch = 1
         cb_params.last_save_ckpt_step = None
         cb_params.latest_ckpt_file = None
+        cb_params.loss_scale_mananger = self._loss_scale_manager
 
         # build callback list
         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)
@@ -850,9 +973,7 @@ class Model:
             train_dataset.__total_batch__ = epoch * sink_size
 
         cb_params.sink_size = sink_size
-        cb_params.cur_step_num = 0
         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
@@ -861,7 +982,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)
@@ -997,7 +1117,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
 
@@ -1084,7 +1203,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)
@@ -1106,7 +1224,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)
@@ -1150,31 +1267,6 @@ class Model:
 
         list_callback.on_train_end(run_context)
 
-    def _wrapper_train(self, callbacks):
-        """
-        This method used to wrap train function with ttp wrapper which will do event notify when
-        exceptions throw.
-
-        Args:
-            callbacks (function): Callbacks passed by train method.
-        """
-
-        if not callbacks:
-            return self._train
-        cbs = callbacks if isinstance(callbacks, list) else [callbacks]
-        obj = None
-        _train_wrapper = None
-        for item in cbs:
-            if isinstance(item, MindIOTTPAdapter):
-                obj = item
-
-        if (obj is not None) and (obj.enable is True):
-            logger.info("MindIO TTP is enable, so we wrapper ttp exception handdler for self train method.")
-            _train_wrapper = obj.wrapper_ttp_persist(self._train)
-
-        return self._train if not _train_wrapper else _train_wrapper
-
-
     def train(self, epoch, train_dataset, callbacks=None, dataset_sink_mode=False, sink_size=-1, initial_epoch=0):
         """
         Training API.
@@ -1240,9 +1332,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
@@ -1283,16 +1376,14 @@ class Model:
         _device_number_check(self._parallel_mode, self._device_number)
 
         callbacks = _append_ccae(callbacks)
-        _train_wrapper = None
         if callbacks:
             self._check_methods_for_custom_callbacks(callbacks, "train")
-        _train_wrapper = self._wrapper_train(callbacks)
-        _train_wrapper(epoch,
-                       train_dataset,
-                       callbacks=callbacks,
-                       dataset_sink_mode=dataset_sink_mode,
-                       sink_size=sink_size,
-                       initial_epoch=initial_epoch)
+        self._train(epoch,
+                    train_dataset,
+                    callbacks=callbacks,
+                    dataset_sink_mode=dataset_sink_mode,
+                    sink_size=sink_size,
+                    initial_epoch=initial_epoch)
 
         # When it's distributed training and using MindRT,
         # the node id should be reset to start from 0.
@@ -1396,7 +1487,7 @@ class Model:
 
         Tutorial Examples:
             - `Advanced Encapsulation: Model - Train and Save Model
-              <https://www.mindspore.cn/tutorials/en/master/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:
@@ -1493,7 +1584,12 @@ 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
+        vlog_print("1", "ME", __file__, sys._getframe().f_lineno, "Begin to init dataset in model.build().")
+        logger.info("Begin to init dataset in model.build() procedure.")
         self._init(train_dataset, valid_dataset, sink_size, epoch)
+        vlog_print("1", "ME", __file__, sys._getframe().f_lineno,
+                   "The model.build() which contains dataset warmup and network compile is success.")
+        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):
         """
@@ -1663,7 +1759,7 @@ class Model:
 
         Tutorial Examples:
             - `Advanced Encapsulation: Model - Train and Save Model
-              <https://www.mindspore.cn/tutorials/en/master/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)