mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/train/callback/_time_monitor.py

@@ -28,35 +28,65 @@ class TimeMonitor(Callback):
     Args:
         data_size (int): How many steps are the intervals between print information each time.
             if the program get `batch_num` during training, `data_size` will be set to `batch_num`,
-            otherwise `data_size` will be used. Default: None.
+            otherwise `data_size` will be used. Default: ``None`` .
+
+        data_time (bool): Whether to sow the average time of fetching data in Host.
+            Note that data fetch and network compute are processed sequentially in non dataset sink mode, while
+            they are asynchronous in dataset sink mode. Default: ``False`` .
 
     Raises:
         ValueError: If data_size is not positive int.
+        TypeError: If data_time is not bool.
 
     Examples:
-        .. 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/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
         >>>
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
         >>> optim = nn.Momentum(net.trainable_params(), 0.01, 0.9)
         >>> model = Model(net, loss_fn=loss, optimizer=optim)
-        >>> data_path = './MNIST_Data'
-        >>> dataset = create_dataset(data_path)
+        >>> # Create the dataset taking MNIST as an example. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/mnist.py
+        >>> dataset = create_dataset()
         >>> time_monitor = TimeMonitor()
         >>> model.train(10, dataset, callbacks=time_monitor)
     """
 
-    def __init__(self, data_size=None):
+    def __init__(self, data_size=None, data_time=False):
         super(TimeMonitor, self).__init__()
         self.data_size = data_size
         self.epoch_time = time.time()
+        self.data_time = data_time
+        self.data_time_sum = 0.0
+        self.data_time_start = 0.0
+        self.data_sink = lambda c: c.original_args()["dataset_sink_mode"]
+        Validator.check_bool(data_time, "data_time")
+
+    def on_train_step_begin(self, run_context):
+        """
+        Record time at the beginning of step.
+
+        Args:
+            run_context (RunContext): Context of the process running. For more details,
+                    please refer to :class:`mindspore.train.RunContext`.
+        """
+        if self.data_time and not self.data_sink(run_context):
+            interval = time.time() - self.data_time_start
+            self.data_time_sum = self.data_time_sum + interval
+
+    def on_train_step_end(self, run_context):
+        """
+        Record time at the end of step.
+
+        Args:
+            run_context (RunContext): Context of the process running. For more details,
+                    please refer to :class:`mindspore.train.RunContext`.
+        """
+        if self.data_time and not self.data_sink(run_context):
+            self.data_time_start = time.time()
 
     def epoch_begin(self, run_context):
         """
@@ -67,6 +97,9 @@ class TimeMonitor(Callback):
                     please refer to :class:`mindspore.train.RunContext`.
         """
         self.epoch_time = time.time()
+        if self.data_time and not self.data_sink(run_context):
+            self.data_time_sum = 0.0
+            self.data_time_start = time.time()
 
     def epoch_end(self, run_context):
         """
@@ -87,5 +120,22 @@ class TimeMonitor(Callback):
         Validator.check_positive_int(step_size)
 
         step_seconds = epoch_seconds / step_size
-        print("{} epoch time: {:5.3f} ms, per step time: {:5.3f} ms".format
-              (mode.title(), epoch_seconds, step_seconds), flush=True)
+
+        train_log = "{} epoch time: {:5.3f} ms, per step time: {:5.3f} ms".format(
+            mode.title(), epoch_seconds, step_seconds)
+
+        if self.data_time and not self.data_sink(run_context):
+            data_step_seconds = self.data_time_sum * 1000 / step_size
+            data_log = " (data time: {:5.3f} ms)".format(data_step_seconds)
+            train_log += data_log
+        elif self.data_time and self.data_sink(run_context):
+            # send info viewer to query epoch message of cur_epoch_num
+            send_info = cb_params["dataset_helper"].get_send_info(run_context)
+            epoch = cb_params["cur_epoch_num"]
+            epoch_send_info = send_info.epoch(epoch)
+            # show average time of fetching data time
+            fetch_data_time = epoch_send_info["fetch_data_time"]
+            data_log = " (data time: {:5.3f} ms)".format(fetch_data_time)
+            train_log += data_log
+
+        print(train_log, flush=True)

