conditional_block_op.cc 12.8 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/pten/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 146
      AssignLocalGradientToParentScope(dev_place, cur_scope, scope,
                                       inside_grads, outside_grads);
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 158 159 160 161 162 163 164 165 166
      const framework::Scope &parent_scope,
      const std::vector<std::string> &inside_grads,
      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 &inside_grad_name = inside_grads[i];
      VLOG(4) << "inside_grad_name = " << inside_grad_name
              << ", outside_grad_name = " << outside_grad_name;
      framework::Variable *inside_var =
          cur_scope.FindLocalVar(inside_grad_name);
      if (inside_var == nullptr) {
Y
Yu Yang 已提交
167 168
        continue;
      }
169 170 171 172 173 174 175 176 177
      framework::Variable *outside_var =
          parent_scope.FindVar(outside_grad_name);
      if (outside_var == nullptr) {
        continue;
      }
      platform::DeviceContext *dev_ctx =
          platform::DeviceContextPool::Instance().Get(place);
      framework::VisitVarType(*inside_var,
                              AssignFunctor(outside_var, *dev_ctx));
Y
Yu Yang 已提交
178 179
    }
  }
180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202

  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 "
203
                              "doesn't match input_var %s.",
204 205 206 207 208 209 210 211 212
                              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, "
213
                              "which doesn't match input_var %s.",
214 215 216 217 218 219 220
                              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 "
221
                              "size as input_var %s.",
222 223 224 225 226 227 228 229 230
                              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 "
231
            "now."));
232 233 234 235 236 237 238 239 240 241 242 243 244
      }
    }
  }

  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());
245
    outside_tensor->mutable_data(place, input_tensor.dtype());
246 247
    const platform::DeviceContext *dev_ctx =
        platform::DeviceContextPool::Instance().Get(place);
248
    pten::funcs::set_constant(*dev_ctx, outside_tensor, 0.0f);
249 250
    outside_tensor->set_lod(input_tensor.lod());
  }
Y
Yu Yang 已提交
251 252 253 254 255
};

class ConditionalBlockGradInferShape : public framework::InferShapeBase {
 public:
  void operator()(framework::InferShapeContext *context) const override {
256 257 258 259
    PADDLE_ENFORCE_EQ(
        context->HasInputs(ConditionalOp::kCondition), true,
        platform::errors::InvalidArgument(
            "Condition must be set in conditional_block_grad_op."));
260 261
    if (context->HasInputs(ConditionalOp::kInputs) &&
        context->HasOutputs(framework::GradVarName(ConditionalOp::kInputs))) {
Z
Zeng Jinle 已提交
262 263
      context->SetOutputsDim(framework::GradVarName(ConditionalOp::kInputs),
                             context->GetInputsDim(ConditionalOp::kInputs));
Y
Yu Yang 已提交
264 265 266 267
    }
  }
};

H
hong 已提交
268 269
template <typename T>
class ConditionalBlockGradMaker : public framework::SingleGradOpMaker<T> {
Y
Yu Yang 已提交
270
 public:
H
hong 已提交
271
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
Y
Yu Yang 已提交
272 273

 protected:
274
  void Apply(GradOpPtr<T> grad_op) const override {
Y
Yu Yang 已提交
275
    grad_op->SetType("conditional_block_grad");
Z
Zeng Jinle 已提交
276
    grad_op->SetInput(ConditionalOp::kCondition,
H
hong 已提交
277 278 279 280 281
                      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 已提交
282
    grad_op->SetInput(framework::GradVarName(ConditionalOp::kOutputs),
H
hong 已提交
283 284 285
                      this->OutputGrad(ConditionalOp::kOutputs));
    grad_op->SetInput(ConditionalOp::kScope,
                      this->Output(ConditionalOp::kScope));
Z
Zeng Jinle 已提交
286
    grad_op->SetOutput(framework::GradVarName(ConditionalOp::kInputs),
H
hong 已提交
287
                       this->InputGrad(ConditionalOp::kInputs, false));
A
Abhinav Arora 已提交
288
    grad_op->SetBlockAttr("sub_block", this->grad_block_[0]);
H
hong 已提交
289 290
    grad_op->SetAttr("is_scalar_condition",
                     this->GetAttr("is_scalar_condition"));
Y
Yu Yang 已提交
291 292 293 294 295 296 297 298
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
REGISTER_OPERATOR(conditional_block, ops::ConditionalBlockOp,
299
                  ops::ConditionalBlockInferShape,
Y
Yu Yang 已提交
300
                  ops::ConditionalBlockOpProtoMaker,
H
hong 已提交
301
                  ops::ConditionalBlockGradMaker<paddle::framework::OpDesc>);
Y
Yu Yang 已提交
302 303
REGISTER_OPERATOR(conditional_block_grad, ops::ConditionalBlockGradOp,
                  ops::ConditionalBlockGradInferShape);