while_op.cc 15.9 KB
Newer Older
C
chengduo 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
// Copyright (c) 2018 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.
Y
Yang Yang(Tony) 已提交
14 15

#include <vector>
Y
Yi Wang 已提交
16 17 18 19
#include "paddle/fluid/framework/executor.h"
#include "paddle/fluid/framework/lod_tensor_array.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/operator.h"
S
sneaxiy 已提交
20
#include "paddle/fluid/framework/var_type.h"
Y
Yi Wang 已提交
21
#include "paddle/fluid/operators/detail/safe_ref.h"
Y
Yang Yang(Tony) 已提交
22 23 24 25 26 27 28

namespace paddle {
namespace operators {

using StepScopeVar = std::vector<framework::Scope *>;
using LoDTensor = framework::LoDTensor;

Y
Yang Yu 已提交
29 30 31 32 33 34
static constexpr char kStepBlock[] = "sub_block";
static constexpr char kCondition[] = "Condition";
static constexpr char kStepScopes[] = "StepScopes";
static constexpr char kX[] = "X";
static constexpr char kXGRAD[] = "X@GRAD";
static constexpr char kOutputs[] = "Out";
Y
Yang Yang(Tony) 已提交
35 36 37 38 39 40 41 42

class WhileOp : public framework::OperatorBase {
 public:
  WhileOp(const std::string &type, const framework::VariableNameMap &inputs,
          const framework::VariableNameMap &outputs,
          const framework::AttributeMap &attrs)
      : framework::OperatorBase(type, inputs, outputs, attrs) {}

43 44 45
 private:
  void RunImpl(const framework::Scope &scope,
               const platform::Place &dev_place) const override {
Y
Yang Yang(Tony) 已提交
46 47 48 49
    PADDLE_ENFORCE_NOT_NULL(scope.FindVar(Input(kCondition)));
    auto &cond = scope.FindVar(Input(kCondition))->Get<LoDTensor>();
    PADDLE_ENFORCE_EQ(cond.dims(), paddle::framework::make_ddim({1}));

D
dzhwinter 已提交
50
    framework::Executor executor(dev_place);
Y
Yu Yang 已提交
51
    auto *block = Attr<framework::BlockDesc *>(kStepBlock);
D
dzhwinter 已提交
52

Y
Yang Yang(Tony) 已提交
53 54 55 56 57
    auto *program = block->Program();

    auto step_scopes =
        scope.FindVar(Output(kStepScopes))->GetMutable<StepScopeVar>();

Y
Yang Yu 已提交
58 59
    PADDLE_ENFORCE(platform::is_cpu_place(cond.place()),
                   "Condition of while op must in CPU memory.");
X
Xin Pan 已提交
60

C
chengduo 已提交
61
    bool is_test = Attr<bool>("is_test");
X
Xin Pan 已提交
62
    auto ctx = executor.Prepare(*program, block->ID());
Y
Yang Yang(Tony) 已提交
63 64 65
    while (cond.data<bool>()[0]) {
      auto &current_scope = scope.NewScope();
      step_scopes->push_back(&current_scope);
C
chengduoZH 已提交
66
      executor.RunPreparedContext(ctx.get(), &current_scope, false, true, true);
C
chengduo 已提交
67 68 69
      if (is_test) {
        scope.DeleteScope(&current_scope);
      }
Y
Yang Yang(Tony) 已提交
70 71 72 73 74 75
    }
  }
};

class WhileOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
76
  void Make() override {
Y
Yang Yu 已提交
77
    AddInput(kX,
Y
Yang Yang(Tony) 已提交
78 79 80 81 82 83 84
             "A set of variables, which are required by operators inside the "
             "block of While Op.")
        .AsDuplicable();
    AddInput(
        kCondition,
        "(Bool) An scalar. When it's False, the While Op will be terminated.")
        .AsDuplicable();
Y
Yang Yang(Tony) 已提交
85
    AddOutput(kOutputs,
Y
Yang Yang(Tony) 已提交
86
              "A set of variables, which will be assigned with values "
Y
Yang Yang(Tony) 已提交
87
              "generated by the operators inside the block of While Op.")
Y
Yang Yang(Tony) 已提交
88 89 90 91 92
        .AsDuplicable();
    AddOutput(kStepScopes,
              "(StepScopeVar) A vector of local scope, which size equals the "
              "step number of While Op. The i'th scope storages temporary "
              "variables generated in the i'th step.");
Y
Yu Yang 已提交
93 94
    AddAttr<framework::BlockDesc *>(kStepBlock,
                                    "The step block inside WhileOp");
C
chengduo 已提交
95
    AddAttr<bool>("is_test", "True if in test phase.").SetDefault(false);
Y
Yang Yang(Tony) 已提交
96 97 98 99 100 101 102 103 104 105 106 107
    AddComment(R"DOC(
)DOC");
  }
};

class WhileGradOp : public framework::OperatorBase {
 public:
  WhileGradOp(const std::string &type, const framework::VariableNameMap &inputs,
              const framework::VariableNameMap &outputs,
              const framework::AttributeMap &attrs)
      : framework::OperatorBase(type, inputs, outputs, attrs) {}

108 109 110
 private:
  void RunImpl(const framework::Scope &scope,
               const platform::Place &dev_place) const override {
C
chengduo 已提交
111 112
    PADDLE_ENFORCE(!Attr<bool>("is_test"),
                   "GradOp is only callable when is_test is false");
113 114 115
    // get device context from pool
    platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
    auto &dev_ctx = *pool.Get(dev_place);
D
dzhwinter 已提交
116
    framework::Executor executor(dev_place);
Y
Yu Yang 已提交
117
    auto *block = Attr<framework::BlockDesc *>(kStepBlock);
Y
Yang Yang(Tony) 已提交
118
    auto *program = block->Program();
X
Xin Pan 已提交
119
    auto ctx = executor.Prepare(*program, block->ID());
Y
Yang Yang(Tony) 已提交
120 121 122 123

    auto *step_scopes =
        scope.FindVar(Input(kStepScopes))->GetMutable<StepScopeVar>();

Y
Yang Yang(Tony) 已提交
124 125 126 127 128 129
    auto outside_og_names = Inputs(framework::GradVarName(kOutputs));
    auto inside_og_names =
        Attr<std::vector<std::string>>("original_output_grad");

    PADDLE_ENFORCE_EQ(outside_og_names.size(), inside_og_names.size());

Y
Yang Yang(Tony) 已提交
130 131
    for (auto cur_scope_iter = step_scopes->rbegin();
         cur_scope_iter != step_scopes->rend(); ++cur_scope_iter) {
Y
Yang Yang(Tony) 已提交
132 133 134 135 136 137 138
      VLOG(3) << "Start backward at time_step "
              << cur_scope_iter - step_scopes->rbegin();
      framework::Scope &cur_scope = **cur_scope_iter;
      // Link OG from outside to inside
      for (size_t i = 0; i < outside_og_names.size(); ++i) {
        auto outside_og_name = outside_og_names[i];
        auto inside_og_name = inside_og_names[i];
139 140
        VLOG(8) << "Linking outside " << outside_og_name << " --> inside "
                << inside_og_name;
C
chengduo 已提交
141 142 143 144
        if (scope.FindVar(outside_og_name) == nullptr) {
          continue;
        }

145 146 147 148 149 150
        auto &og_outside =
            detail::Ref(scope.FindVar(outside_og_name),
                        "Cannot find Outside Gradient %s", outside_og_name);
        auto &og_inside =
            detail::Ref(cur_scope.Var(inside_og_name),
                        "Cannot find inside gradient %s", inside_og_name);
S
sneaxiy 已提交
151
        if (framework::IsType<framework::LoDTensor>(og_outside.Type())) {
Y
Yang Yang(Tony) 已提交
152 153 154 155 156
          auto &outside_tensor = og_outside.Get<framework::LoDTensor>();
          auto &inside_tensor =
              detail::Ref(og_inside.GetMutable<framework::LoDTensor>());
          inside_tensor.set_lod(outside_tensor.lod());
          inside_tensor.ShareDataWith(outside_tensor);
S
sneaxiy 已提交
157 158
        } else if (framework::IsType<framework::LoDTensorArray>(
                       og_outside.Type())) {
Y
Yang Yang(Tony) 已提交
159 160 161
          auto &outside_array = og_outside.Get<framework::LoDTensorArray>();
          auto &inside_array =
              detail::Ref(og_inside.GetMutable<framework::LoDTensorArray>());
162
          VLOG(8) << outside_og_name << " size = " << outside_array.size();
Y
Yang Yang(Tony) 已提交
163 164 165
          inside_array.resize(outside_array.size());

          for (size_t j = 0; j < inside_array.size(); ++j) {
166
            VLOG(8) << j << " " << outside_array[j].numel();
Y
Yang Yang(Tony) 已提交
167 168 169 170 171 172 173
            if (outside_array[j].numel() != 0) {
              inside_array[j].set_lod(outside_array[j].lod());
              inside_array[j].ShareDataWith(outside_array[j]);
            } else {
              PADDLE_ENFORCE_EQ(inside_array[j].numel(), 0);
            }
          }
C
chengduo 已提交
174 175
        } else {
          PADDLE_THROW("Currently only support LoDTensor and LoDTensorArray.");
Y
Yang Yang(Tony) 已提交
176 177
        }
      }
C
chengduoZH 已提交
178 179
      executor.RunPreparedContext(ctx.get(), *cur_scope_iter, false, true,
                                  true);
Y
Yang Yang(Tony) 已提交
180

C
chengduo 已提交
181 182 183
      // The Outputs(kXGRAD) contains the names of the gradient of parameters
      // and inputs.
      auto &pg_ig_names = Outputs(kXGRAD);
Y
Yang Yu 已提交
184
      auto &p_names = Inputs(kX);
C
chengduo 已提交
185 186 187
      PADDLE_ENFORCE_EQ(pg_ig_names.size(), p_names.size());
      for (size_t param_id = 0; param_id < pg_ig_names.size(); ++param_id) {
        if (pg_ig_names[param_id] == framework::kEmptyVarName) {
188
          continue;  // parameter doesn't have gradient
Y
Yang Yang(Tony) 已提交
189 190
        }
        auto inside_grad_name = framework::GradVarName(p_names[param_id]);
Y
Yang Yang(Tony) 已提交
191

C
chengduo 已提交
192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213
        // for some grad_op, their input doesn't have gradient,
        // for example lookup_table_grad_op, the input(Idx) doesn't have
        // gradient.
        auto pg_ig_var = cur_scope.FindVar(inside_grad_name);
        PADDLE_ENFORCE(pg_ig_var != nullptr);
        if (pg_ig_var->IsType<framework::LoDTensorArray>()) {
          auto pg_ig_lod_t_arr =
              pg_ig_var->GetMutable<framework::LoDTensorArray>();
          bool empty = true;
          for (auto &each : *pg_ig_lod_t_arr) {
            if (each.numel() != 0) {
              empty = false;
              break;
            }
          }
          if (empty) {
            LOG(WARNING) << pg_ig_names[param_id]
                         << " is not found in cur_scope.";
            continue;
          }
        }

Y
Yang Yang(Tony) 已提交
214
        //  // TODO(tonyyang-svail): Not sure we need the following
Y
Yang Yang(Tony) 已提交
215 216 217 218 219 220 221 222 223 224 225
        //  // If does not compute gradient of that variable inside rnn,
        //  just
        //  // continue
        //  if (local_var_names.find(inside_grad_name) ==
        //  local_var_names.end()) {
        //    continue;
        //  }

        // zero gradient variable in step 0
        if (cur_scope_iter == step_scopes->rbegin()) {
          auto *var = (*cur_scope_iter)->FindVar(inside_grad_name);
Y
Yang Yang(Tony) 已提交
226
          PADDLE_ENFORCE_NOT_NULL(var, "Can not find var %s", inside_grad_name);
C
chengduo 已提交
227 228 229 230 231
          PADDLE_ENFORCE(var->IsType<framework::LoDTensorArray>() ||
                             var->IsType<LoDTensor>(),
                         "Currently the type of var only can be LoDTensorArray "
                         "or LoDTensor.");

Y
Yang Yang(Tony) 已提交
232 233 234
          if (var->IsType<LoDTensor>()) {
            auto &inside_tensor = var->Get<framework::LoDTensor>();
            framework::AttributeMap attrs;
F
fengjiayi 已提交
235
            attrs["dtype"] = framework::ToDataType(inside_tensor.type());
Y
Yang Yang(Tony) 已提交
236 237 238
            attrs["shape"] = framework::vectorize2int(inside_tensor.dims());
            attrs["value"] = 0.0f;

C
chengduo 已提交
239
            auto var_name = pg_ig_names[param_id];
Y
Yang Yang(Tony) 已提交
240
            auto zero_op = framework::OpRegistry::CreateOp(
Y
Yiqun Liu 已提交
241
                "fill_constant", framework::VariableNameMap{},
242
                {{"Out", {var_name}}}, attrs);
D
dzhwinter 已提交
243
            zero_op->Run(scope, dev_place);
244 245 246
            scope.FindVar(var_name)
                ->GetMutable<framework::LoDTensor>()
                ->set_lod(inside_tensor.lod());
Y
Yang Yang(Tony) 已提交
247 248
          }
        }
Y
Yang Yang(Tony) 已提交
249
        auto new_inside_name = cur_scope.Rename(inside_grad_name);
Y
Yang Yang(Tony) 已提交
250
        auto sum_op = framework::OpRegistry::CreateOp(
C
chengduo 已提交
251 252
            "sum", {{"X", {pg_ig_names[param_id], new_inside_name}}},
            {{"Out", {pg_ig_names[param_id]}}},
253
            framework::AttributeMap{{"use_mkldnn", {false}}});
D
dzhwinter 已提交
254
        sum_op->Run(cur_scope, dev_place);
Y
Yang Yang(Tony) 已提交
255
        cur_scope.Rename(new_inside_name, inside_grad_name);
Y
Yang Yang(Tony) 已提交
256
      }
257 258
      dev_ctx.Wait();
      const_cast<framework::Scope &>(scope).DeleteScope(&cur_scope);
Y
Yang Yang(Tony) 已提交
259 260 261 262 263 264 265 266 267
    }
  }
};

class WhileGradOpDescMaker : public framework::SingleGradOpDescMaker {
 public:
  using framework::SingleGradOpDescMaker::SingleGradOpDescMaker;