mindspore/train/data_sink.py

@@ -40,7 +40,7 @@ def _init_sink_dataset(dataset, sink_size, input_signature, create_info):
     # create transfer_dataset
     is_info_queue = (create_info and sink_size == 1 and dataset_size != 1 and
                      input_signature is None and not dynamic_shape and
-                     context.get_context('device_target') == 'Ascend' and not context.get_context("enable_ge"))
+                     context.get_context('device_target') == 'Ascend')
     transfer_dataset = _exec_datagraph(dataset, sink_size, create_data_info_queue=is_info_queue)
     dataset.__transfer_dataset__ = transfer_dataset
 
@@ -130,29 +130,23 @@ 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.
+        When using data sinking, the dataset will be automatically sent in a loop, and only the step size of sinking
+        `sink_size` needs to be considered. The default value of `sink_size` is ``1``, which means that all data will
+        be sunk every epoch. If `sink_size` is greater than 1, the amount of data sunk per epoch will be the dataset
+        with a size of `sink_size`.
 
     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
             :class:`mindspore.dataset`, such as :class:`mindspore.dataset.ImageFolderDataset`.
-        sink_size (int): Control the amount of data in each sink. `sink_size` must be positive integer. Default: 1.
+        sink_size (int): Control the amount of data in each sink. `sink_size` must be positive integer. Default: ``1`` .
         jit_config (JitConfig): Controls the execution mode(Graph mode/PyNative mode) of the generated function, and Jit
-            config for compile. Default: None, means running in PyNative mode.
+            config for compile. Default: ``None`` , means running in PyNative mode.
         input_signature (Union[Tensor, List or Tuple of Tensors]): The Tensor which describes the input arguments.
             The shape and dtype of the Tensor will be supplied to this function. If input_signature is specified,
             each input to `fn` must be a `Tensor`. And the input parameters of `fn` cannot accept `**kwargs`. The shape
             and dtype of actual inputs should keep the same as input_signature. Otherwise, TypeError will be raised.
-            Default: None.
+            Default: ``None`` .
 
     Returns:
         Function, the generated function will be executed in data sinking mode.
@@ -175,7 +169,7 @@ def data_sink(fn, dataset, sink_size=1, jit_config=None, input_signature=None):
         ...     out = x + y
         ...     return out
         >>>
-        >>> sink_process = ms.train.data_sink(func_net, dataset, sink_size=1)
+        >>> sink_process = ms.data_sink(func_net, dataset, sink_size=1)
         >>> for _ in range(2):
         ...     out = sink_process()
         ...     print(out)
@@ -212,7 +206,7 @@ def data_sink(fn, dataset, sink_size=1, jit_config=None, input_signature=None):
         real_sink_fun = _get_sink_fun(sink_fun, key_info, is_info_queue, dataset, jit_config)
 
         loop = sink_size
-        if jit_config is not None:
+        if jit_config is not None and context.get_context('mode') == context.GRAPH_MODE:
             loop = 1
 
         out = None

mindspore/train/dataset_helper.py

@@ -18,17 +18,19 @@ from __future__ import absolute_import
 import math
 
 from mindspore import _checkparam as Validator
+from mindspore import log as logger
+from mindspore.common._auto_dynamic import is_auto_dynamic, convert_new_shapes
 from mindspore.common.dtype import pytype_to_dtype
 from mindspore.common.api import _cell_graph_executor
 from mindspore.common._utils import is_shape_unknown
 from mindspore.dataset.engine import offload
-import mindspore.dataset as ds
 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_sched
 from mindspore.ops import operations as P
+from mindspore.common.auto_dynamic_shape import _auto_dynamic_shape
 
 
 def _send_data(dataset, epoch_num):
@@ -47,31 +49,31 @@ def _send_data_no_flag(dataset, epoch_num):
 
 def _dynamic_sink_data(dataset, dataset_iter):
     """Special scenario for dataset with sink_size=1."""
