lstm_op.h 5.5 KB
Newer Older
D
dangqingqing 已提交
1 2
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

D
dangqingqing 已提交
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
D
dangqingqing 已提交
6

D
dangqingqing 已提交
7
http://www.apache.org/licenses/LICENSE-2.0
D
dangqingqing 已提交
8

D
dangqingqing 已提交
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. */
D
dangqingqing 已提交
14 15 16

#pragma once
#include "paddle/framework/op_registry.h"
17 18 19
#include "paddle/operators/math/lstm_compute.h"
#include "paddle/operators/math/math_function.h"
#include "paddle/operators/math/sequence2batch.h"
D
dangqingqing 已提交
20 21 22 23 24 25

namespace paddle {
namespace operators {

using framework::LoDTensor;
using framework::Tensor;
26 27 28
template <typename T, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
using EigenMatrix = framework::EigenMatrix<T, MajorType, IndexType>;
D
dangqingqing 已提交
29 30 31 32

template <typename Place, typename T>
class LSTMKernel : public framework::OpKernel<T> {
 public:
D
dangqingqing 已提交
33
  void Compute(const framework::ExecutionContext& ctx) const override {
34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54
    auto* input = ctx.Input<framework::LoDTensor>("Input");
    auto* weight = ctx.Input<framework::Tensor>("Weight");
    auto* bias = ctx.Input<framework::Tensor>("Bias");

    auto* batch_gate = ctx.Output<framework::LoDTensor>("BatchGate");
    batch_gate->mutable_data<T>(ctx.GetPlace());
    auto* hidden_out = ctx.Output<framework::LoDTensor>("Hidden");
    hidden_out->mutable_data<T>(ctx.GetPlace());
    auto* cell_out = ctx.Output<framework::LoDTensor>("Cell");
    cell_out->mutable_data<T>(ctx.GetPlace());

    // Now the function ShareLoD in InferShape is not implemented.
    // So copy LoD here.
    ctx.ShareLoD("Input", "Hidden");
    ctx.ShareLoD("Input", "Cell");

    bool is_reverse = ctx.Attr<bool>("isReverse");
    math::LoDTensor2BatchFunctor<Place, T> to_batch;
    to_batch(ctx.device_context(), *input, *batch_gate, is_reverse);

    auto in_dims = input->dims();
Y
Yu Yang 已提交
55
    int frame_size = static_cast<int>(in_dims[1] / 4);
56
    framework::DDim dims({in_dims[0], frame_size});
D
dangqingqing 已提交
57

58 59 60
    if (bias) {
      Eigen::array<int, 2> extents({{1, 4 * frame_size}});
      Eigen::array<int, 2> offsets({{0, 0}});
D
dangqingqing 已提交
61
      auto b = EigenMatrix<T>::From(*bias);
62 63 64 65 66 67 68 69 70 71 72
      auto gate = EigenMatrix<T>::From(*batch_gate);
      gate.device(ctx.GetEigenDevice<Place>()) =
          gate +
          b.slice(offsets, extents)
              .reshape(Eigen::array<int, 2>({{1, frame_size * 4}}))
              .broadcast(
                  Eigen::array<int, 2>({{static_cast<int>(in_dims[0]), 1}}));
    }

    math::LstmMetaValue<T> lstm_value;
    T* bias_data = const_cast<T*>(bias->data<T>());
Y
Yu Yang 已提交
73
    // the code style in LstmMetaValue will be updated later.
74 75 76 77 78
    lstm_value.checkIg = bias_data + 4 * frame_size;
    lstm_value.checkFg = lstm_value.checkIg + frame_size;
    lstm_value.checkOg = lstm_value.checkFg + frame_size;
    lstm_value.prevStateValue = nullptr;

D
dangqingqing 已提交
79
    framework::LoDTensor batch_out, batch_cell, batch_cell_pre_act;
80 81 82
    batch_out.mutable_data<T>(dims, ctx.GetPlace());
    batch_cell.mutable_data<T>(dims, ctx.GetPlace());
    batch_cell_pre_act.mutable_data<T>(dims, ctx.GetPlace());
83

D
dangqingqing 已提交
84
    auto batch_starts = batch_gate->lod()[0];
Y
Yu Yang 已提交
85
    size_t num_batch = batch_starts.size() - 1;
86 87 88 89
    auto gate_act = ctx.Attr<std::string>("gateActivation");
    auto cell_act = ctx.Attr<std::string>("cellActivation");
    auto cand_act = ctx.Attr<std::string>("candidateActivation");

Y
Yu Yang 已提交
90 91 92
    for (size_t n = 0; n < num_batch; n++) {
      int bstart = static_cast<int>(batch_starts[n]);
      int bend = static_cast<int>(batch_starts[n + 1]);
93

D
dangqingqing 已提交
94 95 96 97
      Tensor gate_t = batch_gate->Slice(bstart, bend);
      Tensor out_t = batch_out.Slice(bstart, bend);
      Tensor cell_t = batch_cell.Slice(bstart, bend);
      Tensor cell_pre_act_t = batch_cell_pre_act.Slice(bstart, bend);
98 99 100 101

      int cur_batch_size = bend - bstart;

      if (n != 0) {
Y
Yu Yang 已提交
102
        int pre_h_start = static_cast<int>(batch_starts[n - 1]);
D
dangqingqing 已提交
103
        int pre_h_end = pre_h_start + cur_batch_size;
D
dangqingqing 已提交
104
        auto pre_hidden_t = batch_out.Slice(pre_h_start, pre_h_end);
105 106
        math::matmul<Place, T>(ctx.device_context(), pre_hidden_t, false,
                               *weight, false, static_cast<T>(1.0), &gate_t,
D
dangqingqing 已提交
107
                               static_cast<T>(1.0));
108
      }
Y
Yu Yang 已提交
109
      // else if : FIXME support the initial hidden and cell
110 111 112 113 114 115 116 117 118

      lstm_value.gateValue = gate_t.data<T>();
      lstm_value.outputValue = out_t.data<T>();
      lstm_value.stateValue = cell_t.data<T>();
      lstm_value.stateActiveValue = cell_pre_act_t.data<T>();
      math::LstmUnitFunctor<Place, T>::compute(ctx.device_context(), lstm_value,
                                               frame_size, cur_batch_size,
                                               gate_act, cell_act, cand_act);
      lstm_value.prevStateValue = lstm_value.stateValue;
D
dangqingqing 已提交
119
    }
120 121 122 123 124 125

    math::Batch2LoDTensorFunctor<Place, T> to_seq;
    batch_out.set_lod(batch_gate->lod());
    // restore the output hidden in LoDTensor from the batch hidden
    to_seq(ctx.device_context(), batch_out, *hidden_out);

126
    batch_cell.set_lod(batch_gate->lod());
127 128
    // restore the output cell state in LoDTensor from the batch cell
    to_seq(ctx.device_context(), batch_cell, *cell_out);
D
dangqingqing 已提交
129
  }
D
dangqingqing 已提交
130 131 132 133 134 135 136 137 138 139
};

template <typename Place, typename T>
class LSTMGradKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {}
};

}  // namespace operators
}  // namespace paddle