diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc index 0be7e3db32c..78b611a1484 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc @@ -84,6 +84,7 @@ std::map> kSessionBackends; std::map> kMindRtBackends; PyObjectIdCache g_pyobj_id_cache; +// A warpper to test function and handle the exception template void PynativeExecutorTry(const std::function &method, T *ret, const Args &... args) { const auto inst = PynativeExecutor::GetInstance(); @@ -2807,6 +2808,7 @@ void GradExecutor::GradNetInner(py::object *ret, const prim::GradOperationPtr &g auto size = args.size(); const auto &cell_id = GetGradCellId(grad->sens_param(), cell, args); MS_LOG(DEBUG) << "GradNet start " << size << " " << cell_id; + // return if cell has been compiled if (!top_cell()->need_compile_graph()) { MS_LOG(DEBUG) << "No need compile graph"; if (!cell_stack_.empty()) { @@ -3169,6 +3171,7 @@ void GradExecutor::RunGradGraph(py::object *ret, const py::object &cell, const p MS_EXCEPTION_IF_NULL(resource); MS_LOG(DEBUG) << "Run resource ptr " << resource.get(); + // Get args for run graph VectorRef arg_list; auto filter_args = FilterTensorArgs(args, has_sens); py::tuple converted_args = ConvertArgs(filter_args); @@ -3178,6 +3181,7 @@ void GradExecutor::RunGradGraph(py::object *ret, const py::object &cell, const p compile::VmEvalFuncPtr run = resource->GetResult(pipeline::kOutput).cast(); MS_EXCEPTION_IF_NULL(run); + // Run graph const auto &backend = MsContext::GetInstance()->backend_policy(); MS_LOG(DEBUG) << "Eval run " << backend; grad_is_running_ = true; @@ -3478,10 +3482,12 @@ void GradExecutor::ClearRes() { std::stack().swap(high_order_stack_); } +// return GradExecutor object Ptr grad_executor_ GradExecutorPtr PynativeExecutor::grad_executor() const { MS_EXCEPTION_IF_NULL(grad_executor_); return grad_executor_; } +// return ForwardExecutor object Ptr forward_executor_ ForwardExecutorPtr PynativeExecutor::forward_executor() const { MS_EXCEPTION_IF_NULL(forward_executor_); return forward_executor_; @@ -3533,6 +3539,7 @@ py::object PynativeExecutor::CheckAlreadyRun(const prim::GradOperationPtr &grad, return grad_executor()->CheckAlreadyRun(grad, cell, args); } +// PynativeExecutor_.__call__ py::object PynativeExecutor::Run(const py::object &cell, const py::tuple &args) { py::object ret; PynativeExecutorTry(grad_executor()->RunGraph, &ret, cell, args); @@ -3604,6 +3611,7 @@ py::object PynativeExecutor::GradMsFunction(const py::object &out, const py::arg return grad_executor()->GradMsFunction(out, args); } +// PynativeExecutor_.grad_net void PynativeExecutor::GradNet(const prim::GradOperationPtr &grad, const py::object &cell, const py::object &weights, const py::object &grad_position, const py::args &args) { py::object ret; diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h index 4d6b1b4d0dd..63d9f98f1c8 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h @@ -398,6 +398,7 @@ class ForwardExecutor { std::string last_target_{"Unknown"}; }; +// The executor for pynative mode, std::enable_shared_from_this is for get shared_ptr of this. class PynativeExecutor : public std::enable_shared_from_this { public: static std::shared_ptr GetInstance() { @@ -449,6 +450,7 @@ class PynativeExecutor : public std::enable_shared_from_this { private: PynativeExecutor() = default; + // member variables static std::shared_ptr executor_; static std::mutex instance_lock_; static ForwardExecutorPtr forward_executor_; diff --git a/mindspore/python/mindspore/common/api.py b/mindspore/python/mindspore/common/api.py index f8c2542e423..7ed9d084f08 100644 --- a/mindspore/python/mindspore/common/api.py +++ b/mindspore/python/mindspore/common/api.py @@ -873,6 +873,7 @@ class _PynativeExecutor: Return: The return object after running grad graph. """ + # set the top cell args = args + tuple(kwargs.values()) return self._executor(obj, args) diff --git a/mindspore/python/mindspore/nn/cell.py b/mindspore/python/mindspore/nn/cell.py index c102bb75f92..114216d2c20 100755 --- a/mindspore/python/mindspore/nn/cell.py +++ b/mindspore/python/mindspore/nn/cell.py @@ -966,6 +966,7 @@ class Cell(Cell_): self.compile(*inputs) new_inputs = [] + # init data for i in inputs: if isinstance(i, Tensor): if i.has_init: @@ -981,6 +982,7 @@ class Cell(Cell_): _check_all_tensor(i): new_inputs.append(i) + # call _cell_graph_executor to run graph and calculate gradients if self._auto_parallel_mode: if new_inputs and isinstance(new_inputs[0], Tensor) and inputs[0].virtual_flag: # get parallel inputs in sink mode, parallel inputs set in _cell_graph_executor.compile diff --git a/mindspore/python/mindspore/ops/composite/base.py b/mindspore/python/mindspore/ops/composite/base.py index 495323ede99..516046ae6d6 100644 --- a/mindspore/python/mindspore/ops/composite/base.py +++ b/mindspore/python/mindspore/ops/composite/base.py @@ -474,6 +474,7 @@ class _Grad(GradOperation_): @ms_function def after_grad(*args): return grad_(fn)(*args) + # If calling Grad in PYNATIVE_MODE, do grad in PYNATIVE_MODE elif self.pynative_: _pynative_executor.set_grad_position(grad_, grad_position) @@ -482,10 +483,12 @@ class _Grad(GradOperation_): if _pynative_executor.check_graph(fn, *args, **kwargs): print("Another grad step is running") self._pynative_forward_run(grad_, args, kwargs, fn) + # call _pynative_executor to run graph and calculate gradients _pynative_executor.grad(grad_, fn, weights, grad_position, *args, **kwargs) out = _pynative_executor(fn, *args, **kwargs) _pynative_executor.clear_grad(fn, *args, **kwargs) return out + # If calling Grad in pure PYNATIVE_MODE, do grad in PYNATIVE_MODE else: grad_.pynative_ = True # after_grad of this branch can't use @ms_function, just directly call grad_ @@ -514,16 +517,18 @@ class _Grad(GradOperation_): new_kwargs.pop('sens') else: args = args[:-1] + # fn is a function if isinstance(fn, FunctionType): if not _pynative_executor.check_run(grad, fn, *args, **new_kwargs): _pynative_executor.set_grad_flag(True) _pynative_executor.new_graph(fn, *args, **new_kwargs) outputs = fn(*args, **new_kwargs) _pynative_executor.end_graph(fn, outputs, *args, **new_kwargs) + # fn is a Cell else: # Check if fn has run already. if not _pynative_executor.check_run(grad, fn, *args, **new_kwargs): - fn.set_grad() + fn.set_grad() # set grad flag fn(*args, **new_kwargs) fn.set_grad(False) diff --git a/mindspore/python/mindspore/ops/functional.py b/mindspore/python/mindspore/ops/functional.py index d5fa7513dbe..57a66c05c62 100644 --- a/mindspore/python/mindspore/ops/functional.py +++ b/mindspore/python/mindspore/ops/functional.py @@ -289,6 +289,8 @@ def grad(fn, grad_position=0, sens_param=False): [[-2.00000000e+00, 6.00000000e+00], [-3.00000000e+00, 8.00000000e+00]])) """ + + # Depending on the parameters, call the module callable class in different ways grad_position = _convert_grad_position_type(grad_position) if sens_param: return grad_by_position_with_sens(fn, None, grad_position)