mindspore 2.0.0a0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/train/callback/_landscape.py

@@ -282,7 +282,7 @@ class SummaryLandscape:
                   For example, in image dataset, You can set num_samples is 2048,
                   which means that 2048 images are used to create loss landscape.
                   Default: 2048.
-                - intervals (List[List[int]): Specifies the interval
+                - intervals (List[List[int]]): Specifies the interval
                   in which the loss landscape. For example: If the user wants to
                   create loss landscape of two training processes, they are 1-5 epoch
                   and 6-10 epoch respectively. They can set [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]].
@@ -294,6 +294,12 @@ class SummaryLandscape:
                 Default: None. The default save path is the same as the summary file.
         """
 
+        executor = None
+        if len(device_ids) > 1:
+            executor = ProcessPoolExecutor(len(device_ids))
+            futures = [executor.submit(self._set_context, i) for i in device_ids]
+            wait(futures, return_when=ALL_COMPLETED)
+
         output_path = os.path.realpath(output) if output is not None else self._summary_dir
         summary_record = SummaryRecord(output_path)
         self._check_device_ids(device_ids)
@@ -321,13 +327,14 @@ class SummaryLandscape:
                 json.dump(data, file)
             os.chmod(json_path, stat.S_IRUSR)
 
-        for interval, landscape in self._list_landscapes(callback_fn=callback_fn, device_ids=device_ids):
+        for interval, landscape in self._list_landscapes(callback_fn=callback_fn, executor=executor,
+                                                         device_ids=device_ids):
             summary_record.add_value(PluginEnum.LANDSCAPE.value, f'landscape_{str(interval)}', landscape)
             summary_record.record(0)
             summary_record.flush()
         summary_record.close()
 
-    def _list_landscapes(self, callback_fn, device_ids=None):
+    def _list_landscapes(self, callback_fn, executor=None, device_ids=None):
         """Create landscape with single device and list all landscape."""
 
         if not os.path.exists(os.path.join(self._ckpt_dir, 'train_metadata.json')):
@@ -343,39 +350,30 @@ class SummaryLandscape:
         kwargs = dict(proz=0.2, landscape_size=data['landscape_size'], device_ids=device_ids, callback_fn=callback_fn)
 
         start = time.time()
-        with ProcessPoolExecutor(max_workers=len(device_ids)) as executor:
-            if len(device_ids) > 1:
-                futures = []
-                for device_id in device_ids:
-                    future = executor.submit(self._set_context, device_id)
-                    futures.append(future)
-                wait(futures, return_when=ALL_COMPLETED)
-
-            kwargs['executor'] = executor if len(device_ids) > 1 else None
-
-            if data['create_landscape']['train']:
-                for i, epochs in enumerate(self._epoch_group.values()):
-                    self._log_message(data['create_landscape'], index=i, interval=epochs)
-                    kwargs['epochs'] = epochs
-                    mid_time = time.time()
-                    landscape_data = self._create_landscape_by_pca(**kwargs)
-                    logger.info("Create landscape end, use time: %s s." % (round(time.time() - mid_time, 6)))
-                    landscape_data.unit = data['unit']
-                    landscape_data.step_per_epoch = data['step_per_epoch']
-                    landscape_data.num_samples = data['num_samples']
-                    yield [epochs[0], epochs[-1]], landscape_data.transform_to_loss_landscape_msg(landscape_data)
-
-            if data['create_landscape']['result']:
-                final_epochs = [list(self._epoch_group.values())[-1][-1]]
-                self._log_message(data['create_landscape'], final_epochs=final_epochs)
-                kwargs['epochs'] = final_epochs
+        kwargs['executor'] = executor
+        if data['create_landscape']['train']:
+            for i, epochs in enumerate(self._epoch_group.values()):
+                self._log_message(data['create_landscape'], index=i, interval=epochs)
+                kwargs['epochs'] = epochs
                 mid_time = time.time()
-                landscape_data = self._create_landscape_by_random(**kwargs)
+                landscape_data = self._create_landscape_by_pca(**kwargs)
                 logger.info("Create landscape end, use time: %s s." % (round(time.time() - mid_time, 6)))
                 landscape_data.unit = data['unit']
                 landscape_data.step_per_epoch = data['step_per_epoch']
                 landscape_data.num_samples = data['num_samples']
-                yield final_epochs, landscape_data.transform_to_loss_landscape_msg(landscape_data)
+                yield [epochs[0], epochs[-1]], landscape_data.transform_to_loss_landscape_msg(landscape_data)
+
+        if data['create_landscape']['result']:
+            final_epochs = [list(self._epoch_group.values())[-1][-1]]
+            self._log_message(data['create_landscape'], final_epochs=final_epochs)
+            kwargs['epochs'] = final_epochs
+            mid_time = time.time()
+            landscape_data = self._create_landscape_by_random(**kwargs)
+            logger.info("Create landscape end, use time: %s s." % (round(time.time() - mid_time, 6)))
+            landscape_data.unit = data['unit']
+            landscape_data.step_per_epoch = data['step_per_epoch']
+            landscape_data.num_samples = data['num_samples']
+            yield final_epochs, landscape_data.transform_to_loss_landscape_msg(landscape_data)
         logger.info("Total use time: %s s." % (round(time.time() - start, 6)))
 
     def _log_message(self, create_landscape, index=None, interval=None, final_epochs=None):

mindspore/train/callback/_loss_monitor.py

@@ -17,7 +17,7 @@ from __future__ import absolute_import
 
 import numpy as np
 
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.train.callback._callback import Callback, _handle_loss
 
 
@@ -41,8 +41,8 @@ class LossMonitor(Callback):
         .. note::
             Before running the following example, you need to customize the network LeNet5 and
             dataset preparation function create_dataset. Refer to
-            `Building a Network <https://www.mindspore.cn/tutorials/en/r2.0.0-alpha/beginner/model.html>`_
-            and `Dataset <https://www.mindspore.cn/tutorials/en/r2.0.0-alpha/beginner/dataset.html>`_ .
+            `Building a Network <https://www.mindspore.cn/tutorials/en/r2.0/beginner/model.html>`_
+            and `Dataset <https://www.mindspore.cn/tutorials/en/r2.0/beginner/dataset.html>`_ .
 
         >>> from mindspore import nn
         >>> from mindspore.train import Model, LossMonitor

mindspore/train/callback/_on_request_exit.py

@@ -19,7 +19,7 @@ import os
 import signal
 
 from mindspore import log
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.train.serialization import load_checkpoint, save_checkpoint, export
 from mindspore.train.callback._callback import Callback
 
@@ -95,8 +95,8 @@ class OnRequestExit(Callback):
             self.train_file_path = os.path.abspath(os.path.join(directory, f"{file_name}_train"))
             self.eval_file_path = os.path.abspath(os.path.join(directory, f"{file_name}_eval"))
         self.sig = Validator.check_isinstance('sig', sig, int)
-        if self.sig == signal.SIGKILL or self.sig == signal.SIGINT:
-            raise ValueError("Not support send exit request by signal SIGKILL or SIGINT.")
+        if hasattr(signal, "SIGKILL") and self.sig == signal.SIGKILL:
+            raise ValueError("Not support send exit request by signal SIGKILL.")
         self.exit = False
 
     def on_train_begin(self, run_context):

mindspore/train/callback/_reduce_lr_on_plateau.py

@@ -20,7 +20,7 @@ import numpy as np
 
 from mindspore.common.tensor import Tensor
 from mindspore.common.parameter import Parameter
-from mindspore._checkparam import Validator, Rel
+from mindspore import _checkparam as Validator
 from mindspore import log as logger
 from mindspore.ops import functional as F, ReduceOp
 from mindspore import nn, ops
@@ -84,8 +84,8 @@ class ReduceLROnPlateau(Callback):
         .. note::
             Before running the following example, you need to customize the network LeNet5 and
             dataset preparation function create_dataset. Refer to
-            `Building a Network <https://www.mindspore.cn/tutorials/en/r2.0.0-alpha/beginner/model.html>`_
-            and `Dataset <https://www.mindspore.cn/tutorials/en/r2.0.0-alpha/beginner/dataset.html>`_ .
+            `Building a Network <https://www.mindspore.cn/tutorials/en/r2.0/beginner/model.html>`_
+            and `Dataset <https://www.mindspore.cn/tutorials/en/r2.0/beginner/dataset.html>`_ .
 
         >>> from mindspore import nn
         >>> from mindspore.train import Model, ReduceLROnPlateau
@@ -102,7 +102,7 @@ class ReduceLROnPlateau(Callback):
                  mode='auto', min_delta=1e-4, cooldown=0, min_lr=0):
         super(ReduceLROnPlateau, self).__init__()
         self.monitor = Validator.check_value_type('monitor', monitor, str)
-        self.factor = Validator.check_float_range(factor, 0.0, 1.0, Rel.INC_NEITHER)
+        self.factor = Validator.check_float_range(factor, 0.0, 1.0, Validator.INC_NEITHER)
         self.patience = Validator.check_non_negative_int(patience)
         self.verbose = Validator.check_bool(verbose)
         self.mode = Validator.check_value_type('mode', mode, str)

mindspore/train/callback/_summary_collector.py

@@ -103,7 +103,7 @@ class SummaryCollector(Callback):
               training computational graph is collected. Default: True.
             - collect_train_lineage (bool): Whether to collect lineage data for the training phase,
               this field will be displayed on the `lineage page \
-              <https://www.mindspore.cn/mindinsight/docs/en/r2.0.0-alpha/lineage_and_scalars_comparison.html>`_
+              <https://www.mindspore.cn/mindinsight/docs/en/r2.0/lineage_and_scalars_comparison.html>`_
               of MindInsight. Default: True.
             - collect_eval_lineage (bool): Whether to collect lineage data for the evaluation phase,
               this field will be displayed on the lineage page of MindInsight. Default: True.
@@ -488,7 +488,26 @@ class SummaryCollector(Callback):
             self._record.add_value('custom_lineage_data', 'custom_lineage_data', packaged_custom_data)
             self._has_saved_custom_data = True
             self._record.record(cb_params.cur_step_num)
+        if not self._dataset_sink_mode:
+            self._collect_tensor_data(cb_params)
 
+        collect_landscape = self._collect_specified_data.get('collect_landscape')
+        if collect_landscape is not None:
+            intervals = collect_landscape.get('intervals')
+            collect_interval = False
+            for interval in intervals:
+                if "cur_step_num" in cb_params:
+                    if cb_params.cur_step_num in interval:
+                        collect_interval = True
+                        break
+
+        if collect_landscape and collect_landscape.get('unit', 'step') == 'step' and collect_interval:
+            self._save_model_params_for_landscape(cb_params)
+
+    def _collect_tensor_data(self, cb_params):
+        """Collect tensor summary data."""
+        if cb_params.mode != ModeEnum.TRAIN.value:
+            return
         if self._first_step:
             self._tensor_collect_range = self._get_tensor_collect_range(cb_params, self._dataset_sink_mode)
             self._collect_at_step_end(cb_params, plugin_filter=None)
@@ -503,19 +522,6 @@ class SummaryCollector(Callback):
             elif current % self._collect_freq == 0:
                 self._collect_at_step_end(cb_params, lambda plugin: plugin != PluginEnum.TENSOR.value)
 
-        collect_landscape = self._collect_specified_data.get('collect_landscape')
-        if collect_landscape is not None:
-            intervals = collect_landscape.get('intervals')
-            collect_interval = False
-            for interval in intervals:
-                if "cur_step_num" in cb_params:
-                    if cb_params.cur_step_num in interval:
-                        collect_interval = True
-                        break
-
-        if collect_landscape and collect_landscape.get('unit', 'step') == 'step' and collect_interval:
-            self._save_model_params_for_landscape(cb_params)
-
     def _get_tensor_collect_range(self, cb_params, dataset_sink_mode):
         """Get tensor collect range."""
         total_step = cb_params.epoch_num
@@ -542,6 +548,7 @@ class SummaryCollector(Callback):
 
     def epoch_end(self, run_context):
         cb_params = run_context.original_args()
+        self._collect_tensor_data(cb_params)
         collect_landscape = self._collect_specified_data.get('collect_landscape')
         if collect_landscape is not None:
             intervals = collect_landscape.get('intervals')
@@ -863,14 +870,14 @@ class SummaryCollector(Callback):
         if regular is not None:
             for parameter in parameters:
                 if re.match(regular, parameter.name):
-                    self._record.add_value(PluginEnum.HISTOGRAM.value, parameter.name+'/auto', parameter.data)
+                    self._record.add_value(PluginEnum.HISTOGRAM.value, parameter.name + '/auto', parameter.data)
             return
 
         # Note: If `histogram_regular` in `self._collect_specified_data` and the value is None,
         # we will collect the first five parameters.
         default_parameter_count = 5
         for parameter in parameters[:default_parameter_count]:
-            self._record.add_value(PluginEnum.HISTOGRAM.value, parameter.name+'/auto', parameter.data)
+            self._record.add_value(PluginEnum.HISTOGRAM.value, parameter.name + '/auto', parameter.data)
 
     @staticmethod
     def _get_learning_rate(optimizer):

mindspore/train/callback/_time_monitor.py

@@ -17,7 +17,7 @@ from __future__ import absolute_import
 
 import time
 
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.train.callback._callback import Callback
 
 
@@ -37,8 +37,8 @@ class TimeMonitor(Callback):
         .. note::
             Before running the following example, you need to customize the network LeNet5 and
             dataset preparation function create_dataset. Refer to
-            `Building a Network <https://www.mindspore.cn/tutorials/en/r2.0.0-alpha/beginner/model.html>`_
-            and `Dataset <https://www.mindspore.cn/tutorials/en/r2.0.0-alpha/beginner/dataset.html>`_ .
+            `Building a Network <https://www.mindspore.cn/tutorials/en/r2.0/beginner/model.html>`_
+            and `Dataset <https://www.mindspore.cn/tutorials/en/r2.0/beginner/dataset.html>`_ .
 
         >>> from mindspore import nn
         >>> from mindspore.train import Model, TimeMonitor

mindspore/train/checkpoint_pb2.py

@@ -19,7 +19,7 @@ DESCRIPTOR = _descriptor.FileDescriptor(
   syntax='proto2',
   serialized_options=None,
   create_key=_descriptor._internal_create_key,
-  serialized_pb=b'\n\x10\x63heckpoint.proto\"b\n\nCheckpoint\x12 \n\x05value\x18\x01 \x03(\x0b\x32\x11.Checkpoint.Value\x1a\x32\n\x05Value\x12\x0b\n\x03tag\x18\x01 \x02(\t\x12\x1c\n\x06tensor\x18\x02 \x02(\x0b\x32\x0c.TensorProto\"H\n\x0bTensorProto\x12\x0c\n\x04\x64ims\x18\x01 \x03(\x03\x12\x13\n\x0btensor_type\x18\x02 \x02(\t\x12\x16\n\x0etensor_content\x18\x03 \x02(\x0c'
+  serialized_pb=b'\n\x10\x63heckpoint.proto\"\x93\x01\n\nCheckpoint\x12 \n\x05value\x18\x01 \x03(\x0b\x32\x11.Checkpoint.Value\x1a\x63\n\x05Value\x12\x0b\n\x03tag\x18\x01 \x02(\t\x12\x1e\n\x06tensor\x18\x02 \x01(\x0b\x32\x0c.TensorProtoH\x00\x12$\n\tmaptensor\x18\x03 \x01(\x0b\x32\x0f.MapTensorProtoH\x00\x42\x07\n\x05value\"H\n\x0bTensorProto\x12\x0c\n\x04\x64ims\x18\x01 \x03(\x03\x12\x13\n\x0btensor_type\x18\x02 \x02(\t\x12\x16\n\x0etensor_content\x18\x03 \x02(\x0c\".\n\x0eMapTensorProto\x12\x1c\n\x06tensor\x18\x01 \x03(\x0b\x32\x0c.TensorProto'
 )
 
 
@@ -42,7 +42,14 @@ _CHECKPOINT_VALUE = _descriptor.Descriptor(
       serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
     _descriptor.FieldDescriptor(
       name='tensor', full_name='Checkpoint.Value.tensor', index=1,
-      number=2, type=11, cpp_type=10, label=2,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='maptensor', full_name='Checkpoint.Value.maptensor', index=2,
+      number=3, type=11, cpp_type=10, label=1,
       has_default_value=False, default_value=None,
       message_type=None, enum_type=None, containing_type=None,
       is_extension=False, extension_scope=None,
@@ -58,9 +65,14 @@ _CHECKPOINT_VALUE = _descriptor.Descriptor(
   syntax='proto2',
   extension_ranges=[],
   oneofs=[
+    _descriptor.OneofDescriptor(
+      name='value', full_name='Checkpoint.Value.value',
+      index=0, containing_type=None,
+      create_key=_descriptor._internal_create_key,
+    fields=[]),
   ],
-  serialized_start=68,
-  serialized_end=118,
+  serialized_start=69,
+  serialized_end=168,
 )
 
 _CHECKPOINT = _descriptor.Descriptor(
@@ -90,8 +102,8 @@ _CHECKPOINT = _descriptor.Descriptor(
   extension_ranges=[],
   oneofs=[
   ],
-  serialized_start=20,
-  serialized_end=118,
+  serialized_start=21,
+  serialized_end=168,
 )
 
 
@@ -136,15 +148,56 @@ _TENSORPROTO = _descriptor.Descriptor(
   extension_ranges=[],
   oneofs=[
   ],
-  serialized_start=120,
-  serialized_end=192,
+  serialized_start=170,
+  serialized_end=242,
+)
+
+
+_MAPTENSORPROTO = _descriptor.Descriptor(
+  name='MapTensorProto',
+  full_name='MapTensorProto',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='tensor', full_name='MapTensorProto.tensor', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=244,
+  serialized_end=290,
 )
 
 _CHECKPOINT_VALUE.fields_by_name['tensor'].message_type = _TENSORPROTO
+_CHECKPOINT_VALUE.fields_by_name['maptensor'].message_type = _MAPTENSORPROTO
 _CHECKPOINT_VALUE.containing_type = _CHECKPOINT
+_CHECKPOINT_VALUE.oneofs_by_name['value'].fields.append(
+  _CHECKPOINT_VALUE.fields_by_name['tensor'])
+_CHECKPOINT_VALUE.fields_by_name['tensor'].containing_oneof = _CHECKPOINT_VALUE.oneofs_by_name['value']
+_CHECKPOINT_VALUE.oneofs_by_name['value'].fields.append(
+  _CHECKPOINT_VALUE.fields_by_name['maptensor'])
+_CHECKPOINT_VALUE.fields_by_name['maptensor'].containing_oneof = _CHECKPOINT_VALUE.oneofs_by_name['value']
 _CHECKPOINT.fields_by_name['value'].message_type = _CHECKPOINT_VALUE
+_MAPTENSORPROTO.fields_by_name['tensor'].message_type = _TENSORPROTO
 DESCRIPTOR.message_types_by_name['Checkpoint'] = _CHECKPOINT
 DESCRIPTOR.message_types_by_name['TensorProto'] = _TENSORPROTO
+DESCRIPTOR.message_types_by_name['MapTensorProto'] = _MAPTENSORPROTO
 _sym_db.RegisterFileDescriptor(DESCRIPTOR)
 
 Checkpoint = _reflection.GeneratedProtocolMessageType('Checkpoint', (_message.Message,), {
@@ -169,5 +222,12 @@ TensorProto = _reflection.GeneratedProtocolMessageType('TensorProto', (_message.
   })
 _sym_db.RegisterMessage(TensorProto)
 
+MapTensorProto = _reflection.GeneratedProtocolMessageType('MapTensorProto', (_message.Message,), {
+  'DESCRIPTOR' : _MAPTENSORPROTO,
+  '__module__' : 'checkpoint_pb2'
+  # @@protoc_insertion_point(class_scope:MapTensorProto)
+  })
+_sym_db.RegisterMessage(MapTensorProto)
+
 
 # @@protoc_insertion_point(module_scope)

mindspore/train/data_sink.py

@@ -23,7 +23,7 @@ from mindspore.train.dataset_helper import _has_dynamic_shape, _check_inputs
 import mindspore.dataset as ds
 from mindspore._c_expression import _set_dataset_mode_config
 from mindspore.parallel._utils import _get_device_num, _need_to_full, _to_full_shapes, _get_pipeline_stages
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 
 
 def _init_sink_dataset(dataset, sink_size, input_signature, create_info):
@@ -52,7 +52,7 @@ def _init_sink_dataset(dataset, sink_size, input_signature, create_info):
         _check_inputs(input_signature, dataset_shapes, dataset_types)
 
     queue_name = transfer_dataset.queue_name
-    if _need_to_full():
+    if _need_to_full() and context.get_context('mode') == context.GRAPH_MODE:
         device_num = _get_device_num() // _get_pipeline_stages()
         dataset_shapes = _to_full_shapes(dataset_shapes, device_num)
     next_op = ops.GetNext(dataset_types, dataset_shapes, len(dataset_types), queue_name)
@@ -90,7 +90,7 @@ def _get_next_op(dataset, ori_next_op, is_info_queue):
     if key in dataset.__sink_aux__.next_ops:
         next_op = dataset.__sink_aux__.next_ops[key]
     else:
-        if _need_to_full():
+        if _need_to_full() and context.get_context('mode') == context.GRAPH_MODE:
             device_num = _get_device_num() // _get_pipeline_stages()
             dataset_shapes = _to_full_shapes(dataset_shapes, device_num)
         next_op = ops.GetNext(dataset_types, dataset_shapes, len(dataset_types), queue_name)
@@ -129,6 +129,18 @@ def data_sink(fn, dataset, sink_size=1, jit_config=None, input_signature=None):
     """
     A wrapper function to generate a function for the input function.
 
