shrink_rnn_memory_op.cc 7.4 KB
Newer Older
L
Luo Tao 已提交
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
Y
Yang Yu 已提交
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
Yang Yu 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
Y
Yang Yu 已提交
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. */
Y
Yi Wang 已提交
14 15 16 17
#include "paddle/fluid/framework/lod_rank_table.h"
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/operators/array_operator.h"
#include "paddle/fluid/operators/math/math_function.h"
Y
Yang Yu 已提交
18 19 20 21

namespace paddle {
namespace operators {

Y
Yang Yu 已提交
22
class ShrinkRNNMemoryOp : public ArrayOp {
Y
Yang Yu 已提交
23
 public:
Y
Yang Yu 已提交
24 25 26 27
  ShrinkRNNMemoryOp(const std::string &type,
                    const framework::VariableNameMap &inputs,
                    const framework::VariableNameMap &outputs,
                    const framework::AttributeMap &attrs)
Y
Yang Yu 已提交
28 29 30
      : ArrayOp(type, inputs, outputs, attrs) {}

  void Run(const framework::Scope &scope,
D
dzhwinter 已提交
31
           const platform::Place &place) const override {
Y
Yang Yu 已提交
32 33 34
    auto *x_var = scope.FindVar(Input("X"));
    PADDLE_ENFORCE(x_var != nullptr, "Input X must be set");
    auto &x_tensor = x_var->Get<framework::LoDTensor>();
D
dzhwinter 已提交
35
    size_t offset = this->GetOffset(scope, place);
Y
Yang Yu 已提交
36 37 38 39
    auto *rank_table_var = scope.FindVar(Input("RankTable"));
    PADDLE_ENFORCE(rank_table_var != nullptr, "RankTable must be set");
    auto &rank_table = rank_table_var->Get<framework::LoDRankTable>();

Y
Yang Yu 已提交
40 41 42 43 44 45
    auto &rank_items = rank_table.items();
    int dst_num_rows =
        std::lower_bound(rank_items.begin(), rank_items.end(), offset,
                         [](const framework::LoDRankTable::TableItem &a,
                            size_t b) { return a.length > b; }) -
        rank_items.begin();
Y
Yang Yu 已提交
46 47

    auto *out_var = scope.FindVar(Output("Out"));
48
    PADDLE_ENFORCE(out_var != nullptr, "Output(Out) must be set.");
Y
Yang Yu 已提交
49
    auto &out_tensor = *out_var->GetMutable<framework::LoDTensor>();
Y
yangyaming 已提交
50 51

    size_t height = dst_num_rows;
Y
yangyaming 已提交
52

53 54 55 56 57
    // do shrink for the top level LoD
    if (x_tensor.lod().size() > 0 &&
        x_tensor.lod()[0].size() > static_cast<size_t>(dst_num_rows)) {
      auto lod_offset = framework::GetSubLoDAndAbsoluteOffset(x_tensor.lod(), 0,
                                                              dst_num_rows, 0);
Y
yangyaming 已提交
58
      height = lod_offset.second.second;
Y
yangyaming 已提交
59
      auto out_lod = out_tensor.mutable_lod();
60
      framework::AppendLoD(out_lod, lod_offset.first);
Y
yangyaming 已提交
61 62
    }

Y
Yang Yu 已提交
63
    if (dst_num_rows != 0) {
Y
yangyaming 已提交
64
      out_tensor.ShareDataWith(x_tensor.Slice(0, height));
Y
Yang Yu 已提交
65 66 67 68
    }
  }
};

Y
Yang Yu 已提交
69
class ShrinkRNNMemoryOpProtoMaker : public framework::OpProtoAndCheckerMaker {
Y
Yang Yu 已提交
70
 public:
71
  ShrinkRNNMemoryOpProtoMaker(OpProto *proto, OpAttrChecker *op_checker)
Y
Yang Yu 已提交
72
      : OpProtoAndCheckerMaker(proto, op_checker) {
73 74 75 76 77 78
    AddInput("X", "(LoDTensor) The RNN step memory to be shrinked.");
    AddInput("RankTable", "(LoDRankTable) The lod_rank_table of dynamic RNN.");
    AddInput("I",
             "(LoDTensor) The step index. The RNN step memory 'X' will be "
             "shrinked to match the size of the input of the index'th step.");
    AddOutput("Out", "(LoDTensor) The shrinked RNN step memory.");
79 80 81 82 83 84 85 86 87 88 89
    AddComment(R"DOC(
This operator is used to shrink output batch of memory defined in dynamic RNN.

Dynamic RNN is able to handle variable-length sequences, in which, sequences in
a mini-batch are sorted by their lengths first. After that, the longest sequence
becomes the first one in the sorted batch, followed by the second longest, the
third longest, and so on. Dynamic RNN then slices a batch input timestep by
timestep from the sorted input. Once any sequence in the input batch reaches its
end, memory defined in dynamicRNN has to shrink its outputs to adapt to the input
batch size for the next time step.
)DOC");
Y
Yang Yu 已提交
90 91 92
  }
};

Y
Yang Yu 已提交
93
class ShrinkRNNMemoryInferShape : public framework::InferShapeBase {
Y
Yang Yu 已提交
94 95 96 97 98 99 100 101 102
 public:
  void operator()(framework::InferShapeContext *context) const override {
    PADDLE_ENFORCE(context->HasInput("X"));
    PADDLE_ENFORCE(context->HasInput("I"));
    PADDLE_ENFORCE(context->HasInput("RankTable"));
    context->SetOutputDim("Out", context->GetInputDim("X"));
  }
};

Y
Yang Yu 已提交
103
class ShrinkRNNMemoryGradOp : public ArrayOp {
Y
Yang Yu 已提交
104
 public:
Y
Yang Yu 已提交
105 106 107 108
  ShrinkRNNMemoryGradOp(const std::string &type,
                        const framework::VariableNameMap &inputs,
                        const framework::VariableNameMap &outputs,
                        const framework::AttributeMap &attrs)
Y
Yang Yu 已提交
109 110 111
      : ArrayOp(type, inputs, outputs, attrs) {}

