conditional_block_op.cc 14.4 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
Y
Yu Yang 已提交
2

L
Luo Tao 已提交
3 4 5
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
Y
Yu Yang 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
Y
Yu Yang 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
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. */
14 15

#include "paddle/fluid/operators/controlflow/conditional_block_op.h"
16

17
#include "paddle/fluid/operators/assign_op.h"
18
#include "paddle/phi/kernels/funcs/math_function.h"
Y
Yu Yang 已提交
19 20 21 22

namespace paddle {
namespace operators {

Z
Zeng Jinle 已提交
23 24 25 26 27 28
const char ConditionalOp::kInputs[] = "Input";
const char ConditionalOp::kOutputs[] = "Out";
const char ConditionalOp::kCondition[] = "Cond";
const char ConditionalOp::kScope[] = "Scope";
const char ConditionalOp::kSkipEagerDeletionVars[] = "skip_eager_deletion_vars";

Y
Yu Yang 已提交
29 30 31 32 33 34 35
class ConditionalBlockOp : public ConditionalOp {
 public:
  ConditionalBlockOp(const std::string &type,
                     const framework::VariableNameMap &inputs,
                     const framework::VariableNameMap &outputs,
                     const framework::AttributeMap &attrs)
      : ConditionalOp(type, inputs, outputs, attrs) {}
36 37 38 39

 private:
  void RunImpl(const framework::Scope &scope,
               const platform::Place &dev_place) const override {
40 41
    bool need_run;
    if (Attr<bool>("is_scalar_condition")) {
42 43 44
      // When is_scalar_condition is True, the conditional variable is a scalar,
      // whether need to execute the operators in sub-block depends on the
      // conditional variable (Cond).
Z
Zeng Jinle 已提交
45
      auto xs = InputTensors(scope, ConditionalOp::kCondition);
46 47
      need_run = ScalarCondition(xs);
    } else {
48 49 50
      // When is_scalar_condition is False, the conditional variable maybe a
      // vector or tensor, whether need to execute the operators in sub-block
      // depends on the input variables (Input).
Z
Zeng Jinle 已提交
51
      auto xs = InputTensors(scope, ConditionalOp::kInputs);
52 53 54 55
      need_run = std::all_of(
          xs.begin(), xs.end(),
          [](const framework::LoDTensor *t) { return t->numel() != 0; });
    }
Y
Yu Yang 已提交
56 57

    if (need_run) {
Z
Zeng Jinle 已提交
58
      auto *scope_var = scope.FindVar(Output(ConditionalOp::kScope));
59
      PADDLE_ENFORCE_NOT_NULL(
60 61 62 63
          scope_var,
          platform::errors::PreconditionNotMet(
              "Expect Scope variable to be set in conditional_block_op, but "
              "got a null Scope variable. Please set the Scope variable."));
Y
Yu Yang 已提交
64 65 66 67
      auto *scopes = scope_var->GetMutable<std::vector<framework::Scope *>>();
      scopes->resize(1);
      scopes->front() = &scope.NewScope();
      auto &cur_scope = *scopes->front();
D
dzhwinter 已提交
68
      framework::Executor exec(dev_place);
Y
Yu Yang 已提交
69
      auto *block = Attr<framework::BlockDesc *>("sub_block");
70 71
      VLOG(3) << "Conditional block.idx = " << block->ID()
              << ", scope = " << &cur_scope;
Z
Zeng Jinle 已提交
72 73 74
      auto &skip_vars =
          Attr<std::vector<std::string>>(ConditionalOp::kSkipEagerDeletionVars);
      exec.Run(*block->Program(), &cur_scope, block->ID(), false, true,
75 76
               skip_vars, /* force_disable_gc */ false,
               /* keep_kid_scopes */ true);
Y
Yu Yang 已提交
77 78 79 80
    }
  }
};

81 82 83 84 85
class ConditionalBlockInferShape : public framework::InferShapeBase {
 public:
  void operator()(framework::InferShapeContext *context) const override {
    PADDLE_ENFORCE_EQ(context->HasInputs(ConditionalOp::kCondition), true,
                      platform::errors::InvalidArgument(
86
                          "conditional_block_op must have condition input."));
87 88 89
  }
};

Y
Yu Yang 已提交
90 91 92 93 94 95 96
class ConditionalBlockGradOp : public ConditionalOp {
 public:
  ConditionalBlockGradOp(const std::string &type,
                         const framework::VariableNameMap &inputs,
                         const framework::VariableNameMap &outputs,
                         const framework::AttributeMap &attrs)
      : ConditionalOp(type, inputs, outputs, attrs) {}
97 98 99 100