+    Note:
+        When using data sinking, the dataset will automatically cycle. At this time,
+        only the total number of training steps (total_step) and the number of steps for each sinking (sink_size)
+        need to be considered. When switching from training by rounds (epochs) to data sinking,
+        the formula is as follows:
+
+        total_step = epochs * dataset_size
+
+        train_sink_step = total_step / sink_size
+
+        After transforming from `mindspore.data_sink`, you need to execute `train_sink_step` step for training.
+
     Args:
         fn (Function): The Python function that will be run with dataset.
         dataset (Dataset): The dataset iterator. The dataset can be generated by dataset generator API in

mindspore/train/dataset_helper.py

@@ -17,7 +17,7 @@ from __future__ import absolute_import
 
 import math
 
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.common.dtype import pytype_to_dtype
 from mindspore.common.api import _cell_graph_executor
 from mindspore.common._utils import is_shape_unknown
@@ -27,8 +27,7 @@ from mindspore import context, nn
 from mindspore.train._utils import _exec_datagraph, _get_types_and_shapes, _construct_tensor_list
 from mindspore.parallel._utils import _get_device_num, _get_global_rank, _need_to_full, \
     _to_full_shapes, _get_pipeline_stages
-from mindspore.parallel._ps_context import _is_role_worker, _is_role_sched, _is_ps_mode, \
-                                           _enable_distributed_mindrt
+from mindspore.parallel._ps_context import _is_role_sched
 from mindspore.ops import operations as P
 
 
