shrink_rnn_memory_op.cc 7.4 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 57 58
    // 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 已提交
59
      height = lod_offset.second.second;
Y
yangyaming 已提交
60
      auto out_lod = out_tensor.mutable_lod();
61
      framework::AppendLoD(out_lod, lod_offset.first);
Y
yangyaming 已提交
62 63
    }

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

Y
Yang Yu 已提交
70
class ShrinkRNNMemoryOpProtoMaker : public framework::OpProtoAndCheckerMaker {
Y
Yang Yu 已提交
71
 public:
72
  ShrinkRNNMemoryOpProtoMaker(OpProto *proto, OpAttrChecker *op_checker)
Y
Yang Yu 已提交
73
      : OpProtoAndCheckerMaker(proto, op_checker) {
74 75 76 77 78 79
    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.");
80 81 82 83 84 85 86 87 88 89 90
    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 已提交
91 92 93
  }
};

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

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

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

Y
Yang Yu 已提交
147
class ShrinkRNNMemoryGradInferShape : public framework::InferShapeBase {
Y
Yang Yu 已提交
148 149 150 151 152 153
 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"));
154
    context->ShareLoD("X", framework::GradVarName("X"));
Y
Yang Yu 已提交
155 156 157
  }
};

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

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

}  // namespace operators
}  // namespace paddle

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