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

mindspore/train/callback/_tft_register.py

@@ -0,0 +1,352 @@
+# Copyright 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.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Checkpoint related classes and functions."""
+
+import os
+from mindspore.train.serialization import save_checkpoint
+from mindspore.parallel._utils import _get_device_num
+from mindspore import _checkparam as Validator
+from mindspore.train.callback._callback import Callback
+from mindspore import context
+from mindspore.common.parameter import Parameter
+from mindspore.communication import get_rank, get_group_size
+from mindspore import log as logger
+from mindspore.train.serialization import _get_cur_rank_dp
+from mindspore._c_expression import _repair_device, _stop_device, _tft_sem_post
+from mindspore._c_expression import clean_tdt_channel
+from mindspore._c_expression import send_recv
+from mindspore._c_expression import CollectiveManager
+from mindspore._c_expression import _get_uce_process_strategy, _get_uce_mem_info
+
+def _get_ckpt_dir(step, ckpt_save_path, is_tmp_file):
+    """ Common func to generate ckpt dir name."""
+    tmp = "_tmp" if is_tmp_file else ""
+    mid_dir = f"tft_saved_checkpoints-step_{str(step)}{tmp}"
+    return os.path.join(ckpt_save_path, mid_dir)
+
+def _save_checkpoint_on_failure(step, save_info, args, cb_ctx):
+    """ Callback used for TFT save ckpt function when errors occur."""
+    logger.info("Enter _save_checkpoint_on_failure function")
+    ckpt_save_path = cb_ctx.ckpt_save_path
+    cb_params = args
+    cur_rank = get_rank()
+    cur_step_num = cb_params.cur_step_num
+    cur_epoch_num = cb_params.cur_epoch_num
+    batch_num = cb_params.batch_num
+    if cur_step_num > step:
+        cur_epoch_num = (step - 1) // batch_num + 1
+    step_num_in_epoch = int((step - 1) % batch_num + 1)
+
+    append_dict = {}
+    append_dict["epoch_num"] = cur_epoch_num
+    append_dict["step_num"] = step
+    append_dict["cur_rank"] = cur_rank
+    append_dict["batch_num"] = batch_num
+    append_dict["__exception_save__"] = True
+
+    append_dict["global_step"] = Parameter([cb_ctx.global_step])
+    outputs = cb_params.net_outputs
+    if isinstance(outputs, (tuple, list)) and len(outputs) >= 3:
+        append_dict["loss_scale"] = outputs[2]
+
+    ckpt_file = f"ttp_rank_{str(cur_rank)}-{str(cur_epoch_num)}_{str(step_num_in_epoch)}.ckpt"
+    cur_ckpt_dir = _get_ckpt_dir(step, ckpt_save_path, True) + "/rank_" + str(cur_rank)
+    os.makedirs(cur_ckpt_dir, exist_ok=True)
+    cur_file = os.path.join(cur_ckpt_dir, ckpt_file)
+    save_checkpoint(cb_params.train_network, cur_file,
+                    integrated_save=False, append_dict=append_dict)
+    logger.info("Finish _save_checkpoint_on_failure function")
+
+def _rename_save_result(step, cb_ctx):
+    """ Callback used for TFT rename function after ckpt save callback was finished and successful."""
+    logger.info("Enter _rename_save_result function")
+    tmp_dir = _get_ckpt_dir(step, cb_ctx.ckpt_save_path, True)
+    fin_dir = _get_ckpt_dir(step, cb_ctx.ckpt_save_path, False)
+
+    os.rename(tmp_dir, fin_dir)
+    logger.info("Finish _rename_save_result function")
+
+def _tft_exit_cb(ctx):
+    logger.error("Enter mindio ttp exit process, which means other ranks occur exception, check other ranks' logs!")
+    _tft_sem_post()
+    os._exit(1)   # pylint: disable=W0212
+
+
+
+def _tft_repair_callback(step, need_rebuild, error_ranks, repair_info, args, cb_ctx):
+    """ Callback used for TFT repair function."""
+    logger.info("Enter _tft_repair_callback repair type: {}".format(repair_info["repair_type"]))
+    if(repair_info["repair_type"] == cb_ctx.tft.RepairType.RT_UCE_HIGHLEVEL.value\
+or repair_info["repair_type"] == cb_ctx.tft.RepairType.RT_UCE_LOWLEVEL.value):
+        logger.info("Enter _tft_repair_callback uce REPARI_DEVICE device_id : {}".format(cb_ctx.device_id))
+        _repair_device(cb_ctx.device_id)
+
+    if(repair_info["repair_type"] == cb_ctx.tft.RepairType.RT_UCE_HIGHLEVEL.value\
+       or repair_info["repair_type"] == cb_ctx.tft.RepairType.RT_SEND.value):
+        logger.info("Enter _tft_repair_callback SEND_RECV repair type: \
+{}, src_rank:{}, dst_rank: {}".format(repair_info["repair_type"], repair_info["src"], repair_info["dst"]))
+        cb_params = args
+        src_rank = repair_info["src"][0]
+        dst_rank = repair_info["dst"][0]
+        send_recv(cb_params.network.trainable_params(), src_rank, dst_rank)
+    logger.info("Finish _tft_repair_callback")
+
+
+def _tft_clean_callback(is_uce_error, ctx):
+    """ Callback used for TFT clean function."""
+    logger.info("Enter _tft_clean_callback")
+    ret = 0
+    if is_uce_error:
+        _get_uce_mem_info(ctx.device_id)
+        err_strategy = _get_uce_process_strategy()
+        logger.info("_tft_clean_callback err_strategy: {}".format(err_strategy))
+        if err_strategy == "RS_UCE_HIGHLEVEL":
+            ret = 0
+        elif err_strategy == "RS_UCE_LOWLEVEL":
+            ret = 2
+        else:
+            ret = 1
+    clean_tdt_channel()
+    logger.info("Enter _tft_clean_callback resume_hccl_comm")
+    CollectiveManager.get_instance().resume_hccl_comm()
+    logger.info("Finish _tft_clean_callback, ret: {}".format(ret))
+    return ret
+
+
+def _tft_stop_callback(cb_ctx):
+    """ Callback used for TFT stop function."""
+    logger.info("Enter _tft_stop_callback device_id: {}".format(cb_ctx.device_id))
+    _stop_device(cb_ctx.device_id)
+    logger.info("Finish _tft_stop_callback")
+
+
+class TFTRegister(Callback):
+    """
+    This callback is used to enable the TFT feature
+    `MindIO TFT <https://www.hiascend.com/document/detail/zh/mindx-dl/60rc2/mindio/mindiottp/mindiottp001.html>`_.
+    This callback will execute TFT operations during training process, such as TFT init, report and exception handle.
+
+    Note:
+        Required for Ascend graph mode only. And sink size must be less than or equal to 1.
+
+    Args:
+        ctrl_rank_id (int): TFT controller's running rank_id, used for init TFT controller.
+        ctrl_ip (str): TFT controller's ip address, used for init TFT controller.
+        ctrl_port (int): TFT controller's ip port, used for init TFT controller and processor.
+        ckpt_save_path (str): Checkpoint save directory when failure occurs, checkpoint file will save to directory
+           named ttp_saved_checkpoints-step_{cur_step_num} under this directory.
+
+    Raises:
+        Exception: TFT init failed.
+        ModuleNotFoundError: Mindio TFT whl package is not installed.
+
+    Examples:
+        >>> import numpy as np
+        >>> import os
+        >>> import math
+        >>> import mindspore as ms
+        >>> import mindspore.dataset as ds
+        >>> from mindspore import nn, ops, Parameter, train
+        >>> from mindspore.communication import init
+        >>> from mindspore.common.initializer import initializer, HeUniform
+        >>> from mindspore.train import Model, TFTRegister
+        >>> from mindspore import dataset as ds
+        >>> ms.set_context(mode=ms.GRAPH_MODE, jit_level='O2')
+        >>> ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.SEMI_AUTO_PARALLEL, pipeline_stages=2)
+        >>> init()
+        >>> ms.set_seed(1)
+        >>> ms.set_auto_parallel_context(strategy_ckpt_config={"save_file":
+        >>>                             "./src_pipeline_strategys/src_strategy_{}.ckpt".format(get_rank())})
+        >>> class MatMulCell(nn.Cell):
+        ...     def __init__(self, param=None, shape=None):
+        ...         super().__init__()
+        ...         if shape is None:
+        ...             shape = [28 * 28, 512]
+        ...         weight_init = HeUniform(math.sqrt(5))
+        ...         self.param = Parameter(initializer(weight_init, shape), name="param")
+        ...         if param is not None:
+        ...             self.param = param
+        ...         self.print = ops.Print()
+        ...         self.matmul = ops.MatMul()
+        ...
+        ...     def construct(self, x):
+        ...         out = self.matmul(x, self.param)
+        ...         self.print("out is:", out)
+        ...         return out
+        >>>
+        >>> class Network(nn.Cell):
+        ...     def __init__(self):
+        ...         super().__init__()
+        ...         self.flatten = nn.Flatten()
+        ...         self.layer1 = MatMulCell()
+        ...         self.relu1 = nn.ReLU()
+        ...         self.layer2 = nn.Dense(512, 512)
+        ...         self.relu2 = nn.ReLU()
+        ...         self.layer3 = nn.Dense(512, 10)
+        ...
+        ...     def construct(self, x):
+        ...         x = self.flatten(x)
+        ...         x = self.layer1(x)
+        ...         x = self.relu1(x)
+        ...         x = self.layer2(x)
+        ...         x = self.relu2(x)
+        ...         logits = self.layer3(x)
+        ...         return logits
+        >>>
+        >>> net = Network()
+        >>> net.layer1.pipeline_stage = 0
+        >>> net.relu1.pipeline_stage = 0
+        >>> net.layer2.pipeline_stage = 0
+        >>> net.relu2.pipeline_stage = 1
+        >>> net.layer3.pipeline_stage = 1
+        >>>
+        >>> def create_dataset(batch_size):
+        ...     dataset_path = os.getenv("DATA_PATH")
+        ...     dataset = ds.MnistDataset(dataset_path)
+        ...     image_transforms = [
+        ...         ds.vision.Rescale(1.0 / 255.0, 0),
+        ...         ds.vision.Normalize(mean=(0.1307,), std=(0.3081,)),
+        ...         ds.vision.HWC2CHW()
+        ...     ]
+        ...     label_transform = ds.transforms.TypeCast(ms.int32)
+        ...     dataset = dataset.map(image_transforms, 'image')
+        ...     dataset = dataset.map(label_transform, 'label')
+        ...     dataset = dataset.batch(batch_size)
+        ...     return dataset
+        >>>
+        >>> data_set = create_dataset(32)
+        >>>
+        >>> optimizer = nn.SGD(net.trainable_params(), 1e-2)
+        >>> optimizer_wrapper = nn.OptTFTWrapper(optimizer)
+        >>> loss_fn = nn.CrossEntropyLoss()
+        >>>
+        >>> net_with_loss = nn.PipelineCell(nn.WithLossCell(net, loss_fn), 4)
+        >>> net_with_loss.set_train()
+        >>> model = Model(net_with_loss, optimizer=optimizer)
+        >>> tft_cb = TFTRegister("192.168.0.1", 2000, "./tft_checkpoint/")
+        >>> loss_cb = train.LossMonitor(1)
+        >>> model.train(1, dataset, callbacks=[tft_cb, loss_cb])
+    """
+
+    def __init__(self, ctrl_rank_id, ctrl_ip, ctrl_port, ckpt_save_path):
+        super(TFTRegister, self).__init__()
+
+        tft_env = os.getenv("MS_ENABLE_TFT", "")
+        if ("TTP:1" not in tft_env) and ("UCE:1" not in tft_env):
+            raise ValueError("MindIO TFT regitster need custom switch on[MS_ENABLE_TFT='{TTP:1,UCE:1}']!")
+        mode = context.get_context("mode")
+        device_target = context.get_context("device_target")
+        if device_target != "Ascend" or mode != context.GRAPH_MODE:
+            raise ValueError("MindIO adataper only support on Ascend device with GRAPH Mode!")
+
+        # let it raise errors if not install mindio_tft package
+        from mindio_ttp import framework_ttp as tft
+        self.tft = tft
+        self.global_step = 0
+        Validator.check_non_negative_int(ctrl_port)
+        self.has_init_replica = False
+        self._controller_ip = ctrl_ip
+        self._controller_rank_id = ctrl_rank_id
+        self._controller_port = ctrl_port
+        self.device_id = context.get_context("device_id")
+        self._init_tft()
+        self.ckpt_save_path = ckpt_save_path
+
+    def _set_tft_optimizer_replica(self, run_context):
+        """ set Mindio TFT optimizer replica info, used internal. """
+        cur_rank = get_rank()
+        cb_params = run_context.original_args()
+        train_network = cb_params.train_network
+        # in data_parallel mode, every ranks has same train parameters
+        if context.get_auto_parallel_context("parallel_mode") == "data_parallel":
+            group_size = get_group_size()
+            dp = tuple(range(group_size))
+        else:
+            param_layout_dict = train_network.parameter_layout_dict
+            dp = _get_cur_rank_dp(param_layout_dict) if param_layout_dict else _get_cur_rank_dp(train_network)
+        logger.warning(f"Set TFT replica with dp: {dp}.")
+        replica_info = [
+            {
+                "type": 1,
+                "rank_list": dp,
+                "replica_cnt": len(dp),
+                "replica_shift": 0
+            }
+        ]
+        self.tft.tft_set_optimizer_replica(cur_rank, replica_info)
+
+    def _init_tft(self):
+        """ Init Mindio TFT, used internal. """
+        logger.info("Begin to init tft.")
+        self.tft.tft_register_save_ckpt_handler(_save_checkpoint_on_failure, self)
+        self.tft.tft_register_rename_handler(_rename_save_result, self)
+        self.tft.tft_register_exit_handler(_tft_exit_cb, self)
+        self.tft.tft_register_stop_handler(_tft_stop_callback, self)
+        self.tft.tft_register_clean_handler(_tft_clean_callback, self)
+        self.tft.tft_register_repair_handler(_tft_repair_callback, self)
+
+        world_size = _get_device_num()
+        cur_rank = get_rank()
+        enable_local_copy = False
+        enable_arf = False
+        enable_zit = False
+        enable_tls = False
+        tls_key_dir = ""
+
+        if cur_rank == self._controller_rank_id:
+            logger.info(f"Begin to start tft controller on rank_id:{cur_rank}")
+            self.tft.tft_init_controller(cur_rank, world_size, enable_local_copy, enable_arf, enable_zit)
+            self.tft.tft_start_controller(self._controller_ip, self._controller_port, enable_tls, tls_key_dir)
+            logger.info("Finish start tft controller.")
+
+        logger.info("Begin to start tft processor.")
+        self.tft.tft_init_processor(cur_rank, world_size, enable_local_copy, enable_tls, tls_key_dir)
+        self.tft.tft_start_processor(self._controller_ip, self._controller_port)
+        logger.info("Finished start tft processor.")
+
+    def on_train_step_end(self, run_context):
+        """
+        And report status to MindIO TFT after every step finished.
+
+        Args:
+            run_context (RunContext): Context of the train running. Refer to
+                                      :class:`mindspore.train.RunContext` for detail.
+        """
+        if self.has_init_replica is False:
+            self.has_init_replica = True
+            self._set_tft_optimizer_replica(run_context)
+        cb_params = run_context.original_args()
+        if cb_params.optimizer is not None:
+            self.global_step = int(cb_params.optimizer.global_step.data)
+        else:
+            self.global_step = int(cb_params.network.optimizer.global_step.data)
+        logger.info("START Set optimizer finish step status to TFT. step: {}".format(cb_params.cur_step_num))
+        self.tft.tft_end_updating_os(cb_params.cur_step_num)
+        logger.info("END Set optimizer finish step status to TFT.")
+
+
+    def on_train_begin(self, run_context):
+        cb_params = run_context.original_args()
+        sink_size = cb_params.get("sink_size", 0)
+        if sink_size > 1:
+            raise ValueError("TFT feature doesn't support sink_size > 1.")
+        logger.info("Set set args to TFT.")
+        self.tft.tft_set_step_args(cb_params)
+
+    def end(self, run_context):
+        cur_rank = get_rank()
+        if cur_rank == self._controller_rank_id:
+            self.tft.tft_destroy_controller()
+        self.tft.tft_destroy_processor()

