diff --git a/mindspore/ccsrc/backend/session/ascend_session.cc b/mindspore/ccsrc/backend/session/ascend_session.cc index 6d106bde455..4dac2a5d9e4 100644 --- a/mindspore/ccsrc/backend/session/ascend_session.cc +++ b/mindspore/ccsrc/backend/session/ascend_session.cc @@ -1383,5 +1383,14 @@ GraphId AscendSession::CompileGraphImpl(NotNull func_graph, const RunInfer(func_graph, inputs); return CompileGraphImpl(func_graph); } + +void AscendSession::SyncStream() { + auto runtime_instance = device::KernelRuntimeManager::Instance().GetKernelRuntime(kAscendDevice, device_id_); + MS_EXCEPTION_IF_NULL(runtime_instance); + auto ret = runtime_instance->SyncStream(); + if (!ret) { + MS_LOG(ERROR) << "Sync stream error!"; + } +} } // namespace session } // namespace mindspore diff --git a/mindspore/ccsrc/backend/session/ascend_session.h b/mindspore/ccsrc/backend/session/ascend_session.h index c4e815c8f7e..100882159e3 100644 --- a/mindspore/ccsrc/backend/session/ascend_session.h +++ b/mindspore/ccsrc/backend/session/ascend_session.h @@ -48,6 +48,7 @@ class AscendSession : public SessionBasic { void Init(uint32_t device_id) override; // get graph id of final graph GraphId GetFinalRunGraph() const override { return final_graph_id_; } + void SyncStream() override; protected: GraphId CompileGraphImpl(const AnfNodePtrList &lst, const AnfNodePtrList &outputs) override; diff --git a/mindspore/ccsrc/backend/session/gpu_session.cc b/mindspore/ccsrc/backend/session/gpu_session.cc index 30822176501..994c5c5f6d5 100644 --- a/mindspore/ccsrc/backend/session/gpu_session.cc +++ b/mindspore/ccsrc/backend/session/gpu_session.cc @@ -508,6 +508,15 @@ void GPUSession::CleanValueNodeDeviceAddr(const std::shared_ptr &ke MS_EXCEPTION_IF_NULL(runtime_instance); runtime_instance->CleanValueNodeDeviceAddr(kernel_graph.get()); } + +void GPUSession::SyncStream() { + auto runtime_instance = device::KernelRuntimeManager::Instance().GetSingleKernelRuntime(kGPUDevice, device_id_); + MS_EXCEPTION_IF_NULL(runtime_instance); + auto ret = runtime_instance->SyncStream(); + if (!ret) { + MS_LOG(ERROR) << "Sync stream error!"; + } +} } // namespace gpu } // namespace session } // namespace mindspore diff --git a/mindspore/ccsrc/backend/session/gpu_session.h b/mindspore/ccsrc/backend/session/gpu_session.h index e3dc825c213..024a9b4fb7d 100644 --- a/mindspore/ccsrc/backend/session/gpu_session.h +++ b/mindspore/ccsrc/backend/session/gpu_session.h @@ -32,6 +32,7 @@ class GPUSession : public SessionBasic { GPUSession() = default; ~GPUSession() override = default; void Init(uint32_t device_id) override; + void SyncStream() override; protected: GraphId CompileGraphImpl(const AnfNodePtrList &lst, const AnfNodePtrList &outputs) override; diff --git a/mindspore/ccsrc/backend/session/session_basic.h b/mindspore/ccsrc/backend/session/session_basic.h index 9c95464bc96..b9e38bf6c28 100644 --- a/mindspore/ccsrc/backend/session/session_basic.h +++ b/mindspore/ccsrc/backend/session/session_basic.h @@ -66,9 +66,8 @@ class SessionBasic : public std::enable_shared_from_this { } virtual void Init(uint32_t device_id) { device_id_ = device_id; } - void InitExecutor(const std::string &device_name, uint32_t device_id); - + virtual void SyncStream() {} virtual ~SessionBasic() { summary_callback_ = nullptr; } GraphId CompileGraph(const GraphSegmentPtr &segment, const AnfNodePtrList &outputs); diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc index 28ccb484fe5..fc9d51ad538 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc @@ -2113,6 +2113,13 @@ void PynativeExecutor::GradNet(const GradOperationPtr &grad, const py::object &c PynativeExecutorTry(this, &PynativeExecutor::GradNetInner, grad, cell, weights, args); } +void PynativeExecutor::Sync() { + if (session == nullptr) { + MS_EXCEPTION(NotExistsError) << "No session has been created!"; + } + session->SyncStream(); +} + REGISTER_PYBIND_DEFINE(PynativeExecutor_, ([](const py::module *m) { (void)py::class_>(*m, "PynativeExecutor_") .def_static("get_instance", &PynativeExecutor::GetInstance, "PynativeExecutor get_instance.") @@ -2121,6 +2128,7 @@ REGISTER_PYBIND_DEFINE(PynativeExecutor_, ([](const py::module *m) { .def("check_graph", &PynativeExecutor::CheckGraph, "pynative check a grad graph.") .def("grad_net", &PynativeExecutor::GradNet, "pynative grad graph.") .def("clear", &PynativeExecutor::Clear, "pynative clear status.") + .def("sync", &PynativeExecutor::Sync, "pynative sync stream.") .def("__call__", &PynativeExecutor::Run, py::arg("args"), py::arg("phase") = py::str(""), "Executor run function.") .def("set_grad_flag", &PynativeExecutor::set_grad_flag, py::arg("flag") = py::bool_(false), diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h index eea958b5e5c..f41034017c4 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h @@ -96,6 +96,8 @@ class PynativeExecutor : public std::enable_shared_from_this { void Clean(); // Destrcut call void ClearRes(); + // Sync stream + void Sync(); private: PynativeExecutor() = default; diff --git a/mindspore/common/api.py b/mindspore/common/api.py index 5e8f830a3be..f7f0dd8aa13 100644 --- a/mindspore/common/api.py +++ b/mindspore/common/api.py @@ -314,6 +314,9 @@ class _PynativeExecutor: def clear(self, flag=""): self._executor.clear(flag) + def sync(self): + self._executor.sync() + def set_grad_flag(self, flag): self._executor.set_grad_flag(flag) diff --git a/mindspore/train/dataset_helper.py b/mindspore/train/dataset_helper.py index 78c8cc14926..709c3cd326c 100644 --- a/mindspore/train/dataset_helper.py +++ b/mindspore/train/dataset_helper.py @@ -67,6 +67,7 @@ def connect_network_with_dataset(network, dataset_helper): >>> net = Net() >>> net_with_get_next = connect_network_with_dataset(net, dataset_helper) """ + class _DataWrapper(nn.Cell): """ Wraps the input network with a dataset which automatically fetches data with 'GetNext' function from the @@ -163,16 +164,20 @@ class DatasetHelper: if context.get_context("enable_ge"): iterclass = _DatasetIterGE else: - if context.get_context("device_target") == "Ascend": - iterclass = _DatasetIterMSLoopSink - elif context.get_context("device_target") == "GPU": - ms_role = os.getenv("MS_ROLE") - if ms_role in ("MS_PSERVER", "MS_SCHED"): - iterclass = _DatasetIterPSLite - else: + if context.get_context("mode") == context.GRAPH_MODE: + if context.get_context("device_target") == "Ascend": iterclass = _DatasetIterMSLoopSink - elif context.get_context("device_target") == "CPU": - raise RuntimeError("Currently dataset sink mode is not supported when the device target is CPU.") + elif context.get_context("device_target") == "GPU": + ms_role = os.getenv("MS_ROLE") + if ms_role in ("MS_PSERVER", "MS_SCHED"): + iterclass = _DatasetIterPSLite + else: + iterclass = _DatasetIterMSLoopSink + elif context.get_context("device_target") == "CPU": + raise RuntimeError( + "Currently dataset sink mode is not supported when the device target is CPU.") + else: + iterclass = _DatasetIterPyNative self.iter = iterclass(dataset, sink_size, epoch_num) else: iterclass = _DatasetIterNormal @@ -281,6 +286,20 @@ class _DatasetIterGE(_DatasetIter): self.op = op +class _DatasetIterPyNative(_DatasetIter): + """Iter for MS(enable_loop_sink=False).""" + + def __init__(self, dataset, sink_size, epoch_num): + super().__init__(dataset, sink_size, epoch_num) + if sink_size > 0: + self.sink_count = sink_size + else: + self.sink_count = dataset.get_dataset_size() + + def op(): + return tuple() + + self.op = op class _DatasetIterMSLoopSink(_DatasetIter): """Iter for context (device_target=Ascend)""" @@ -329,6 +348,7 @@ class _DatasetIterPSLite(_DatasetIter): def op(): return _construct_tensor_list(self.dataset_types, self.dataset_shapes, batch_expand_num=1) + self.op = op diff --git a/mindspore/train/model.py b/mindspore/train/model.py index 76b8d30a965..5ce36b339c3 100755 --- a/mindspore/train/model.py +++ b/mindspore/train/model.py @@ -23,7 +23,7 @@ from mindspore import log as logger from ..common.tensor import Tensor from ..nn.metrics import get_metrics from .._checkparam import check_input_data, check_output_data, Validator -from .callback import _InternalCallbackParam, RunContext, _CallbackManager +from .callback import _InternalCallbackParam, RunContext, _CallbackManager, Callback from .. import context from ..parallel._utils import _get_parallel_mode, _get_device_num, _get_global_rank, \ _get_parameter_broadcast, _device_number_check, _parameter_broadcast_check @@ -35,6 +35,7 @@ from ..context import ParallelMode from ..parallel._cost_model_context import _set_multi_subgraphs from .dataset_helper import DatasetHelper, connect_network_with_dataset from . import amp +from ..common.api import _pynative_exec def _transfer_tensor_to_tuple(inputs): @@ -47,6 +48,11 @@ def _transfer_tensor_to_tuple(inputs): return inputs +class _StepSync(Callback): + def step_end(self, run_context): + _pynative_exec.sync() + + class Model: """ High-Level API for Training or Testing. @@ -365,6 +371,9 @@ class Model: cb_params.device_number = self._device_number cb_params.train_dataset = train_dataset cb_params.list_callback = self._transform_callbacks(callbacks) + if context.get_context("mode") == context.PYNATIVE_MODE: + cb_params.list_callback.insert(0, _StepSync()) + callbacks = cb_params.list_callback cb_params.train_dataset_element = None cb_params.network = self._network if _is_role_pserver() or _is_role_sched(): @@ -374,8 +383,8 @@ class Model: with _CallbackManager(callbacks) as list_callback: if not dataset_sink_mode: self._train_process(epoch, train_dataset, list_callback, cb_params) - elif context.get_context("device_target") == "CPU" or context.get_context("mode") == context.PYNATIVE_MODE: - logger.warning("The CPU or PyNative mode cannot support dataset sink mode currently." + elif context.get_context("device_target") == "CPU": + logger.warning("The CPU cannot support dataset sink mode currently." "So the training process will be performed with dataset not sink.") self._train_process(epoch, train_dataset, list_callback, cb_params) else: @@ -417,7 +426,7 @@ class Model: run_context = RunContext(cb_params) list_callback.begin(run_context) - + is_graph = (context.get_context("mode") == context.GRAPH_MODE) # used to stop training for early stop, such as stopAtTIme or stopATStep should_stop = False dataset_helper = None @@ -441,7 +450,10 @@ class Model: cb_params.train_dataset_element = inputs list_callback.step_begin(run_context) outputs = self._train_network(*inputs) - cb_params.cur_step_num += dataset_helper.sink_size() + if is_graph: + cb_params.cur_step_num += dataset_helper.sink_size() + else: + cb_params.cur_step_num += 1 cb_params.net_outputs = outputs list_callback.step_end(run_context)