 private:
  void RunImpl(const framework::Scope &scope,
               const platform::Place &dev_place) const override {
101 102
    bool need_run;
    if (Attr<bool>("is_scalar_condition")) {
Z
Zeng Jinle 已提交
103
      auto xs = this->InputTensors(scope, ConditionalOp::kCondition);
104 105
      need_run = ScalarCondition(xs);
    } else {
Z
Zeng Jinle 已提交
106
      auto xs = this->InputTensors(scope, ConditionalOp::kInputs);
107 108 109 110
      need_run = std::all_of(
          xs.begin(), xs.end(),
          [](const framework::LoDTensor *t) { return t->numel() != 0; });
    }
Y
Yu Yang 已提交
111

112 113 114
    const auto &inputs = Inputs(ConditionalOp::kInputs);
    const auto &outside_grads =
        Outputs(framework::GradVarName(ConditionalOp::kInputs));
Y
Yu Yang 已提交
115
    if (need_run) {
116 117 118 119 120 121
      std::vector<std::string> inside_grads;
      inside_grads.reserve(inputs.size());
      for (auto &in : inputs) {
        inside_grads.emplace_back(framework::GradVarName(in));
      }

Z
Zeng Jinle 已提交
122
      auto *scope_var = scope.FindVar(Input(ConditionalOp::kScope));
123
      PADDLE_ENFORCE_NOT_NULL(
124 125 126 127
          scope_var,
          platform::errors::PreconditionNotMet(
              "Expect Scope variable to be set in conditional_block_op, but "
              "got a null Scope variable. Please set the Scope variable."));
Y
Yu Yang 已提交
128
      auto &scopes = scope_var->Get<std::vector<framework::Scope *>>();
129 130 131 132 133
      PADDLE_ENFORCE_GT(
          scopes.size(), 0,
          platform::errors::InvalidArgument(
              "Expect Scope variable contains at least 1 scope, but got: %d",
              scopes.size()));
Y
Yu Yang 已提交
134 135
      framework::Scope &cur_scope = *scopes[0];

D
dzhwinter 已提交
136
      framework::Executor exec(dev_place);
Y
Yu Yang 已提交
137
      auto *block = Attr<framework::BlockDesc *>("sub_block");
Y
Yu Yang 已提交
138

139 140
      VLOG(3) << "Conditional Grad block.idx = " << block->ID()
              << ", scope = " << &cur_scope;
141
      exec.Run(*block->Program(), &cur_scope, block->ID(), false, true,
142 143
               inside_grads, /* force_disable_gc */ false,
               /* keep_kid_scopes */ false);
Y
Yu Yang 已提交
144

145
      AssignLocalGradientToParentScope(dev_place, cur_scope, scope,
146
                                       inside_grads, outside_grads, inputs);
147
      return;
Y
Yu Yang 已提交
148
    }
149 150

    AssignZeroToParentScope(dev_place, scope, inputs, outside_grads);
Y
Yu Yang 已提交
151 152 153
  }

 private:
154
  void AssignLocalGradientToParentScope(
D
dzhwinter 已提交
155
      const platform::Place &place, const framework::Scope &cur_scope,
156 157
      const framework::Scope &parent_scope,
      const std::vector<std::string> &inside_grads,
158 159 160 161
      const std::vector<std::string> &outside_grads,
      const std::vector<std::string> &inputs) const {
    std::vector<std::string> assign_zero_outside_grads;
    std::vector<std::string> assign_zero_inputs;
162 163 164 165 166 167 168 169 170 171
    for (size_t i = 0; i < outside_grads.size(); ++i) {
      const std::string &outside_grad_name = outside_grads[i];
      const std::string &inside_grad_name = inside_grads[i];
      VLOG(4) << "inside_grad_name = " << inside_grad_name
              << ", outside_grad_name = " << outside_grad_name;
      framework::Variable *outside_var =
          parent_scope.FindVar(outside_grad_name);
      if (outside_var == nullptr) {
        continue;
      }
172 173 174 175 176 177 178
      framework::Variable *inside_var =
          cur_scope.FindLocalVar(inside_grad_name);
      if (inside_var == nullptr) {
        assign_zero_outside_grads.emplace_back(outside_grad_name);
        assign_zero_inputs.emplace_back(inputs[i]);
        continue;
      }
179 180 181 182
      platform::DeviceContext *dev_ctx =
          platform::DeviceContextPool::Instance().Get(place);
      framework::VisitVarType(*inside_var,
                              AssignFunctor(outside_var, *dev_ctx));
Y
Yu Yang 已提交
183
    }
184 185 186 187
    // Assign zero to the grad_vars that are in outside_grads but not in
    // inside_grads
    AssignZeroToParentScope(place, parent_scope, assign_zero_inputs,
                            assign_zero_outside_grads);
Y
Yu Yang 已提交
188
  }
189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211