mindspore/train/callback/_time_monitor.py

@@ -30,7 +30,7 @@ class TimeMonitor(Callback):
             if the program get `batch_num` during training, `data_size` will be set to `batch_num`,
             otherwise `data_size` will be used. Default: ``None`` .
 
-        data_time (bool): Whether to sow the average time of fetching data in Host.
+        data_time (bool): Whether to show 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`` .
 
@@ -43,13 +43,13 @@ class TimeMonitor(Callback):
         >>> 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
+        >>> # https://gitee.com/mindspore/docs/blob/master/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)
         >>> # 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()
         >>> time_monitor = TimeMonitor()
         >>> model.train(10, dataset, callbacks=time_monitor)

mindspore/train/data_sink.py

@@ -130,10 +130,11 @@ 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 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`.
+        When using data sinking, the dataset will be automatically looped to the device. The device side can cache up
+        to 100 batches of data and occupy no more than 2GB of memory. At this time, only the number of steps for each
+        sinking `sink_size` needs to be considered. `sink_size` defaults to ``1``, indicating that each epoch only
+        takes one batch of data from the cache for training and outputs a loss. If `sink_size` is greater than 1, each
+        epoch takes out `sink_size` batches of data from the cache for training and outputs a loss.
 
     Args:
         fn (Function): The Python function that will be run with dataset.
@@ -145,7 +146,7 @@ def data_sink(fn, dataset, sink_size=1, jit_config=None, input_signature=None):
         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.
+            and dtype of actual inputs should keep the same as `input_signature`. Otherwise, TypeError will be raised.
             Default: ``None`` .
 
     Returns:

mindspore/train/dataset_helper.py

@@ -16,18 +16,20 @@
 from __future__ import absolute_import
 
 import math
+import copy
 
 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.api import _cell_graph_executor, _is_args_fullmode, ARG_SPECIFIED
 from mindspore.common._utils import is_shape_unknown
 from mindspore.dataset.engine import offload
 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
+    _to_full_shapes, _get_pipeline_stages, _change_symbols_for_parallel, _is_in_auto_parallel_mode, \
+    _origin_shapes, _dynamic_shape_for_dataset
 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
@@ -95,6 +97,20 @@ class _DataWrapper(nn.Cell):
         self.add_flags(**flags)
         self.get_next = P.GetNext(
             dataset_types, dataset_shapes, len(dataset_types), queue_name)
+        if network.get_inputs() is not None:
+            network_inputs = network.get_inputs()
+            is_fullmode = _is_args_fullmode(network_inputs, False)
+            if is_fullmode:
+                symbol_inputs = [getattr(inp, "symbolic_shape", None) for inp in network.get_inputs()]
+            else:
+                symbol_inputs = [None for _ in dataset_shapes]
+                arg_specified = network_inputs.get(ARG_SPECIFIED, [])
+                for idx, inp in arg_specified:
+                    symbol_inputs[idx] = getattr(inp, "symbolic_shape", None)
+            symbols_for_parallel = _change_symbols_for_parallel(dataset_shapes, copy.deepcopy(symbol_inputs))
+            if any((s is not None for s in symbols_for_parallel)):
+                self.get_next.add_prim_attr("symbols", symbol_inputs)
+                self.get_next.add_prim_attr("symbols_for_parallel", symbols_for_parallel)
         self.network = network
         self._get_attr_from_cell(network)
 
@@ -129,7 +145,15 @@ def _generate_network_with_dataset(network, dataset_helper, queue_name):
         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)
+        if _is_in_auto_parallel_mode():
+            # here, the dataset shapes has been processed by full_shape(), so need to resume it to original shape
+            # the _check_inputs() will change static origin_shape to dynamic shape
+            # after _check_inputs(), convert dataset_shapes to dynamic shape
+            origin_shape = _origin_shapes(dataset_shapes)
+            _check_inputs(network.get_inputs(), origin_shape, dataset_types)
+            dataset_shapes = _dynamic_shape_for_dataset(dataset_shapes, origin_shape)
+        else:
+            _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(
@@ -141,6 +165,13 @@ def _check_inputs(network_shapes, dataset_shapes, dataset_types):
     """
     Check if set inputs are correct.
     """
+    if not _is_args_fullmode(network_shapes, False):
+        temp_network_shapes = [None for _ in dataset_shapes]
+        arg_specified = network_shapes.get(ARG_SPECIFIED, [])
+        for idx, inp in arg_specified:
+            temp_network_shapes[idx] = inp
+        network_shapes = temp_network_shapes
+
     for tensor_index, ele_dataset_shape in enumerate(dataset_shapes):
         if network_shapes[tensor_index] is None:
             continue
@@ -182,7 +213,8 @@ def _get_dataset_aux(dataset):
 def connect_network_with_dataset(network, dataset_helper):
     """
     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).
+    <https://mindspore.cn/docs/en/master/model_train/train_process/train_optimize.html>`_,
+    (dataset_sink_mode=True).
 
     Args:
         network (Cell): The training network for dataset.