 protected:
Y
Yu Yang 已提交
268
  std::unique_ptr<framework::OpDesc> Apply() const override {
F
Update  
fengjiayi 已提交
269 270 271 272 273 274 275
    auto *while_grad = new framework::OpDesc();
    while_grad->SetType("while_grad");
    while_grad->SetInput(kX, Input(kX));
    while_grad->SetInput(kOutputs, Output(kOutputs));
    while_grad->SetInput(kStepScopes, Output(kStepScopes));

    auto *grad_block = this->grad_block_[0];
Y
Yu Yang 已提交
276 277
    auto *fwd_block = grad_block->ForwardBlock();
    auto *parent_block = grad_block->ParentBlock();
278 279 280

    // Not all of IGs will be generated by inner gradient operators of while op.
    // Ignore IGs that is not generated by the inside block.
F
Update  
fengjiayi 已提交
281 282 283 284
    std::unordered_set<std::string> inner_op_outputs;
    for (const auto *op : grad_block->AllOps()) {
      for (auto &oname : op->OutputArgumentNames()) {
        inner_op_outputs.insert(oname);
285 286
      }
    }
F
Update  
fengjiayi 已提交
287
    auto igs = InputGrad(kX, /*do not drop empty gradient*/ false);
288
    for (auto &each_ig : igs) {
F
Update  
fengjiayi 已提交
289
      if (inner_op_outputs.find(each_ig) == inner_op_outputs.end()) {
290
        VLOG(8) << "Ignore " << each_ig;
291 292 293
        each_ig = framework::kEmptyVarName;
      }
    }
F
Update  
fengjiayi 已提交
294
    while_grad->SetOutput(framework::GradVarName(kX), igs);
Y
Yang Yang(Tony) 已提交
295 296 297 298

    // OG should be re-calculated by step blocks, since many outputs of while op
    // do not need to calculate gradients.
    std::unordered_set<std::string> block_ins;
F
fengjiayi 已提交
299 300 301 302 303 304 305
    block_ins.reserve(Input(kX).size() + Output(kOutputs).size());
    for (auto &p : Input(kX)) {
      block_ins.insert(p);
    }
    for (auto &o : Output(kOutputs)) {
      block_ins.insert(o);
    }
Y
Yu Yang 已提交
306
    std::unordered_set<std::string> output_grads;
F
Update  
fengjiayi 已提交
307 308 309 310
    for (const auto *op : grad_block->AllOps()) {
      for (auto &input_name : op->InputArgumentNames()) {
        // If the input of Op has been recorded or is generated by the forward
        // block, do not make it as input again.
Y
Yu Yang 已提交
311 312 313

        // The input is located in I/O or other op's outputs or the variable is
        // located in grad_block's parents
F
Update  
fengjiayi 已提交
314
        if (block_ins.find(input_name) != block_ins.end() ||
Y
Yu Yang 已提交
315 316
            (fwd_block->FindVarRecursive(input_name) != nullptr ||
             parent_block->FindVarRecursive(input_name) != nullptr)) {
Y
Yang Yang(Tony) 已提交
317 318
          continue;
        }
C
chengduo 已提交
319

Y
Yu Yang 已提交
320
        output_grads.insert(input_name);
Y
Yang Yang(Tony) 已提交
321
      }
F
Update  
fengjiayi 已提交
322
      for (auto &output_name : op->OutputArgumentNames()) {
Y
Yang Yang(Tony) 已提交
323
        block_ins.insert(output_name);
Y
Yang Yang(Tony) 已提交
324 325
      }
    }
Y
Yang Yang(Tony) 已提交
326

Y
Yu Yang 已提交
327 328 329 330 331
    std::vector<std::string> output_grads_list;
    output_grads_list.resize(output_grads.size());
    std::copy(output_grads.begin(), output_grads.end(),
              output_grads_list.begin());
    while_grad->SetInput(framework::GradVarName(kOutputs), output_grads_list);
F
Update  
fengjiayi 已提交
332 333

    while_grad->SetAttrMap(this->Attrs());
A
Abhinav Arora 已提交
334
    while_grad->SetBlockAttr(kStepBlock, grad_block);
Y
Yang Yang(Tony) 已提交
335 336
    // record the original output gradient names, since the gradient name of
    // while operator could be renamed.
Y
Yu Yang 已提交
337
    while_grad->SetAttr("original_output_grad", output_grads_list);
Y
Yang Yang(Tony) 已提交
338

F
Update  
fengjiayi 已提交
339
    return std::unique_ptr<framework::OpDesc>(while_grad);
Y
Yang Yang(Tony) 已提交
340 341 342
  }
};

Y
Yang Yang(Tony) 已提交
343 344
class WhileGradOpVarTypeInference : public framework::VarTypeInference {
 public:
Y
Yu Yang 已提交
345 346
  void operator()(const framework::OpDesc &op_desc,
                  framework::BlockDesc *block) const override {
Y
Yang Yu 已提交
347
    auto p_names = op_desc.Input(kX);
C
chengduo 已提交
348
    auto pg_ig_names = op_desc.Output(framework::GradVarName(kX));
Y
Yang Yang(Tony) 已提交
349 350 351

    for (size_t i = 0; i < p_names.size(); ++i) {
      auto &p_var = detail::Ref(block->FindVarRecursive(p_names[i]));
C
chengduo 已提交
352
      auto *g_var = block->FindVarRecursive(pg_ig_names[i]);
Y
Yang Yang(Tony) 已提交
353
      if (g_var != nullptr) {  // Gradient could be @EMPTY@
C
chengduo 已提交
354
        VLOG(5) << "Setting " << pg_ig_names[i] << " following " << p_names[i]
Y
Yang Yang(Tony) 已提交
355 356 357 358 359 360 361 362 363 364 365
                << " type: " << p_var.GetType();
        g_var->SetType(p_var.GetType());
        g_var->SetDataType(p_var.GetDataType());
      }
    }
  }
};

class WhileGradOpShapeInference : public framework::InferShapeBase {
 public:
  void operator()(framework::InferShapeContext *ctx) const override {
Y
Yang Yu 已提交
366 367
    ctx->HasInputs(kX);
    ctx->HasOutputs(framework::GradVarName(kX));
Y
Yang Yang(Tony) 已提交
368 369 370
    ctx->HasInputs(kOutputs);
    ctx->HasInputs(framework::GradVarName(kOutputs));

Y
Yang Yu 已提交
371
    auto p_names = ctx->Inputs(kX);
C
chengduo 已提交
372
    auto pg_ig_names = ctx->Outputs(kXGRAD);
Y
Yang Yu 已提交
373
    auto var_types = ctx->GetInputsVarType(kX);
Y
Yang Yang(Tony) 已提交
374 375 376
    std::vector<std::string> names_to_set;
    std::vector<framework::DDim> dims_to_set;
    for (size_t i = 0; i < p_names.size(); ++i) {
C
chengduo 已提交
377
      if (pg_ig_names[i] == framework::kEmptyVarName) {
Y
Yang Yang(Tony) 已提交
378 379
        continue;
      }
Y
Yang Yu 已提交
380
      auto dims = ctx->GetInputsElementDim(kX, i);
381
      if (var_types[i] == framework::proto::VarType::LOD_TENSOR) {
C
chengduo 已提交
382
        names_to_set.push_back(pg_ig_names[i]);
F
fengjiayi 已提交
383
        dims_to_set.push_back(dims);
384
      } else if (var_types[i] == framework::proto::VarType::LOD_TENSOR_ARRAY) {
Y
Yang Yang(Tony) 已提交
385
        // not sure how to set the dim of LOD_TENSOR_ARRAY
C
chengduo 已提交
386
        names_to_set.push_back(pg_ig_names[i]);
F
fengjiayi 已提交
387
        dims_to_set.push_back(dims);
Y
Yang Yang(Tony) 已提交
388 389 390 391 392 393
      }
    }
    ctx->SetDims(names_to_set, dims_to_set);
  }
};

Y
Yang Yang(Tony) 已提交
394 395 396 397 398 399
}  // namespace operators
}  // namespace paddle

REGISTER_OPERATOR(while, paddle::operators::WhileOp,
                  paddle::operators::WhileOpMaker,
                  paddle::operators::WhileGradOpDescMaker);
Y
Yang Yang(Tony) 已提交
400 401 402
REGISTER_OPERATOR(while_grad, paddle::operators::WhileGradOp,
                  paddle::operators::WhileGradOpShapeInference,
                  paddle::operators::WhileGradOpVarTypeInference);