-    _, dataset_shapes = dataset_iter.types_shapes()
     if hasattr(dataset_iter, "sink_size") and \
        dataset_iter.sink_size == 1 and \
        dataset.get_dataset_size() != 1 and \
        hasattr(dataset_iter, "sink_count") and \
-       dataset_iter.sink_count == 1 and \
-       not _has_dynamic_shape(dataset_shapes) and \
-       not ds.config.get_dynamic_shape():
+       dataset_iter.sink_count == 1:
         return True
     return False
 
 
-def _dynamic_sink_exception_scenario(dataset_iter):
+def _dynamic_sink_exception_scenario(dataset_iter, is_dynamic):
     """The exception scenario for dynamic data is not applicable."""
-    _, dataset_shapes = dataset_iter.types_shapes()
-
-    if _has_dynamic_shape(dataset_shapes) or context.get_context("mode") != context.GRAPH_MODE:
+    if context.get_context("mode") != context.GRAPH_MODE or is_dynamic:
         return True
     return False
 
 
-def _dynamic_sink_scenario(dataset, dataset_iter):
+def _dynamic_sink_scenario(dataset, dataset_iter, is_dynamic):
     """Special scenario with dynamic shape and sink_size=1."""
     flag = False
-    if _dynamic_sink_data(dataset, dataset_iter) and not _dynamic_sink_exception_scenario(dataset_iter):
+
+    # This is used only for test
+    if is_auto_dynamic():
+        return False
+
+    if _dynamic_sink_data(dataset, dataset_iter) and not _dynamic_sink_exception_scenario(dataset_iter, is_dynamic):
         flag = True
 
     return flag
@@ -93,8 +95,7 @@ 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
+        self._get_attr_from_cell(network)
 
     def construct(self):
         outputs = self.get_next()
@@ -120,11 +121,18 @@ def _generate_network_with_dataset(network, dataset_helper, queue_name):
     Generate new network with network and dataset info.
     """
     dataset_types, dataset_shapes = dataset_helper.types_shapes()
+
+    # This is used only for test
+    if is_auto_dynamic():
+        new_shapes = convert_new_shapes(dataset_shapes)
+        return _generate_dataset_sink_mode_net(network, new_shapes, dataset_types, queue_name)
+
     if network.get_inputs() and None not in network.get_inputs():
         _check_inputs(network.get_inputs(), dataset_shapes, dataset_types)
     elif context.get_context("mode") == context.PYNATIVE_MODE:
         dataset_shapes = tuple([(-2,)] * len(dataset_shapes))
-    network = _generate_dataset_sink_mode_net(network, dataset_shapes, dataset_types, queue_name)
+    network = _generate_dataset_sink_mode_net(
+        network, dataset_shapes, dataset_types, queue_name)
     return network
 
 
@@ -172,29 +180,20 @@ def _get_dataset_aux(dataset):
 
 def connect_network_with_dataset(network, dataset_helper):
     """
-    Connect the `network` with dataset in `dataset_helper`.
-
-    This function wraps the input network with 'GetNext' so that the data can be fetched automatically from the
-    data channel corresponding to the 'queue_name' and passed to the input network during forward computation.
-
-    Note:
-        In the case of running the network on Ascend/GPU in graph mode, this function will wrap the input network with
-        :class:`mindspore.ops.GetNext`. In other cases, the input network will be returned with no change.
-        The :class:`mindspore.ops.GetNext` is required to get data only in sink mode,
-        so this function is not applicable to no-sink mode.
-        when dataset_helper's dataset_sink_mode is True, it can only be connected to one network.
+    Connect the `network` with dataset in `dataset_helper`. Only supported in `sink mode
+    <https://mindspore.cn/tutorials/experts/en/r2.2/optimize/execution_opt.html>`_, (dataset_sink_mode=True).
 
     Args:
         network (Cell): The training network for dataset.
         dataset_helper (DatasetHelper): A class to process the MindData dataset, it provides the type, shape and queue