@@ -64,8 +63,7 @@ def _dynamic_sink_exception_scenario(dataset_iter):
     """The exception scenario for dynamic data is not applicable."""
     _, dataset_shapes = dataset_iter.types_shapes()
 
-    if _has_dynamic_shape(dataset_shapes) or (_is_role_worker() and _is_ps_mode()) or \
-       context.get_context("mode") != context.GRAPH_MODE:
+    if _has_dynamic_shape(dataset_shapes) or context.get_context("mode") != context.GRAPH_MODE:
         return True
     return False
 
@@ -95,6 +93,8 @@ class _DataWrapper(nn.Cell):
         self.get_next = P.GetNext(
             dataset_types, dataset_shapes, len(dataset_types), queue_name)
         self.network = network
+        if isinstance(network, nn.Cell) and network.jit_config_dict:
+            self._jit_config_dict = network.jit_config_dict
 
     def construct(self):
         outputs = self.get_next()
@@ -133,6 +133,8 @@ def _check_inputs(network_shapes, dataset_shapes, dataset_types):
     Check if set inputs are correct.
     """
     for tensor_index, ele_dataset_shape in enumerate(dataset_shapes):
+        if network_shapes[tensor_index] is None:
+            continue
         set_inputs_shape = list(network_shapes[tensor_index].shape)
         inputs_shape = list(ele_dataset_shape)
         if dataset_types[tensor_index] != network_shapes[tensor_index].dtype:
@@ -145,8 +147,6 @@ def _check_inputs(network_shapes, dataset_shapes, dataset_types):
             raise ValueError(
                 f"The {tensor_index + 1}th input dims of 'set_inputs' must be the same as network's input, "
                 f"but got 'set_inputs': {len(set_inputs_shape)} and network's input: {len(inputs_shape)}.")
-        if network_shapes[tensor_index] is None:
-            break
         for index, ele_shape in enumerate(ele_dataset_shape):
             if network_shapes[tensor_index].shape[index] != -1:
                 if set_inputs_shape[index] != ele_shape:
@@ -396,7 +396,7 @@ class DatasetHelper:
 
     def _reset(self, step, epoch):
         """Reset the dataset to the provided step and epoch."""
-        self.iter._reset(step, epoch)  # pylint: disable=W0212
+        self.iter._reset(step, epoch)  # pylint: disable=protected-access
 
     def get_data_info(self):
         """
