recurrent_op.cc 9.3 KB
Newer Older
Y
Yan Chunwei 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

   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. */

15
#include "paddle/operators/recurrent_op.h"
Y
Yan Chunwei 已提交
16 17 18 19 20

#include <cstring>
#include <sstream>

#include "paddle/framework/op_registry.h"
Y
Yan Chunwei 已提交
21
#include "paddle/operators/net_op.h"
Y
Yan Chunwei 已提交
22 23 24 25

namespace paddle {
namespace operators {

D
dongzhihong 已提交
26 27 28 29
using Scope = framework::Scope;
using Variable = framework::Variable;
using Tensor = framework::Tensor;

Y
Yu Yang 已提交
30 31
void RecurrentAlgorithm::InferShape(const Scope& scope) const {
  seq_len_ = scope.FindVar((arg_->inlinks[0]).external)
Y
Yan Chunwei 已提交
32 33 34
                 ->GetMutable<Tensor>()
                 ->dims()[0];
  CreateScopes(scope);
35
  auto step_scopes = GetStepScopes(scope);
36 37
  rnn::SegmentInputs(step_scopes, arg_->inlinks, seq_len_,
                     true /*infer_shape_mode*/);
D
dangqingqing 已提交
38
  InitMemories(step_scopes[0], true /*infer_shape_mode*/);
Y
Yu Yang 已提交
39
  Variable* net = scope.FindVar(arg_->step_net);
Y
Yan Chunwei 已提交
40
  PADDLE_ENFORCE(net != nullptr, "failed to get step net");
Y
Yan Chunwei 已提交
41

Y
Yan Chunwei 已提交
42 43
  for (size_t i = 0; i < seq_len_; i++) {
    if (i > 0) {
44 45
      rnn::LinkMemories(step_scopes, arg_->memories, i, -1,
                        true /*infer_shape_mode*/);
Y
Yan Chunwei 已提交
46
    }
Y
Yu Yang 已提交
47
    net->GetMutable<NetOp>()->InferShape(*step_scopes[i]);
Y
Yan Chunwei 已提交
48
  }
49 50
  rnn::ConcatOutputs(step_scopes, arg_->outlinks, seq_len_,
                     true /*infer_shape_mode*/);
Y
Yan Chunwei 已提交
51 52
}

Y
Yu Yang 已提交
53
void RecurrentAlgorithm::Run(const Scope& scope,
Y
Yan Chunwei 已提交
54 55
                             const platform::DeviceContext& dev_ctx) const {
  auto step_scopes = GetStepScopes(scope);
56 57
  rnn::SegmentInputs(step_scopes, arg_->inlinks, seq_len_,
                     false /*infer_shape_mode*/);
D
dangqingqing 已提交
58
  InitMemories(step_scopes[0], false /*infer_shape_mode*/);
Y
Yu Yang 已提交
59
  Variable* net = scope.FindVar(arg_->step_net);
D
dangqingqing 已提交
60

Y
Yan Chunwei 已提交
61
  for (size_t step_id = 0; step_id < seq_len_; step_id++) {
Y
Yan Chunwei 已提交
62
    // create output alias variables
Y
Yan Chunwei 已提交
63
    if (step_id > 0) {
64 65
      rnn::LinkMemories(step_scopes, arg_->memories, step_id, -1,
                        false /*infer_shape_mode*/);
Y
Yan Chunwei 已提交
66
    }
Y
Yu Yang 已提交
67
    net->GetMutable<NetOp>()->Run(*step_scopes[step_id], dev_ctx);
Y
Yan Chunwei 已提交
68
  }
69 70
  rnn::ConcatOutputs(step_scopes, arg_->outlinks, seq_len_,
                     false /*infer_shape_mode*/);
Y
Yan Chunwei 已提交
71 72
}

Y
Yu Yang 已提交
73
void RecurrentAlgorithm::CreateScopes(const Scope& scope) const {
Y
Yan Chunwei 已提交
74
  // TODO(superjom) Only two scopes are needed for inference, this case will be
Y
Yan Chunwei 已提交
75
  // supported later.
Y
Yan Chunwei 已提交
76 77 78 79 80 81 82 83 84 85
  auto step_scopes_var = scope.FindVar(arg_->step_scopes);
  PADDLE_ENFORCE(step_scopes_var != nullptr, "");
  auto step_scopes = step_scopes_var->GetMutable<std::vector<Scope*>>();

  // Now all variables in scope must be created outside of op.
  auto net_var = scope.FindVar(arg_->step_net);
  PADDLE_ENFORCE(net_var != nullptr, "no stepnet called %s in scope",
                 arg_->step_net);
  auto net_op = net_var->GetMutable<NetOp>();
  PADDLE_ENFORCE(!net_op->outputs_.empty(), "net_op has no outputs");
Y
Yan Chunwei 已提交
86 87 88

  if (seq_len_ > step_scopes->size()) {
    for (size_t i = step_scopes->size(); i < seq_len_; ++i) {
Y
Yu Yang 已提交
89
      auto& step_scope = scope.NewScope();
Y
Yan Chunwei 已提交
90

Y
Yan Chunwei 已提交
91
      // create step net's temp inputs
Y
Yan Chunwei 已提交
92
      for (auto& input : net_op->inputs_) {
93
        // the weight are located in parent scope
Y
Yu Yang 已提交
94 95 96 97 98
        for (auto& var_name : input.second) {
          if (!step_scope.FindVar(var_name)) {
            step_scope.NewVar(var_name)->GetMutable<Tensor>();
          }
        }
Y
Yan Chunwei 已提交
99
      }
Y
Yan Chunwei 已提交
100 101
      // create stepnet's outputs
      for (const auto& output : net_op->outputs_) {
Y
Yu Yang 已提交
102 103 104
        for (auto& var_name : output.second) {
          step_scope.NewVar(var_name);
        }
Y
Yan Chunwei 已提交
105
      }
Y
Yu Yang 已提交
106
      step_scopes->emplace_back(&step_scope);
Y
Yan Chunwei 已提交
107 108 109 110
    }
  }
}

D
dangqingqing 已提交
111
void RecurrentAlgorithm::InitMemories(Scope* step_scope,
D
dangqingqing 已提交
112
                                      bool infer_shape_mode) const {
Y
Yan Chunwei 已提交
113
  for (auto& attr : arg_->memories) {
114
    Tensor* pre_mem = step_scope->NewVar(attr.pre_var)->GetMutable<Tensor>();
Y
Yu Yang 已提交
115
    PADDLE_ENFORCE(step_scope->FindVar(attr.boot_var) != nullptr,
116
                   "memory [%s]'s boot variable [%s] not exists", attr.var,
Y
Yan Chunwei 已提交
117
                   attr.boot_var);
118
    Tensor* boot_mem = step_scope->FindVar(attr.boot_var)->GetMutable<Tensor>();
D
dangqingqing 已提交
119
    if (infer_shape_mode) {
120
      pre_mem->Resize(boot_mem->dims());
Y
Yan Chunwei 已提交
121
      PADDLE_ENFORCE_EQ(pre_mem->dims().size(), 2);
122 123 124
    } else {
      pre_mem->ShareDataWith<float>(*boot_mem);
    }
Y
Yan Chunwei 已提交
125 126 127
  }
}

128 129 130 131 132 133 134 135 136
const rnn::ArgumentName RecurrentOp::kArgName{
    "step_net", "step_scopes",  "inlinks",
    "outlinks", "inlink_alias", "outlink_alias",
    "memories", "pre_memories", "boot_memories"};

const rnn::ArgumentName RecurrentGradientOp::kArgName{
    "step_net",    "step_scopes",  "outlink@grad",
    "inlink@grad", "inlink_alias", "outlink_alias",
    "memories",    "pre_memories", "boot_memories@grad"};
Y
Yan Chunwei 已提交
137 138 139 140 141 142 143 144

void RecurrentOp::Init() {
  OperatorBase::Init();
  std::unique_ptr<rnn::Argument> arg(new rnn::Argument());
  rnn::InitArgument(kArgName, arg.get(), *this);
  alg_.Init(std::move(arg));
}

D
dongzhihong 已提交
145 146
class RecurrentAlgorithmProtoAndCheckerMaker
    : public framework::OpProtoAndCheckerMaker {
147
 public:
D
dongzhihong 已提交
148 149
  RecurrentAlgorithmProtoAndCheckerMaker(framework::OpProto* proto,
                                         framework::OpAttrChecker* op_checker)
Y
Yan Chunwei 已提交
150 151 152
      : OpProtoAndCheckerMaker(proto, op_checker) {
    const auto& name = RecurrentOp::kArgName;
    // inputs and outputs stored in proto
D
dangqingqing 已提交
153 154
    AddInput(name.inlinks,
             "the inputs that need to be segmented for each step.")
Y
Yu Yang 已提交
155 156 157
        .SetMultiple();
    AddInput(name.boot_memories, "variables to initialize memories.")
        .SetMultiple();
Y
Yan Chunwei 已提交
158 159
    AddInput(name.step_net, "network shared by all steps.");

D
dangqingqing 已提交
160
    AddOutput(name.outlinks, "the outputs that need to concated for all steps.")
Y
Yu Yang 已提交
161
        .SetMultiple();
Y
Yan Chunwei 已提交
162 163 164 165 166 167 168 169 170 171 172 173 174 175
    AddOutput(name.step_scopes, "step scopes");

    // Attributes stored in AttributeMap
    AddAttr<std::vector<std::string>>(name.inlink_alias, "alias of inlinks");
    AddAttr<std::vector<std::string>>(name.outlink_alias, "alias of outlinks");
    AddAttr<std::vector<std::string>>(name.pre_memories,
                                      "names of pre-memories");
    AddAttr<std::vector<std::string>>(name.memories, "names of memories");

    AddComment("This is a recurrent group operator.");
  }
};

void RecurrentGradientAlgorithm::Run(
Y
Yu Yang 已提交
176
    const Scope& scope, const platform::DeviceContext& dev_ctx) const {
Y
Yan Chunwei 已提交
177
  auto step_scopes = GetStepScopes(scope);
178 179
  rnn::SegmentInputs(step_scopes, arg_->inlinks, seq_len_,
                     false /*infer_shape_mode*/);
Y
Yu Yang 已提交
180
  Variable* net = scope.FindVar(arg_->step_net);
Y
Yan Chunwei 已提交
181 182 183
  PADDLE_ENFORCE(net != nullptr, "failed to get step net");
  for (int step_id = seq_len_ - 1; step_id >= 0; --step_id) {
    if (static_cast<size_t>(step_id) != seq_len_ - 1) {
184 185
      rnn::LinkMemories(step_scopes, arg_->memories, step_id, 1,
                        false /*infer_shape_mode*/);
Y
Yan Chunwei 已提交
186
    }
Y
Yu Yang 已提交
187
    net->GetMutable<NetOp>()->Run(*step_scopes[step_id], dev_ctx);
Y
Yan Chunwei 已提交
188
  }
189
  LinkBootMemoryGradients(step_scopes[0], false);
190 191
  rnn::ConcatOutputs(step_scopes, arg_->outlinks, seq_len_,
                     false /*infer_shape_mode*/);
Y
Yan Chunwei 已提交
192 193 194
}

void RecurrentGradientAlgorithm::LinkBootMemoryGradients(
D
dangqingqing 已提交
195
    Scope* step_scope, bool infer_shape_mode) const {
Y
Yan Chunwei 已提交
196
  for (auto& attr : arg_->memories) {
D
dangqingqing 已提交
197
    PADDLE_ENFORCE(step_scope->FindVar(attr.var) != nullptr,
198
                   "memory variable [%s] does not exists", attr.var);
Y
Yu Yang 已提交
199
    PADDLE_ENFORCE(step_scope->FindVar(attr.boot_var) != nullptr,
200
                   "boot variable [%s] does not exists", attr.boot_var);
D
dangqingqing 已提交
201
    Tensor* mem_grad = step_scope->NewVar(attr.var)->GetMutable<Tensor>();
Y
Yan Chunwei 已提交
202
    Tensor* boot_mem_grad =
203
        step_scope->NewVar(attr.boot_var)->GetMutable<Tensor>();
D
dangqingqing 已提交
204
    if (infer_shape_mode) {
205 206 207 208
      boot_mem_grad->Resize(mem_grad->dims());
    } else {
      boot_mem_grad->ShareDataWith<float>(*mem_grad);
    }
Y
Yan Chunwei 已提交
209 210 211
  }
}

Y
Yu Yang 已提交
212 213
void RecurrentGradientAlgorithm::InferShape(const Scope& scope) const {
  seq_len_ = scope.FindVar((arg_->inlinks[0]).external)
Y
Yan Chunwei 已提交
214 215 216
                 ->GetMutable<Tensor>()
                 ->dims()[0];
  auto step_scopes = GetStepScopes(scope);
217 218
  rnn::SegmentInputs(step_scopes, arg_->inlinks, seq_len_,
                     true /*infer_shape_mode*/);
Y
Yu Yang 已提交
219
  Variable* net = scope.FindVar(arg_->step_net);
Y
Yan Chunwei 已提交
220 221 222
  PADDLE_ENFORCE(net != nullptr, "failed to get step net");
  for (int step_id = seq_len_ - 1; step_id >= 0; --step_id) {
    if (static_cast<size_t>(step_id) != seq_len_ - 1) {
223 224
      rnn::LinkMemories(step_scopes, arg_->memories, step_id, 1,
                        true /*infer_shape_mode*/);
Y
Yan Chunwei 已提交
225
    }
Y
Yu Yang 已提交
226
    net->GetMutable<NetOp>()->InferShape(*step_scopes[step_id]);
Y
Yan Chunwei 已提交
227
  }
228 229
  rnn::ConcatOutputs(step_scopes, arg_->outlinks, seq_len_,
                     true /*infer_shape_mode*/);
D
dangqingqing 已提交
230
  LinkBootMemoryGradients(step_scopes[0], true /*infer_shape_mode*/);
Y
Yan Chunwei 已提交
231 232 233 234 235 236 237 238 239 240 241 242
}

void RecurrentGradientOp::Init() {
  OperatorBase::Init();
  std::unique_ptr<rnn::Argument> arg(new rnn::Argument());
  rnn::InitArgument(kArgName, arg.get(), *this);
  alg_.Init(std::move(arg));
}

}  // namespace operators
}  // namespace paddle

243
REGISTER_OP(recurrent_op, paddle::operators::RecurrentOp,
Y
Yan Chunwei 已提交
244
            paddle::operators::RecurrentAlgorithmProtoAndCheckerMaker);