[Dy2stat] Add Support for paddle.grad (#33110)

This PR made these changes to support double grad:

1. Translate `paddle.grad` to `paddle.static.gradients` to support double grad for dy2stat.
2. Fix IfElseTransformer bug which may not change value if "Store before Load" variable is in "Store" statement is in IfElse conditional statement
3. Add `DOut` to support double grad variables in `run_program_op`
4. Add support for renaming for double grads for `jit.save/load`

Author: zhhsplendid

paddle/fluid/operators/run_program_op.cc

@@ -83,6 +83,13 @@ class RunProgramOpMaker : public framework::OpProtoAndCheckerMaker {
               "contains at most one scope."
               "NOTE: Do not use Scope directly because Scope output is not "
               "currently supported.");
+    AddOutput("DOut",
+              "(vector<LoDTensor>)"
+              "The output tensors for GRAD Tensors in RunProgram forward "
+              "operator, the forward operator contains GRAD Tensors when it "
+              "computes double grad.")
+        .AsDuplicable()
+        .AsDispensable();
     AddAttr<BlockDesc*>("global_block",
                         "(BlockDesc *)"
                         "The global block of executed program desc.");
@@ -154,6 +161,7 @@ class RunProgramGradOpMaker : public framework::SingleGradOpMaker<T> {
     grad_op->SetInput("Params", this->Input("Params"));
     grad_op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
     grad_op->SetInput("OutScope", this->Output("OutScope"));
+    grad_op->SetInput("DOut", this->Output("DOut"));
     grad_op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
     grad_op->SetOutput(framework::GradVarName("Params"),
                        this->InputGrad("Params"));

paddle/fluid/operators/run_program_op.h

@@ -131,6 +131,9 @@ static void ShareVarsIntoScope(const std::vector<Variable *> &vars,
                                const std::vector<std::string> &var_names,
                                framework::Scope *scope) {
   for (size_t i = 0; i < vars.size(); ++i) {
+    if (var_names[i] == "Fake_var") {
+      continue;
+    }
     auto *var = scope->Var(var_names[i]);
     CheckInputVarStatus(*vars[i], var_names[i]);
     VariableShare(*vars[i], var);
@@ -141,9 +144,9 @@ static void ShareVarsFromScope(const std::vector<Variable *> &vars,
                                const std::vector<std::string> &var_names,
                                framework::Scope *scope) {
   for (size_t i = 0; i < vars.size(); ++i) {
-    if (var_names[i] == framework::kEmptyVarName) {
-      VLOG(2) << "find variable name is " << framework::kEmptyVarName
-              << ", skip it!";
+    if (var_names[i] == framework::kEmptyVarName ||
+        var_names[i] == "Fake_var") {
+      VLOG(2) << "find variable name is " << var_names[i] << ", skip it!";
       continue;
     }
     // NOTE: Here skip not found var is dangerous, if a bug is caused here,
@@ -170,9 +173,11 @@ class RunProgramOpKernel : public framework::OpKernel<T> {
     auto &input_vars = ctx.MultiInputVar("X");
     auto &param_vars = ctx.MultiInputVar("Params");
     auto output_vars = ctx.MultiOutputVar("Out");
+    auto dout_vars = ctx.MultiOutputVar("DOut");
 
     auto input_var_names = ctx.InputNames("X");
     auto output_var_names = ctx.OutputNames("Out");
+    auto dout_var_names = ctx.OutputNames("DOut");
 
     // current program may not hold parameters
     std::vector<std::string> param_names;
@@ -195,7 +200,7 @@ class RunProgramOpKernel : public framework::OpKernel<T> {
     // Step 2. prepare executor and init persistable variables
     framework::Executor exe(ctx.GetPlace());
     auto exe_ctx = framework::GetExecutorInfoFromCache(
-        exe, ctx, {output_var_names}, /*is_grad=*/false);
+        exe, ctx, {output_var_names, dout_var_names}, /*is_grad=*/false);
 
     // NOTE(Aurelius84): While training some models, forward can be called many
     // times and then apply backpropagation all at once, such as Reinforcement
@@ -219,6 +224,7 @@ class RunProgramOpKernel : public framework::OpKernel<T> {
 
     // Step 4. Get Output
     details::ShareVarsFromScope(output_vars, output_var_names, &scope);
+    details::ShareVarsFromScope(dout_vars, dout_var_names, &scope);
 
     // Debug info: scope info when run end
     VLOG(3) << framework::GenScopeTreeDebugInfo(out_scope_vec->front());

python/paddle/fluid/dygraph/dygraph_to_static/ast_transformer.py

@@ -25,6 +25,7 @@
 from paddle.fluid.dygraph.dygraph_to_static.break_continue_transformer import BreakTransformOptimizer
 from paddle.fluid.dygraph.dygraph_to_static.call_transformer import CallTransformer
 from paddle.fluid.dygraph.dygraph_to_static.cast_transformer import CastTransformer
+from paddle.fluid.dygraph.dygraph_to_static.grad_transformer import GradTransformer
 from paddle.fluid.dygraph.dygraph_to_static.ifelse_transformer import IfElseTransformer
 from paddle.fluid.dygraph.dygraph_to_static.list_transformer import ListTransformer
 from paddle.fluid.dygraph.dygraph_to_static.logical_transformer import LogicalTransformer
@@ -86,6 +87,7 @@ def transfer_from_node_type(self, node_wrapper):
             PrintTransformer,  # print statement
             CallTransformer,  # transform call recursively
             CastTransformer,  # type casting statement
+            GradTransformer,  # transform paddle.grad to paddle.gradients
         ]
 
         for index, transformer in enumerate(transformers):

python/paddle/fluid/dygraph/dygraph_to_static/grad_transformer.py

@@ -0,0 +1,87 @@
+#   Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+from __future__ import print_function
+
+import gast
+import warnings
+
+from paddle.fluid.dygraph.dygraph_to_static.static_analysis import AstNodeWrapper
+from paddle.fluid.dygraph.dygraph_to_static import utils
+
+
+class GradTransformer(gast.NodeTransformer):
+    """
+    A class transforms dygraph paddle.grad to static graph paddle.gradients. The
+    transformation is applied to support double grad mode.
+    """
+
+    def __init__(self, wrapper_root):
+        assert isinstance(
+            wrapper_root, AstNodeWrapper
+        ), "Input non-AstNodeWrapper node for the initialization of GradTransformer."
+        self.wrapper_root = wrapper_root
+        self.root = wrapper_root.node
+
+    def transform(self):
+        self.visit(self.root)
+
+    def visit_Call(self, node):
+        self.generic_visit(node)
+        if not is_grad_api_node(node):
+            return node
+
+        dygraph_grad_parameters = [
+            "outputs", "inputs", "grad_outputs", "retain_graph", "create_graph",
+            "only_inputs", "allow_unused", "no_grad_vars"
+        ]
+        to_static_grad_param = {
+            "outputs": "targets",
+            "inputs": "inputs",
+            "grad_outputs": "target_gradients",
+            "no_grad_vars": "no_grad_set"
+        }
+        static_keywords = []
+
+        for kw in node.keywords:
+            if kw.arg not in dygraph_grad_parameters or kw.arg not in to_static_grad_param:
+                warnings.warn("paddle.grad has unsupported parameter in jit: " +
+                              kw.arg + ", jit will discard it")
+                continue
+            dygraph_grad_parameters.remove(kw.arg)
+            kw.arg = to_static_grad_param[kw.arg]
+            static_keywords.append(kw)
+
+        for i in range(len(node.args)):
+            arg_name = dygraph_grad_parameters[i]
+            if arg_name not in to_static_grad_param:
+                warnings.warn("paddle.grad has unsupported parameter in jit: " +
+                              kw.arg + ", jit will discard it")
+                continue
+            kw = gast.keyword(
+                arg=to_static_grad_param[arg_name], value=node.args[i])
+            static_keywords.append(kw)
+
+        node.func = gast.parse('paddle.static.gradients').body[0].value
+        node.keywords = static_keywords
+        node.args = []
+        return node
+
+
+def is_grad_api_node(node):
+    assert isinstance(node, gast.Call)
+    api_name = utils.ast_to_source_code(node.func).strip()
+    if utils.is_paddle_api(node):
+        return api_name.endswith("grad")
+    return False

python/paddle/fluid/dygraph/dygraph_to_static/ifelse_transformer.py

@@ -402,7 +402,7 @@ def _modified_vars(child_dict, parent_dict):
             var for var in _vars_with_store(child_dict) if var in parent_dict
         ])
 
-    def _vars_loaded_before_store(ids_dict):
+    def _vars_loaded(ids_dict):
         """
         gast.Param is also a kind of `load` semantic.
         """
@@ -411,8 +411,6 @@ def _vars_loaded_before_store(ids_dict):
             for ctx in ctxs:
                 if isinstance(ctx, (gast.Load, gast.Param)):
                     new_dict[k].append(ctx)
-                elif isinstance(ctx, gast.Store):
-                    break
         return new_dict
 
     # modified vars
@@ -439,8 +437,12 @@ def _vars_loaded_before_store(ids_dict):
     new_vars_in_body_and_orelse = body_new_vars & orelse_new_vars
 
     # 3. new var is created only in one of If.body or If.orelse node, and it used as gast.Load firstly after gast.If node.
+    # TODO(zhhsplendid): the _vars_loaded can be optimized as _vars_loaded_before_store. Because if a variable is stored before load,
+    # the value would change by the store statement, we don't have to return to change the value. However, analysis is
+    # complex because if the IfElse is nested and outer IfElse store statement may not run at all. We will put this optimization
+    # as the future TODO
     used_vars_after_ifelse = set(
-        [var for var in _vars_loaded_before_store(after_ifelse_vars_dict)])
+        [var for var in _vars_loaded(after_ifelse_vars_dict)])
     new_vars_to_create = new_vars_in_one_of_body_or_orelse & used_vars_after_ifelse | new_vars_in_body_and_orelse
 
     # 4. generate return_ids of if/else node.

python/paddle/fluid/dygraph/dygraph_to_static/partial_program.py

@@ -135,6 +135,7 @@ def __init__(self, main_program, inputs, outputs, parameters=None):
         self._origin_main_program = self._verify_program(main_program)
         self._inner_scope = core.Scope()
         # Set default mode to train
+        self._double_grads = self._get_double_grads(self._origin_main_program)
         self.training = True
 
     @LazyInitialized
@@ -192,24 +193,44 @@ def _prune_unused_params(self, program):
         """
         required_params = []
         for param in self._params:
+            found_param = False
             for block in program.blocks:
-                if param.name in block.vars:
-                    required_params.append(param)
+                for op in block.ops:
+                    if param.name in op.input_arg_names or param.name in op.output_arg_names:
+                        required_params.append(param)
+                        found_param = True
+                        break
+                if found_param:
                     break
 
         self._params = required_params
 
+    def _get_double_grads(self, program):
+        double_grads = []
+        for block in program.blocks:
+            for name in block.vars:
+                if "@GRAD" in name:
+                    var_desc = block.vars[name].desc
+                    var_base = core.VarBase(var_desc.dtype(),
+                                            var_desc.shape(),
+                                            var_desc.name(),
+                                            var_desc.type(), False)
+                    double_grads.append(var_base)
+        return double_grads
+
     def forward(self, inputs):
         in_vars, out_vars, tmp_scope_vec = self._prepare(inputs)
-
         framework._dygraph_tracer().trace_op(
             type='run_program',
             inputs={
                 'X': valid_vars(in_vars),
                 'Params': valid_vars(self._params)
             },
-            outputs={'Out': valid_vars(out_vars),
-                     'OutScope': tmp_scope_vec},
+            outputs={
+                'Out': valid_vars(out_vars),
+                'OutScope': tmp_scope_vec,
+                'DOut': valid_vars(self._double_grads)
+            },
             attrs={
                 'global_block': self.program.desc.block(0),
                 'start_op_index': 0,