@@ -420,9 +420,7 @@ class _DatasetIter:
             ds.config.set_dynamic_shape(True)
         if not hasattr(dataset, '__transfer_dataset__'):
             if hasattr(dataset, '__loop_size__'):
-                # PS mode does not support loop sink and need get the real sink size.
-                if not (_is_role_worker() and _is_ps_mode()) or _enable_distributed_mindrt():
-                    self.sink_size = dataset.__loop_size__
+                self.sink_size = dataset.__loop_size__
             create_data_info_queue = (sink_size == 1 and self.sink_count == 1 and dataset.get_dataset_size() != 1
                                       and not self.dynamic_shape)
             dataset.__transfer_dataset__ = _exec_datagraph(dataset, self.sink_size,
@@ -442,7 +440,7 @@ class _DatasetIter:
         self.continue_send = dataset.__transfer_dataset__.continue_send
         self.get_data_info = dataset.__transfer_dataset__.get_data_info
         if hasattr(dataset.__transfer_dataset__, "_reset"):
-            self._reset = dataset.__transfer_dataset__._reset  # pylint: disable=W0212
+            self._reset = dataset.__transfer_dataset__._reset  # pylint: disable=protected-access
 
     def __iter__(self):
         self.index = 0
@@ -476,9 +474,6 @@ class _DatasetIter:
         sink_size = 1
         if hasattr(self.dataset, '__loop_size__'):
             sink_size = self.dataset.__loop_size__
-        elif _is_role_worker() and _is_ps_mode() and not _enable_distributed_mindrt():
-            # PS mode does not support loop sink.
-            sink_size = 1
         else:
             if context.get_context("enable_ge") or context.get_context("device_target") == "Ascend" \
                     or context.get_context("device_target") == "GPU":

mindspore/train/loss_scale_manager.py

@@ -15,7 +15,7 @@
 """Loss scale manager abstract class."""
 from __future__ import absolute_import
 
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore import nn
 
 
@@ -58,17 +58,11 @@ class FixedLossScaleManager(LossScaleManager):
 
     Examples:
         >>> import mindspore as ms
-        >>> from mindspore import nn
+        >>> from mindspore import amp, nn
         >>>
         >>> net = Net()
-        >>> #1) Drop the parameter update if there is an overflow
-        >>> loss_scale_manager = ms.FixedLossScaleManager()
-        >>> optim = nn.Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9)
-        >>> model = ms.Model(net, loss_scale_manager=loss_scale_manager, optimizer=optim)
-        >>>
-        >>> #2) Execute parameter update even if overflow occurs
         >>> loss_scale = 1024.0