-            name of the dataset to wrap the `GetNext`.
+            name of the dataset.
 
     Returns:
-        Cell, a new network wrapped with 'GetNext' in the case of running the task on Ascend in graph mode, otherwise
-        it is the input network.
+        Cell, a new network containing the type, shape and queue name of the dataset info.
 
     Raises:
         RuntimeError: If the API was not called in dataset sink mode.
+
     Supported Platforms:
         ``Ascend`` ``GPU``
 
@@ -208,7 +207,7 @@ def connect_network_with_dataset(network, dataset_helper):
         >>> train_dataset = ds.NumpySlicesDataset(data=data).batch(32)
         >>> dataset_helper = ms.DatasetHelper(train_dataset, dataset_sink_mode=True)
         >>> net = nn.Dense(10, 5)
-        >>> net_with_get_next = ms.connect_network_with_dataset(net, dataset_helper)
+        >>> net_with_dataset = ms.connect_network_with_dataset(net, dataset_helper)
     """
     dataset_iter = dataset_helper.iter
     dataset = dataset_iter.dataset
@@ -219,6 +218,7 @@ def connect_network_with_dataset(network, dataset_helper):
             "The API 'connect_network_with_dataset' should be called in dataset sink mode.")
 
     if _is_role_sched():
+        network.add_flags(sink_mode=True)
         return network
 
     if not hasattr(aux, '__network__'):
@@ -227,13 +227,15 @@ def connect_network_with_dataset(network, dataset_helper):
     if aux.__network__ is not network:
         raise ValueError(
             "The dataset has been connected to other network, please check the code.")
-
+    is_dynamic = bool(network.get_inputs())
     queue_name = dataset.__transfer_dataset__.queue_name
-    if _dynamic_sink_scenario(dataset, dataset_iter) and not context.get_context("enable_ge"):
+    if _dynamic_sink_scenario(dataset, dataset_iter, is_dynamic):
         dataset_types, dataset_shapes = dataset_helper.get_data_info()
         dataset_types = [pytype_to_dtype(x) for x in dataset_types]
-
+        if not is_dynamic:
+            dataset_shapes = _auto_dynamic_shape.auto_dynamic_generate_compile_args(dataset_shapes, True)
         key = str(dataset_types) + str(dataset_shapes)
+        _auto_dynamic_shape.update_phase_and_compile_args(dataset_shapes, key, True, aux)
         if hasattr(aux, '__network_manage__') and key in aux.__network_manage__:
             network = aux.__network_manage__[key]
         else:
@@ -248,21 +250,22 @@ def connect_network_with_dataset(network, dataset_helper):
             else:
                 aux.__network_manage__ = dict()
             aux.__network_manage__[key] = network
+        network.add_flags(sink_mode=True)
         return network
 
     if hasattr(aux, '__sink_network__'):
         network = aux.__sink_network__
     else:
-        if not context.get_context("enable_ge") and context.get_context("device_target") in ("Ascend", "GPU"):
+        if context.get_context("device_target") in ("Ascend", "GPU"):
             network = offload.check_add_offload_sink_mode(
                 dataset, dataset_helper, network)
             network = _generate_network_with_dataset(
                 network, dataset_helper, queue_name)
             aux.__sink_network__ = network
 
-    if _dynamic_sink_data(dataset, dataset_iter) and _dynamic_sink_exception_scenario(dataset_iter):
+    if _dynamic_sink_data(dataset, dataset_iter) and _dynamic_sink_exception_scenario(dataset_iter, is_dynamic):
         dataset_helper.get_data_info()
-
+    network.add_flags(sink_mode=True)
     return network
 
 
@@ -281,7 +284,7 @@ class DatasetHelper:
                            :class:`mindspore.dataset`, such as :class:`mindspore.dataset.ImageFolderDataset`.
         dataset_sink_mode (bool): If the value is True, GetNext is employed to fetch the data at device through the
                                   dataset pipeline, otherwise fetch the data at host by iterating through the dataset.
-                                  Default: True.
+                                  Default: ``True``.
         sink_size (int): Control the amount of data in each sink.
                           If sink_size=-1, sink the complete dataset for each epoch.
                           If sink_size>0, sink sink_size data for each epoch.
@@ -303,7 +306,7 @@ class DatasetHelper:
         >>> for next_element in set_helper:
         ...     # `next_element` includes data and label, using data to run the net
         ...     data = next_element[0]
-        ...     net(data)
+        ...     result = net(data)
     """
 
     def __init__(self, dataset, dataset_sink_mode=True, sink_size=-1, epoch_num=1):