  void AssignZeroToParentScope(
      const platform::Place &place, const framework::Scope &scope,
      const std::vector<std::string> &inputs,
      const std::vector<std::string> &outside_grads) const {
    for (size_t i = 0; i < outside_grads.size(); ++i) {
      const std::string &outside_grad_name = outside_grads[i];
      const std::string &input_name = inputs[i];
      VLOG(4) << "input_name = " << input_name
              << ", outside_grad_name = " << outside_grad_name;
      framework::Variable *input_var = scope.FindVar(input_name);
      if (input_var == nullptr) {
        continue;
      }
      framework::Variable *outside_var = scope.FindVar(outside_grad_name);
      if (outside_var == nullptr) {
        continue;
      }

      if (input_var->IsType<framework::LoDTensor>()) {
        PADDLE_ENFORCE_EQ(outside_var->IsType<framework::LoDTensor>(), true,
                          platform::errors::InvalidArgument(
                              "Type of outside_var %s is NOT LoDTensor, which "
212
                              "doesn't match input_var %s.",
213 214 215 216 217 218 219 220 221
                              outside_grad_name, input_name));
        AssignZeroToOutsideTensor(
            place, scope, input_var->Get<framework::LoDTensor>(),
            outside_var->GetMutable<framework::LoDTensor>());
      } else if (input_var->IsType<framework::LoDTensorArray>()) {
        PADDLE_ENFORCE_EQ(outside_var->IsType<framework::LoDTensorArray>(),
                          true,
                          platform::errors::InvalidArgument(
                              "Type of outside_var %s is NOT LoDTensorArray, "
222
                              "which doesn't match input_var %s.",
223 224 225 226 227 228 229
                              outside_grad_name, input_name));
        const auto &input_tensors = input_var->Get<framework::LoDTensorArray>();
        auto *outside_tensors =
            outside_var->GetMutable<framework::LoDTensorArray>();
        PADDLE_ENFORCE_EQ(input_tensors.size(), outside_tensors->size(),
                          platform::errors::InvalidArgument(
                              "LoDTensorArray outside_var %s doen't have same "
230
                              "size as input_var %s.",
231 232 233 234 235 236 237 238 239
                              outside_grad_name, input_name));
        for (size_t j = 0; j < input_tensors.size(); ++j) {
          AssignZeroToOutsideTensor(place, scope, input_tensors[j],
                                    &((*outside_tensors)[j]));
        }
      } else {
        // TODO(huihuangzheng): add support for SelectedRows
        PADDLE_THROW(platform::errors::InvalidArgument(
            "Conditional block grad op doesn't support non-LoDTensor output "
240
            "now."));
241 242 243 244 245 246 247 248 249 250 251 252 253
      }
    }
  }