-        >>> loss_scale_manager = ms.FixedLossScaleManager(loss_scale, False)
+        >>> loss_scale_manager = amp.FixedLossScaleManager(loss_scale, False)
         >>> optim = nn.Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9, loss_scale=loss_scale)
         >>> model = ms.Model(net, loss_scale_manager=loss_scale_manager, optimizer=optim)
     """
@@ -131,12 +125,15 @@ class DynamicLossScaleManager(LossScaleManager):
         scale_factor (int): Coefficient of increase and decrease. Default: 2.
         scale_window (int): Maximum continuous normal steps when there is no overflow. Default: 2000.
 
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
     Examples:
         >>> import mindspore as ms
-        >>> from mindspore import nn
+        >>> from mindspore import amp, nn
         >>>
         >>> net = Net()
-        >>> loss_scale_manager = ms.DynamicLossScaleManager()
+        >>> loss_scale_manager = amp.DynamicLossScaleManager()
         >>> optim = nn.Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9)
         >>> model = ms.Model(net, loss_scale_manager=loss_scale_manager, optimizer=optim)
     """

mindspore/train/metrics/__init__.py

@@ -113,7 +113,7 @@ def get_metric_fn(name, *args, **kwargs):
     Gets the metric method based on the input name.
 
     Args:
-        name (str): The name of metric method. Names can be obtained by `mindspore.train.names` .
+        name (str): The name of metric method. Names can be obtained by :func:`mindspore.train.names` .
             object for the currently supported metrics.
         args: Arguments for the metric function.
         kwargs: Keyword arguments for the metric function.