@@ -316,22 +319,19 @@ class DatasetHelper:
             sink_size = dataset.get_dataset_size()
 
         if dataset_sink_mode:
-            if context.get_context("enable_ge"):
-                iterclass = _DatasetIterGE
-            else:
-                if context.get_context("mode") == context.GRAPH_MODE:
-                    if _is_role_sched():
-                        iterclass = _DatasetIterPSServer
-                    elif (context.get_context("device_target") == "Ascend") or \
-                         (context.get_context("device_target") == "GPU"):
-                        iterclass = _DatasetIterMSLoopSink
-                    else:
-                        target = context.get_context("device_target")
-                        raise RuntimeError("Currently dataset sink mode is not supported when the device "
-                                           "target is {}, please set dataset_sink_mode to False "
-                                           "in Model.train()".format(target))
+            if context.get_context("mode") == context.GRAPH_MODE:
+                if _is_role_sched():
+                    iterclass = _DatasetIterPSServer
+                elif (context.get_context("device_target") == "Ascend") or \
+                        (context.get_context("device_target") == "GPU"):
+                    iterclass = _DatasetIterMSLoopSink
                 else:
-                    iterclass = _DatasetIterPyNative
+                    target = context.get_context("device_target")
+                    raise RuntimeError("Currently dataset sink mode is not supported when the device "
+                                       "target is {}, please set dataset_sink_mode to False "
+                                       "in Model.train()".format(target))
+            else:
+                iterclass = _DatasetIterPyNative
             self.iter = iterclass(dataset, sink_size, epoch_num)
         else:
             iterclass = _DatasetIterNormal
@@ -383,28 +383,108 @@ class DatasetHelper:
         return self.iter.get_sink_size()
 
     def stop_send(self):
-        """Stop send data about data sink."""
+        """
+        Stop send data about data sink.
+
+        Examples:
+            >>> import mindspore as ms
+            >>> import numpy as np
+            >>> # Define a dataset pipeline
+            >>> def generator():
+            ...    for i in range(5):
+            ...        yield (np.ones((32, 10)),)
+            >>> train_dataset = ms.dataset.GeneratorDataset(generator, ["data"])
+            >>> dataset_helper = ms.DatasetHelper(train_dataset, dataset_sink_mode=True, sink_size=-1)
+            >>> dataset_helper.stop_send()
+        """
         self.iter.stop_send()
 
     def release(self):
-        """Free up resources about data sink."""
+        """
+        Free up resources about data sink.
+
+        Examples:
+            >>> import numpy as np
+            >>> import mindspore as ms
+            >>> from mindspore import nn
+            >>> from mindspore import dataset as ds
+            >>>
+            >>> data = {"x": np.float32(np.random.rand(64, 10)), "y": np.random.randint(0, 5, (64,))}
+            >>> train_dataset = ds.NumpySlicesDataset(data=data).batch(32)
+            >>> dataset_helper = ms.DatasetHelper(train_dataset, dataset_sink_mode=True)
+            >>> dataset_helper.release()
+        """
         self.iter.release()
 
     def continue_send(self):
-        """Continue to send data to device at the beginning of epoch."""
+        """
+        Continue to send data to device at the beginning of epoch.
+
+        Examples:
+            >>> import numpy as np
+            >>> import mindspore as ms
+            >>> from mindspore import nn
+            >>> from mindspore import dataset as ds
+            >>>
+            >>> data = {"x": np.float32(np.random.rand(64, 10)), "y": np.random.randint(0, 5, (64,))}
+            >>> train_dataset = ds.NumpySlicesDataset(data=data).batch(32)
+            >>> dataset_helper = ms.DatasetHelper(train_dataset, dataset_sink_mode=True)
+            >>> dataset_helper.continue_send()
+        """
         self.iter.continue_send()
 