@@ -230,29 +262,36 @@ def connect_network_with_dataset(network, dataset_helper):
             "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, is_dynamic):
+    # In pipeline parallel, some stages have no GetNext, should not get in.
+    use_pipeline_parallel = (context.get_auto_parallel_context("pipeline_stages") > 1)
+    if _dynamic_sink_scenario(dataset, dataset_iter, is_dynamic) and not use_pipeline_parallel:
         dataset_types, dataset_shapes = dataset_helper.get_data_info()
+        # Need to do full_batch for shapes which also do in the _DatasetIterMSLoopSink
+        if _need_to_full():
+            dataset_shapes = _to_full_shapes(dataset_shapes, _get_device_num() // _get_pipeline_stages())
         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:
-            if _need_to_full():
-                device_num = _get_device_num() // _get_pipeline_stages()
-                dataset_shapes = _to_full_shapes(dataset_shapes, device_num)
-
-            network = _generate_dataset_sink_mode_net(
-                network, dataset_shapes, dataset_types, queue_name)
-            if hasattr(aux, '__network_manage__'):
-                aux.__network_manage__ = aux.__network_manage__
+
+        if hasattr(aux, "__shape_type__") and aux.__shape_type__ != key:
+            _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:
-                aux.__network_manage__ = dict()
-            aux.__network_manage__[key] = network
-        network.add_flags(sink_mode=True)
-        return network
+                if _need_to_full():
+                    device_num = _get_device_num() // _get_pipeline_stages()
+                    dataset_shapes = _to_full_shapes(dataset_shapes, device_num)
+
+                network = _generate_dataset_sink_mode_net(
+                    network, dataset_shapes, dataset_types, queue_name)
+                if hasattr(aux, '__network_manage__'):
+                    aux.__network_manage__ = aux.__network_manage__
+                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__
@@ -263,8 +302,11 @@ def connect_network_with_dataset(network, dataset_helper):
             network = _generate_network_with_dataset(
                 network, dataset_helper, queue_name)
             aux.__sink_network__ = network
+            dataset_types, dataset_shapes = dataset_helper.types_shapes()
+            aux.__shape_type__ = str(dataset_types) + str(dataset_shapes)
 
-    if _dynamic_sink_data(dataset, dataset_iter) and _dynamic_sink_exception_scenario(dataset_iter, is_dynamic):
+    if _dynamic_sink_data(dataset, dataset_iter) and _dynamic_sink_exception_scenario(dataset_iter, is_dynamic) and \
+        not use_pipeline_parallel:
         dataset_helper.get_data_info()
     network.add_flags(sink_mode=True)
     return network
@@ -282,7 +324,7 @@ class DatasetHelper:
 
     Args:
         dataset (Dataset): The dataset iterator. The dataset can be generated by dataset generator API in
-                           :class:`mindspore.dataset`, such as :class:`mindspore.dataset.ImageFolderDataset`.
+                           `mindspore.dataset` module, 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``.
@@ -458,6 +500,7 @@ class DatasetHelper:
             '''
             Inner class for parsing send info.
             '''
+
             def __init__(self, send_info, run_context):
                 self.info_ = {}
                 self.sink_size = run_context.original_args()["batch_num"]

mindspore/train/loss_scale_manager.py

@@ -62,7 +62,7 @@ class FixedLossScaleManager(LossScaleManager):
         >>> from mindspore import amp, nn
         >>>
         >>> # 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_scale = 1024.0
         >>> loss_scale_manager = amp.FixedLossScaleManager(loss_scale, False)
@@ -136,7 +136,7 @@ class DynamicLossScaleManager(LossScaleManager):
         >>> from mindspore import amp, nn
         >>>
         >>> # 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_scale_manager = amp.DynamicLossScaleManager()
         >>> optim = nn.Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9)

