shrink_rnn_memory_op.cc 7.5 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
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
      : ArrayOp(type, inputs, outputs, attrs) {}

30 31 32
 private:
  void RunImpl(const framework::Scope &scope,
               const platform::Place &place) const override {
Y
Yang Yu 已提交
33 34 35
    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 已提交
36
    size_t offset = this->GetOffset(scope, place);
Y
Yang Yu 已提交
37 38 39 40
    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 已提交
41 42 43 44 45 46
    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 已提交
47 48

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

    size_t height = dst_num_rows;
Y
yangyaming 已提交
53

54 55 56
    // 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)) {
57 58 59 60 61
      auto lod_offset = framework::GetSubLoDAndAbsoluteOffset(x_tensor.lod(), 0,
                                                              dst_num_rows, 0);
      height = lod_offset.second.second;
      auto out_lod = out_tensor.mutable_lod();
      framework::AppendLoD(out_lod, lod_offset.first);
Y
yangyaming 已提交
62 63
    }

64
    if (dst_num_rows != 0) {
D
dzhwinter 已提交
65 66 67 68
      out_tensor.mutable_data(place, x_tensor.type());
      auto dev_ctx = platform::DeviceContextPool::Instance().Get(place);
      framework::TensorCopy(x_tensor.Slice(0, height), place, *dev_ctx,
                            &out_tensor);
Y
Yang Yu 已提交
69 70 71 72
    }
  }
};

Y
Yang Yu 已提交
73
class ShrinkRNNMemoryOpProtoMaker : public framework::OpProtoAndCheckerMaker {
Y
Yang Yu 已提交
74
 public:
Y
Yu Yang 已提交
75
  void Make() override {
76 77 78 79 80 81
    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.");
82 83 84 85 86 87 88 89 90 91 92
    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 已提交
93 94 95
  }
};

Y
Yang Yu 已提交
96
class ShrinkRNNMemoryInferShape : public framework::InferShapeBase {
Y
Yang Yu 已提交
97 98 99 100 101 102 103 104 105
 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 已提交
106
class ShrinkRNNMemoryGradOp : public ArrayOp {
Y
Yang Yu 已提交
107
 public:
Y
Yang Yu 已提交
108 109 110 111
  ShrinkRNNMemoryGradOp(const std::string &type,
                        const framework::VariableNameMap &inputs,
                        const framework::VariableNameMap &outputs,
                        const framework::AttributeMap &attrs)
Y
Yang Yu 已提交
112 113
      : ArrayOp(type, inputs, outputs, attrs) {}

114 115 116
 private:
  void RunImpl(const framework::Scope &scope,
               const platform::Place &place) const override {
Y
Yang Yu 已提交
117
    auto *dout_var = scope.FindVar(Input(framework::GradVarName("Out")));
Y
Yang Yu 已提交
118
    auto *dx_var = scope.FindVar(Output(framework::GradVarName("X")));
Y
Yang Yu 已提交
119 120 121 122 123 124 125 126 127
    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 已提交
128
    // get device context from pool
Y
Yang Yu 已提交
129 130
    platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
    auto &dev_ctx = *pool.Get(place);
D
dzhwinter 已提交
131

Y
Yang Yu 已提交
132 133 134 135 136
    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];
137 138
      auto slice = dx_tensor.Slice(0, static_cast<int>(height));
      framework::TensorCopy(dout_tensor, dout_tensor.place(), dev_ctx, &slice);
Y
Refine  
Yang Yu 已提交
139
      if (dx_tensor.dims()[0] > height) {
Y
Yang Yu 已提交
140
        auto rest_tensor = dx_tensor.Slice(
Y
Refine  
Yang Yu 已提交
141
            static_cast<int>(height), static_cast<int>(dx_tensor.dims()[0]));
Y
Yang Yu 已提交
142 143 144
        math::set_constant(dev_ctx, &rest_tensor, 0.0f);
      }
    }
145
    dx_tensor.set_lod(x_tensor.lod());
Y
Yang Yu 已提交
146 147 148
  }
};

Y
Yang Yu 已提交
149
class ShrinkRNNMemoryGradInferShape : public framework::InferShapeBase {
Y
Yang Yu 已提交
150 151 152 153
 public:
  void operator()(framework::InferShapeContext *context) const override {
    PADDLE_ENFORCE(context->HasInput("X"));
    PADDLE_ENFORCE(context->HasOutput(framework::GradVarName("X")));
154 155 156

    context->ShareDim("X", /*->*/ framework::GradVarName("X"));
    context->ShareLoD("X", /*->*/ framework::GradVarName("X"));
Y
Yang Yu 已提交
157 158 159
  }
};

Y
Yang Yu 已提交
160
class ShrinkRNNGradOpMaker : public framework::SingleGradOpDescMaker {
Y
Yang Yu 已提交
161 162 163 164
 public:
  using framework::SingleGradOpDescMaker::SingleGradOpDescMaker;

 protected:
Y
Yu Yang 已提交
165 166
  std::unique_ptr<framework::OpDesc> Apply() const override {
    auto *op = new framework::OpDesc();
Y
Yang Yu 已提交
167
    op->SetType("shrink_rnn_memory_grad");
Y
Yang Yu 已提交
168 169 170 171
    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 已提交
172
    return std::unique_ptr<framework::OpDesc>(op);
Y
Yang Yu 已提交
173 174 175 176 177 178 179
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
Y
Yang Yu 已提交
180 181 182 183 184
REGISTER_OPERATOR(shrink_rnn_memory, ops::ShrinkRNNMemoryOp,
                  ops::ShrinkRNNMemoryInferShape,
                  ops::ShrinkRNNMemoryOpProtoMaker, ops::ShrinkRNNGradOpMaker);
REGISTER_OPERATOR(shrink_rnn_memory_grad, ops::ShrinkRNNMemoryGradOp,
                  ops::ShrinkRNNMemoryGradInferShape);