-    def _reset(self, step, epoch):
+    def _reset(self, step, dataset_size):
         """Reset the dataset to the provided step and epoch."""
-        self.iter._reset(step, epoch)  # pylint: disable=protected-access
+        self.iter._reset(step, dataset_size)  # pylint: disable=protected-access
 
+    # pylint: disable=missing-docstring
     def get_data_info(self):
-        """
-        In sink mode, it returns the types and shapes of the current data.
-        Generally, it works in dynamic shape scenarios.
-        """
+        # In sink mode, it returns the types and shapes of the current data.
+        # Generally, it works in dynamic shape scenarios.
         return self.iter.get_data_info()
 
+    # pylint: disable=missing-docstring
+    def get_send_info(self, run_context):
+        # In sink mode, it returns the send information of dataset at this moment.
+        # Send information includes number of send batches, time summary of fetching data on host
+        # and time summary of sending data.
+        class InfoViewer:
+            '''
+            Inner class for parsing send info.
+            '''
+            def __init__(self, send_info, run_context):
+                self.info_ = {}
+                self.sink_size = run_context.original_args()["batch_num"]
+                if run_context.original_args().get("train_dataset", None) is not None:
+                    self.dataset_size = run_context.original_args()["train_dataset"].get_dataset_size()
+                elif run_context.original_args().get("valid_dataset", None) is not None:
+                    self.dataset_size = run_context.original_args()["valid_dataset"].get_dataset_size()
+                else:
+                    raise RuntimeError("Could not find a proper dataset to estimate dataset size.")
+                if not send_info:
+                    epoch = 1
+                    self.info_[epoch] = {'fetch_data_num': 0, 'fetch_data_time': 0, 'first_data_time': 0}
+                else:
+                    for info_per_epoch in send_info:
+                        epoch, fetch_data_num, first_data_time, fetch_data_time = info_per_epoch
+                        if fetch_data_num > 1:
+                            fetch_data_time = (fetch_data_time - first_data_time) / (fetch_data_num - 1) * 1000.
+                        self.info_[epoch] = {'fetch_data_num': fetch_data_num,
+                                             'fetch_data_time': fetch_data_time,
+                                             'first_data_time': first_data_time}
+
+            def epoch(self, epoch):
+                if self.sink_size == self.dataset_size:
+                    return self.info_[epoch]
+                global_step = epoch * self.sink_size
+                data_epoch = math.ceil(global_step / self.dataset_size)
+                return self.info_[data_epoch]
+
+        # send info struct:[epoch, data_num_per_epoch, first_data_time, accumulate_data_time]
+        # for example [1, 1875, 0.421, 0.362]
+        send_info = self.iter.get_send_info()
+        return InfoViewer(send_info, run_context)
+
 
 class _DatasetIter:
     """Base iter for dataset helper"""
@@ -415,14 +495,20 @@ class _DatasetIter:
         self.sink_count = self.get_sink_count(dataset)
         self.dataset_types, self.dataset_shapes = _get_types_and_shapes(
             dataset)
-        self.dynamic_shape = _has_dynamic_shape(self.dataset_shapes) or ds.config.get_dynamic_shape()
-        if self.dynamic_shape:
-            ds.config.set_dynamic_shape(True)
+
+        if dataset.get_init_step() % sink_size != 0:
+            init_epoch = dataset.get_init_step() // sink_size
+            init_step = init_epoch * sink_size
+            logger.warning("Init global step must be the end of the epoch in sink mode, "
+                           "but got: {0}. Reset it to the end of epoch {1} at step {2}."
+                           .format(dataset.get_init_step(), init_epoch, init_step))
+            dataset.set_init_step(init_step)
+
         if not hasattr(dataset, '__transfer_dataset__'):
             if hasattr(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)
+            create_data_info_queue = (
+                sink_size == 1 and self.sink_count == 1 and dataset.get_dataset_size() != 1)
             dataset.__transfer_dataset__ = _exec_datagraph(dataset, self.sink_size,
                                                            create_data_info_queue=create_data_info_queue)
 
@@ -439,6 +525,7 @@ class _DatasetIter:
         self.release = dataset.__transfer_dataset__.release
         self.continue_send = dataset.__transfer_dataset__.continue_send
         self.get_data_info = dataset.__transfer_dataset__.get_data_info
+        self.get_send_info = dataset.__transfer_dataset__.get_send_info
         if hasattr(dataset.__transfer_dataset__, "_reset"):
             self._reset = dataset.__transfer_dataset__._reset  # pylint: disable=protected-access
 
@@ -475,8 +562,7 @@ class _DatasetIter:
         if hasattr(self.dataset, '__loop_size__'):
             sink_size = self.dataset.__loop_size__
         else:
-            if context.get_context("enable_ge") or context.get_context("device_target") == "Ascend" \
-                    or context.get_context("device_target") == "GPU":
+            if context.get_context("device_target") == "Ascend" or context.get_context("device_target") == "GPU":
                 if self.sink_size > 0:
                     sink_size = self.sink_size
                 else:
@@ -484,24 +570,6 @@ class _DatasetIter:
         return sink_size
 
 
-class _DatasetIterGE(_DatasetIter):
-    """Iter for GE."""
-
-    def __init__(self, dataset, sink_size, epoch_num):
-        super().__init__(dataset, sink_size, epoch_num)
-        self.sink_count = self.get_sink_count(dataset)
-        batch_expand_num = 1
-        if _need_to_full():
-            batch_expand_num = _get_device_num() // _get_pipeline_stages()
-        tensor_list_run = _construct_tensor_list(
-            self.dataset_types, self.dataset_shapes, batch_expand_num)
-
-        def op():
-            return tensor_list_run
-
-        self.op = op
-
-
 class _DatasetIterPyNative(_DatasetIter):
     """Iter for context (mode=PYNATIVE_MODE)."""
 

mindspore/train/loss_scale_manager.py

@@ -51,16 +51,19 @@ class FixedLossScaleManager(LossScaleManager):
     inherits from :class:`mindspore.amp.LossScaleManager`.
 
     Args:
-        loss_scale (float): Magnification factor of gradients. Note that if `drop_overflow_update` is set to False,
-            the value of `loss_scale` in optimizer should be set to the same as here. Default: 128.0.
-        drop_overflow_update (bool): Whether to execute optimizer if there is an overflow. If True, the optimizer will
-            not executed when overflow occurs. Default: True.
+        loss_scale (float): Magnification factor of gradients. Note that if `drop_overflow_update` is set to ``False`` ,
+            the value of `loss_scale` in optimizer should be set to the same as here. Default: ``128.0`` .
+        drop_overflow_update (bool): Whether to execute optimizer if there is an overflow.
+            If ``True`` , the optimizer will
+            not executed when overflow occurs. Default: ``True`` .
 
     Examples:
         >>> import mindspore as ms
         >>> from mindspore import amp, nn
         >>>
-        >>> net = Net()
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
         >>> loss_scale = 1024.0
         >>> 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)
@@ -121,9 +124,9 @@ class DynamicLossScaleManager(LossScaleManager):
     adjusted, inherits from :class:`mindspore.amp.LossScaleManager`.
 
     Args:
-        init_loss_scale (float): Initialize loss scale. Default: 2**24.
-        scale_factor (int): Coefficient of increase and decrease. Default: 2.
-        scale_window (int): Maximum continuous normal steps when there is no overflow. Default: 2000.
+        init_loss_scale (float): Initialize loss scale. Default: ``2 ** 24`` .
+        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``
@@ -132,7 +135,9 @@ class DynamicLossScaleManager(LossScaleManager):
         >>> import mindspore as ms
         >>> from mindspore import amp, nn
         >>>
-        >>> net = Net()
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
         >>> 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)