  void Run(const framework::Scope &scope,
D
dzhwinter 已提交
112
           const platform::Place &place) const override {
Y
Yang Yu 已提交
113
    auto *dout_var = scope.FindVar(Input(framework::GradVarName("Out")));
Y
Yang Yu 已提交
114
    auto *dx_var = scope.FindVar(Output(framework::GradVarName("X")));
Y
Yang Yu 已提交
115 116 117 118 119 120 121 122 123
    PADDLE_ENFORCE(dx_var != nullptr, "Input Gradient should not be nullptr");
    auto *x_var = scope.FindVar(Input("X"));
    PADDLE_ENFORCE(x_var != nullptr);

    auto &x_tensor = x_var->Get<framework::LoDTensor>();
    auto &dx_tensor = *dx_var->GetMutable<framework::LoDTensor>();
    dx_tensor.Resize(x_tensor.dims());
    dx_tensor.mutable_data(x_tensor.place(), x_tensor.type());

D
dzhwinter 已提交
124
    // get device context from pool
Y
Yang Yu 已提交
125 126
    platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
    auto &dev_ctx = *pool.Get(place);
D
dzhwinter 已提交
127

Y
Yang Yu 已提交
128 129 130 131 132
    if (dout_var == nullptr) {  // dx_tensor fill zero
      math::set_constant(dev_ctx, &dx_tensor, 0.0f);
    } else {
      auto &dout_tensor = dout_var->Get<framework::LoDTensor>();
      auto height = dout_tensor.dims()[0];
D
dzhwinter 已提交
133
      auto slice = dx_tensor.Slice(0, static_cast<int>(height));
134
      framework::Copy(dout_tensor, dout_tensor.place(), dev_ctx, &slice);
Y
Refine  
Yang Yu 已提交
135
      if (dx_tensor.dims()[0] > height) {
Y
Yang Yu 已提交
136
        auto rest_tensor = dx_tensor.Slice(
Y
Refine  
Yang Yu 已提交
137
            static_cast<int>(height), static_cast<int>(dx_tensor.dims()[0]));
Y
Yang Yu 已提交
138 139 140
        math::set_constant(dev_ctx, &rest_tensor, 0.0f);
      }
    }
141
    dx_tensor.set_lod(x_tensor.lod());
Y
Yang Yu 已提交
142 143 144
  }
};

Y
Yang Yu 已提交
145
class ShrinkRNNMemoryGradInferShape : public framework::InferShapeBase {
Y
Yang Yu 已提交
146 147 148 149 150 151
 public:
  void operator()(framework::InferShapeContext *context) const override {
    PADDLE_ENFORCE(context->HasInput("X"));
    PADDLE_ENFORCE(context->HasOutput(framework::GradVarName("X")));
    context->SetOutputDim(framework::GradVarName("X"),
                          context->GetInputDim("X"));
152
    context->ShareLoD("X", framework::GradVarName("X"));
Y
Yang Yu 已提交
153 154 155
  }
};

Y
Yang Yu 已提交
156
class ShrinkRNNGradOpMaker : public framework::SingleGradOpDescMaker {
Y
Yang Yu 已提交
157 158 159 160
 public:
  using framework::SingleGradOpDescMaker::SingleGradOpDescMaker;

 protected:
Y
Yu Yang 已提交
161 162
  std::unique_ptr<framework::OpDesc> Apply() const override {
    auto *op = new framework::OpDesc();
Y
Yang Yu 已提交
163
    op->SetType("shrink_rnn_memory_grad");
Y
Yang Yu 已提交
164 165 166 167
    op->SetInput("X", Input("X"));
    op->SetInput(framework::GradVarName("Out"), OutputGrad("Out"));
    op->SetOutput(framework::GradVarName("X"), InputGrad("X"));
    op->SetAttrMap(Attrs());
Y
Yu Yang 已提交
168
    return std::unique_ptr<framework::OpDesc>(op);
Y
Yang Yu 已提交
169 170 171 172 173 174 175
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
Y
Yang Yu 已提交
176 177 178 179 180
REGISTER_OPERATOR(shrink_rnn_memory, ops::ShrinkRNNMemoryOp,
                  ops::ShrinkRNNMemoryInferShape,
                  ops::ShrinkRNNMemoryOpProtoMaker, ops::ShrinkRNNGradOpMaker);
REGISTER_OPERATOR(shrink_rnn_memory_grad, ops::ShrinkRNNMemoryGradOp,
                  ops::ShrinkRNNMemoryGradInferShape);