mindspore/train/metrics/accuracy.py

@@ -32,9 +32,9 @@ class Accuracy(EvaluationBase):
         {\text{true_positive} + \text{true_negative} + \text{false_positive} + \text{false_negative}}
 
     Args:
-        eval_type (str): The metric to calculate the accuracy over a dataset. Supports 'classification' and
-          'multilabel'. 'classification' means the dataset label is single. 'multilabel' means the dataset has multiple
-          labels. Default: ``'classification'`` .
+        eval_type (str): The metric to calculate the accuracy over a dataset. Supports ``'classification'`` and
+          ``'multilabel'``. ``'classification'`` means the dataset label is single.
+          ``'multilabel'`` means the dataset has multiple labels. Default: ``'classification'`` .
 
     Supported Platforms:
         ``Ascend`` ``GPU`` ``CPU``
@@ -67,19 +67,19 @@ class Accuracy(EvaluationBase):
     @rearrange_inputs
     def update(self, *inputs):
         """
-        Updates the local variables. For 'classification', if the index of the maximum of the predict value
-        matches the label, the predict result is correct. For 'multilabel', the predict value match the label,
+        Updates the local variables. For ``'classification'``, if the index of the maximum of the predict value
+        matches the label, the predict result is correct. For ``'multilabel'``, the predict value match the label,
         the predict result is correct.
 
         Args:
             inputs: Logits and labels. `y_pred` stands for logits, `y` stands for labels. `y_pred` and `y` must be a
                 `Tensor`, a list or an array.
 
-                - For the 'classification' evaluation type, `y_pred` is a list of floating numbers in range
+                - For the ``'classification'`` evaluation type, `y_pred` is a list of floating numbers in range
                   :math:`[0, 1]` and the shape is :math:`(N, C)` in most cases (not strictly), where :math:`N`
                   is the number of cases and :math:`C` is the number of categories. `y` must be in one-hot format
                   that shape is :math:`(N, C)`, or can be transformed to one-hot format that shape is :math:`(N,)`.
-                - For 'multilabel' evaluation type, the value of `y_pred` and `y` can only be 0 or 1, indices with 1
+                - For ``'multilabel'`` evaluation type, the value of `y_pred` and `y` can only be 0 or 1, indices with 1
                   indicate the positive category. The shape of `y_pred` and `y` are both :math:`(N, C)`.
 
         Raises:

mindspore/train/metrics/confusion_matrix.py

@@ -33,10 +33,10 @@ class ConfusionMatrix(Metric):
         num_classes (int): Number of classes in the dataset.
         normalize (str): Normalization mode for confusion matrix. Default: ``"no_norm"`` . Choose from:
 
-            - **"no_norm"** (None) - No Normalization is used. Default: ``None``.
-            - **"target"** (str) - Normalization based on target value.
-            - **"prediction"** (str) - Normalization based on predicted value.
-            - **"all"** (str) - Normalization over the whole matrix.
+            - ``"no_norm"`` : No Normalization is used. Default: ``None``.
+            - ``"target"`` : Normalization based on target value.
+            - ``"prediction"`` : Normalization based on predicted value.
+            - ``"all"`` : Normalization over the whole matrix.
 
         threshold (float): The threshold used to compare with the input tensor. Default: ``0.5`` .
 
@@ -224,8 +224,10 @@ class ConfusionMatrixMetric(Metric):
 
                 - y_pred (ndarray): The batch data shape is :math:`(N, C, ...)`
                   or :math:`(N, ...)`, representing onehot format
-                  or category index format respectively. As for classification tasks, y_pred should have the shape [BN]
-                  where N is larger than 1. As for segmentation tasks, the shape should be [BNHW] or [BNHWD].
+                  or category index format respectively. As for classification tasks,
+                  y_pred should have the shape :math:`(B, N)`
+                  where N is larger than 1. As for segmentation tasks,
+                  the shape should be :math:`(B, N, H, W)` or :math:`(B, N, H, W, D)`.
                 - y (ndarray): It must be one-hot format. The batch data shape is :math:`(N, C, ...)`.
 
         Raises: