mindspore/ccsrc/pipeline/pynative/pynative_execute.cc

@@ -84,6 +84,7 @@ std::map<std::string, std::shared_ptr<session::SessionBasic>> kSessionBackends;
 std::map<std::string, std::shared_ptr<compile::MindRTBackend>> kMindRtBackends;
 PyObjectIdCache g_pyobj_id_cache;
 
+// A warpper to test function and handle the exception
 template <typename T, typename... Args>
 void PynativeExecutorTry(const std::function<void(T *ret, const Args &...)> &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<compile::VmEvalFuncPtr>();
   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<TopCellInfoPtr>().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;

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<PynativeExecutor> {
  public:
   static std::shared_ptr<PynativeExecutor> GetInstance() {
@@ -449,6 +450,7 @@ class PynativeExecutor : public std::enable_shared_from_this<PynativeExecutor> {
  private:
   PynativeExecutor() = default;
 
+  // member variables
   static std::shared_ptr<PynativeExecutor> executor_;
   static std::mutex instance_lock_;
   static ForwardExecutorPtr forward_executor_;

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)
 

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

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)
 

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)