mindspore/train/metrics/bleu_score.py

@@ -18,7 +18,7 @@ from __future__ import absolute_import
 from collections import Counter
 import numpy as np
 
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.train.metrics.metric import Metric, rearrange_inputs
 
 

mindspore/train/metrics/confusion_matrix.py

@@ -17,7 +17,7 @@ from __future__ import absolute_import
 
 import numpy as np
 
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.train.metrics.metric import Metric, rearrange_inputs
 
 

mindspore/train/metrics/cosine_similarity.py

@@ -17,7 +17,7 @@ from __future__ import absolute_import
 
 import numpy as np
 
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.train.metrics.metric import Metric, rearrange_inputs
 
 

mindspore/train/metrics/dice.py

@@ -17,7 +17,7 @@ from __future__ import absolute_import
 
 import numpy as np
 
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.train.metrics.metric import Metric, rearrange_inputs
 
 
@@ -69,7 +69,7 @@ class Dice(Metric):
     @rearrange_inputs
     def update(self, *inputs):
         r"""
-        Updates the internal evaluation result :math:`y\_pred` and :math:`y`.
+        Updates the internal evaluation result `y_pred` and `y`.
 
         Args:
             inputs (tuple): Input `y_pred` and `y`. `y_pred` and `y` are Tensor, list or numpy.ndarray. `y_pred` is the

mindspore/train/metrics/fbeta.py

@@ -18,7 +18,7 @@ from __future__ import absolute_import
 import sys
 import numpy as np
 
-from mindspore._checkparam import Validator as validator
+from mindspore import _checkparam as validator
 from mindspore.train.metrics.metric import Metric, rearrange_inputs, _check_onehot_data