  void AssignZeroToOutsideTensor(const platform::Place &place,
                                 const framework::Scope &cur_scope,
                                 const framework::LoDTensor &input_tensor,
                                 framework::LoDTensor *outside_tensor) const {
    if (!input_tensor.IsInitialized() || input_tensor.numel() == 0) {
      return;
    }
    VLOG(4) << "Assigning zero to " << outside_tensor;
    outside_tensor->Resize(input_tensor.dims());
254
    outside_tensor->mutable_data(place, input_tensor.dtype());
255 256
    const platform::DeviceContext *dev_ctx =
        platform::DeviceContextPool::Instance().Get(place);
257
    phi::funcs::set_constant(*dev_ctx, outside_tensor, 0.0f);
258 259
    outside_tensor->set_lod(input_tensor.lod());
  }
Y
Yu Yang 已提交
260 261 262 263 264
};

class ConditionalBlockGradInferShape : public framework::InferShapeBase {
 public:
  void operator()(framework::InferShapeContext *context) const override {
265 266 267 268
    PADDLE_ENFORCE_EQ(
        context->HasInputs(ConditionalOp::kCondition), true,
        platform::errors::InvalidArgument(
            "Condition must be set in conditional_block_grad_op."));
269 270
    if (context->HasInputs(ConditionalOp::kInputs) &&
        context->HasOutputs(framework::GradVarName(ConditionalOp::kInputs))) {
Z
Zeng Jinle 已提交
271 272
      context->SetOutputsDim(framework::GradVarName(ConditionalOp::kInputs),
                             context->GetInputsDim(ConditionalOp::kInputs));
Y
Yu Yang 已提交
273 274 275 276
    }
  }
};

277 278 279 280 281 282 283
class ConditionalBlockGradInferVarType : public framework::VarTypeInference {
 public:
  void operator()(framework::InferVarTypeContext *ctx) const override {
    // NOTE(Aurelius84): VarType of Output is LoDTensor by default. In case of
    // Input is {Tensor, LoDTensorArray}, we need synchronous the Input's
    // VarType into Input@GRAD to avoid generating {Tensor, Tensor} as
    // Input@GRAD.
284 285 286 287 288 289 290 291 292 293 294 295
    auto input_size = ctx->InputSize(ConditionalOp::kInputs);
    auto output_size =
        ctx->OutputSize(framework::GradVarName(ConditionalOp::kInputs));
    PADDLE_ENFORCE_EQ(input_size, output_size,
                      platform::errors::InvalidArgument(
                          "input_size and output_size should be equal for "
                          "conditional_block_grad_op."));
    for (size_t i = 0; i < output_size; ++i) {
      ctx->SyncTypeAndDataType(ConditionalOp::kInputs,
                               framework::GradVarName(ConditionalOp::kInputs),
                               i);
    }
296 297 298
  }
};

H
hong 已提交
299 300
template <typename T>
class ConditionalBlockGradMaker : public framework::SingleGradOpMaker<T> {
Y
Yu Yang 已提交
301
 public:
H
hong 已提交
302
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
Y
Yu Yang 已提交
303 304

 protected:
305
  void Apply(GradOpPtr<T> grad_op) const override {
Y
Yu Yang 已提交
306
    grad_op->SetType("conditional_block_grad");
Z
Zeng Jinle 已提交
307
    grad_op->SetInput(ConditionalOp::kCondition,
H
hong 已提交
308 309 310 311 312
                      this->Input(ConditionalOp::kCondition));
    grad_op->SetInput(ConditionalOp::kInputs,
                      this->Input(ConditionalOp::kInputs));
    grad_op->SetInput(ConditionalOp::kOutputs,
                      this->Output(ConditionalOp::kOutputs));
Z
Zeng Jinle 已提交
313
    grad_op->SetInput(framework::GradVarName(ConditionalOp::kOutputs),
H
hong 已提交
314 315 316
                      this->OutputGrad(ConditionalOp::kOutputs));
    grad_op->SetInput(ConditionalOp::kScope,
                      this->Output(ConditionalOp::kScope));
Z
Zeng Jinle 已提交
317
    grad_op->SetOutput(framework::GradVarName(ConditionalOp::kInputs),
H
hong 已提交
318
                       this->InputGrad(ConditionalOp::kInputs, false));
A
Abhinav Arora 已提交
319
    grad_op->SetBlockAttr("sub_block", this->grad_block_[0]);
H
hong 已提交
320 321
    grad_op->SetAttr("is_scalar_condition",
                     this->GetAttr("is_scalar_condition"));
Y
Yu Yang 已提交
322 323 324 325 326 327 328 329
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
REGISTER_OPERATOR(conditional_block, ops::ConditionalBlockOp,
330
                  ops::ConditionalBlockInferShape,
Y
Yu Yang 已提交
331
                  ops::ConditionalBlockOpProtoMaker,
H
hong 已提交
332
                  ops::ConditionalBlockGradMaker<paddle::framework::OpDesc>);
Y
Yu Yang 已提交
333
REGISTER_OPERATOR(conditional_block_grad, ops::ConditionalBlockGradOp,
334 335
                  ops::ConditionalBlockGradInferShape,
                  ops::ConditionalBlockGradInferVarType);