lstm_op.h 12.3 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

namespace paddle {
namespace operators {

D
dangqingqing 已提交
24 25 26
using LoDTensor = framework::LoDTensor;
using Tensor = framework::Tensor;

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

template <typename Place, typename T>
class LSTMKernel : public framework::OpKernel<T> {
 public:
D
dangqingqing 已提交
34
  void Compute(const framework::ExecutionContext& ctx) const override {
D
dangqingqing 已提交
35 36 37
    auto* input = ctx.Input<LoDTensor>("Input");
    auto* weight = ctx.Input<Tensor>("Weight");
    auto* bias = ctx.Input<Tensor>("Bias");
38

D
dangqingqing 已提交
39
    auto* batch_gate = ctx.Output<LoDTensor>("BatchGate");
40
    batch_gate->mutable_data<T>(ctx.GetPlace());
D
dangqingqing 已提交
41
    auto* hidden_out = ctx.Output<LoDTensor>("Hidden");
42
    hidden_out->mutable_data<T>(ctx.GetPlace());
D
dangqingqing 已提交
43
    auto* cell_out = ctx.Output<LoDTensor>("Cell");
44 45 46 47 48 49 50 51 52
    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;
D
dangqingqing 已提交
53 54
    auto& device_ctx = ctx.device_context();
    to_batch(device_ctx, *input, *batch_gate, true, is_reverse);
55 56

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

60 61 62
    if (bias) {
      Eigen::array<int, 2> extents({{1, 4 * frame_size}});
      Eigen::array<int, 2> offsets({{0, 0}});
D
dangqingqing 已提交
63
      auto b = EigenMatrix<T>::From(*bias);
64 65 66 67 68 69 70 71 72 73
      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;
D
dangqingqing 已提交
74 75 76
    if (bias) {
      T* bias_data = const_cast<T*>(bias->data<T>());
      // the code style in LstmMetaValue will be updated later.
77

D
dangqingqing 已提交
78 79 80 81 82 83 84 85
      lstm_value.checkIg = bias_data + 4 * frame_size;
      lstm_value.checkFg = lstm_value.checkIg + frame_size;
      lstm_value.checkOg = lstm_value.checkFg + frame_size;
    } else {
      lstm_value.checkIg = nullptr;
      lstm_value.checkFg = nullptr;
      lstm_value.checkOg = nullptr;
    }
86 87
    lstm_value.prevStateValue = nullptr;

D
dangqingqing 已提交
88 89
    // Use the local variable as here.
    LoDTensor batch_hidden, batch_cell;
90
    auto* batch_cell_pre_act = ctx.Output<LoDTensor>("BatchCellPreAct");
D
dangqingqing 已提交
91
    batch_hidden.mutable_data<T>(dims, ctx.GetPlace());
92
    batch_cell.mutable_data<T>(dims, ctx.GetPlace());
93
    batch_cell_pre_act->mutable_data<T>(dims, ctx.GetPlace());
94

D
dangqingqing 已提交
95
    auto batch_starts = batch_gate->lod()[0];
Y
Yu Yang 已提交
96
    size_t num_batch = batch_starts.size() - 1;
97 98 99 100
    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 已提交
101 102 103
    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]);
104

D
dangqingqing 已提交
105
      Tensor gate_t = batch_gate->Slice(bstart, bend);
D
dangqingqing 已提交
106
      Tensor out_t = batch_hidden.Slice(bstart, bend);
D
dangqingqing 已提交
107
      Tensor cell_t = batch_cell.Slice(bstart, bend);
108
      Tensor cell_pre_act_t = batch_cell_pre_act->Slice(bstart, bend);
109 110 111 112

      int cur_batch_size = bend - bstart;

      if (n != 0) {
Y
Yu Yang 已提交
113
        int pre_h_start = static_cast<int>(batch_starts[n - 1]);
D
dangqingqing 已提交
114
        int pre_h_end = pre_h_start + cur_batch_size;
D
dangqingqing 已提交
115 116 117
        auto pre_hidden_t = batch_hidden.Slice(pre_h_start, pre_h_end);
        math::matmul<Place, T>(device_ctx, pre_hidden_t, false, *weight, false,
                               static_cast<T>(1.0), &gate_t,
D
dangqingqing 已提交
118
                               static_cast<T>(1.0));
119
      }
Y
Yu Yang 已提交
120
      // else if : FIXME support the initial hidden and cell
121 122 123 124 125

      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>();
D
dangqingqing 已提交
126
      math::LstmUnitFunctor<Place, T>::compute(device_ctx, lstm_value,
127 128 129
                                               frame_size, cur_batch_size,
                                               gate_act, cell_act, cand_act);
      lstm_value.prevStateValue = lstm_value.stateValue;
D
dangqingqing 已提交
130
    }
131 132

    math::Batch2LoDTensorFunctor<Place, T> to_seq;
D
dangqingqing 已提交
133
    batch_hidden.set_lod(batch_gate->lod());
134
    // restore the output hidden in LoDTensor from the batch hidden
D
dangqingqing 已提交
135
    to_seq(device_ctx, batch_hidden, *hidden_out);
136

137
    batch_cell.set_lod(batch_gate->lod());
138
    // restore the output cell state in LoDTensor from the batch cell
D
dangqingqing 已提交
139
    to_seq(device_ctx, batch_cell, *cell_out);
D
dangqingqing 已提交
140
  }
D
dangqingqing 已提交
141 142 143 144 145
};

template <typename Place, typename T>
class LSTMGradKernel : public framework::OpKernel<T> {
 public:
D
dangqingqing 已提交
146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165
  void Compute(const framework::ExecutionContext& ctx) const override {
    auto* input = ctx.Input<LoDTensor>("Input");
    auto* weight = ctx.Input<Tensor>("Weight");
    auto* bias = ctx.Input<Tensor>("Bias");

    auto* hidden_out = ctx.Input<LoDTensor>("Hidden");
    auto* cell_out = ctx.Input<LoDTensor>("Cell");

    auto* batch_gate = ctx.Input<LoDTensor>("BatchGate");
    auto* batch_cell_pre_act = ctx.Input<LoDTensor>("BatchCellPreAct");

    auto* hidden_g = ctx.Input<LoDTensor>(framework::GradVarName("Hidden"));
    auto* cell_g = ctx.Input<LoDTensor>(framework::GradVarName("Cell"));

    auto* in_g = ctx.Output<LoDTensor>(framework::GradVarName("Input"));
    auto* weight_g = ctx.Output<Tensor>(framework::GradVarName("Weight"));
    auto* bias_g = ctx.Output<Tensor>(framework::GradVarName("Bias"));

    auto& device_ctx = ctx.device_context();
    if (weight_g) {
166
      weight_g->mutable_data<T>(ctx.GetPlace());
D
dangqingqing 已提交
167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232
      math::SetConstant<Place, T> zero;
      zero(device_ctx, weight_g, static_cast<T>(0.0));
    }

    auto in_dims = input->dims();
    auto out_dims = hidden_g->dims();
    int frame_size = static_cast<int>(in_dims[1] / 4);
    PADDLE_ENFORCE_EQ(frame_size, out_dims[1]);

    math::LstmMetaValue<T> lstm_value;
    if (bias) {
      T* bias_data = const_cast<T*>(bias->data<T>());
      lstm_value.checkIg = bias_data + 4 * frame_size;
      lstm_value.checkFg = lstm_value.checkIg + frame_size;
      lstm_value.checkOg = lstm_value.checkFg + frame_size;
    } else {
      lstm_value.checkIg = nullptr;
      lstm_value.checkFg = nullptr;
      lstm_value.checkOg = nullptr;
    }

    math::LstmMetaGrad<T> lstm_grad;
    if (bias && bias_g) {
      T* bias_g_data = const_cast<T*>(bias_g->mutable_data<T>(ctx.GetPlace()));
      lstm_grad.checkIgGrad = bias_g_data + 4 * frame_size;
      lstm_grad.checkFgGrad = lstm_grad.checkIgGrad + frame_size;
      lstm_grad.checkOgGrad = lstm_grad.checkFgGrad + frame_size;
    } else {
      lstm_grad.checkIgGrad = nullptr;
      lstm_grad.checkFgGrad = nullptr;
      lstm_grad.checkOgGrad = nullptr;
    }

    math::LoDTensor2BatchFunctor<Place, T> to_batch;

    // use the local variable as here.
    LoDTensor batch_hidden;
    batch_hidden.mutable_data<T>(out_dims, ctx.GetPlace());
    batch_hidden.set_lod(batch_gate->lod());
    to_batch(device_ctx, *hidden_out, batch_hidden, false);

    LoDTensor batch_hidden_g;
    batch_hidden_g.mutable_data<T>(out_dims, ctx.GetPlace());
    batch_hidden_g.set_lod(batch_gate->lod());
    to_batch(device_ctx, *hidden_g, batch_hidden_g, false);

    LoDTensor batch_cell;
    batch_cell.mutable_data<T>(out_dims, ctx.GetPlace());
    batch_cell.set_lod(batch_gate->lod());
    to_batch(device_ctx, *cell_out, batch_cell, false);

    LoDTensor batch_cell_g;
    batch_cell_g.mutable_data<T>(out_dims, ctx.GetPlace());
    batch_cell_g.set_lod(batch_gate->lod());
    to_batch(device_ctx, *cell_g, batch_cell_g, false);

    LoDTensor batch_gate_g;
    batch_gate_g.mutable_data<T>(batch_gate->dims(), ctx.GetPlace());
    batch_gate_g.set_lod(batch_gate->lod());

    auto gate_act = ctx.Attr<std::string>("gateActivation");
    auto cell_act = ctx.Attr<std::string>("cellActivation");
    auto cand_act = ctx.Attr<std::string>("candidateActivation");

    auto batch_starts = batch_gate->lod()[0];
    size_t num_batch = batch_starts.size() - 1;
233
    for (int n = static_cast<int>(num_batch) - 1; n >= 0; n--) {
D
dangqingqing 已提交
234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286
      int bstart = static_cast<int>(batch_starts[n]);
      int bend = static_cast<int>(batch_starts[n + 1]);

      Tensor gate = batch_gate->Slice(bstart, bend);
      Tensor cell = batch_cell.Slice(bstart, bend);
      Tensor cell_pre_act = batch_cell_pre_act->Slice(bstart, bend);
      lstm_value.gateValue = gate.data<T>();
      lstm_value.stateValue = cell.data<T>();
      lstm_value.stateActiveValue = cell_pre_act.data<T>();

      Tensor out_g = batch_hidden_g.Slice(bstart, bend);
      Tensor gate_g = batch_gate_g.Slice(bstart, bend);
      Tensor cell_g = batch_cell_g.Slice(bstart, bend);
      lstm_grad.stateGrad = cell_g.data<T>();
      lstm_grad.gateGrad = gate_g.data<T>();
      lstm_grad.outputGrad = out_g.data<T>();

      if (n != 0) {
        int bstart_pre = static_cast<int>(batch_starts[n - 1]);
        Tensor cell_pre = batch_cell.Slice(bstart_pre, bstart);
        Tensor cell_pre_g = batch_cell_g.Slice(bstart_pre, bstart);
        lstm_value.prevStateValue = cell_pre.data<T>();
        lstm_grad.prevStateGrad = cell_pre_g.data<T>();
      } else {
        lstm_value.prevStateValue = nullptr;
        lstm_grad.prevStateGrad = nullptr;
      }

      int cur_batch_size = bend - bstart;
      math::LstmUnitGradFunctor<Place, T>::compute(
          device_ctx, lstm_value, lstm_grad, frame_size, cur_batch_size,
          gate_act, cell_act, cand_act);

      if (n != 0) {
        int pre_h_start = static_cast<int>(batch_starts[n - 1]);
        int pre_h_end = pre_h_start + cur_batch_size;
        auto pre_hidden_g = batch_hidden_g.Slice(pre_h_start, pre_h_end);
        math::matmul<Place, T>(device_ctx, gate_g, false, *weight, true,
                               static_cast<T>(1.0), &pre_hidden_g,
                               static_cast<T>(1.0));
        if (weight_g) {
          /* backward weight */
          auto pre_hidden = batch_hidden.Slice(pre_h_start, pre_h_end);
          math::matmul<Place, T>(device_ctx, pre_hidden, true, gate_g, false,
                                 static_cast<T>(1.0), weight_g,
                                 static_cast<T>(1.0));
        }
      }
    }

    math::Batch2LoDTensorFunctor<Place, T> to_seq;
    if (in_g) {
      /* backward data */
287
      in_g->mutable_data<T>(ctx.GetPlace());
D
dangqingqing 已提交
288 289 290 291
      to_seq(device_ctx, batch_gate_g, *in_g);
    }
    if (bias && bias_g) {
      /* backward bias */
292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312
      // Following Eigen computation failed for double type on GPU device.
      // bias_g->mutable_data<T>(ctx.GetPlace());
      // Tensor bias_mat;
      // bias_mat.ShareDataWith(*bias_g);
      // bias_mat.Resize({1, 4 * frame_size});

      // auto bias_g_e = EigenVector<T>::Flatten(bias_mat);
      // auto gate_g_e = EigenMatrix<T>::From(batch_gate_g);
      // Eigen::array<int, 1> dims{{0}};
      // bias_g_e.device(ctx.GetEigenDevice<Place>()) = gate_g_e.sum(dims);

      int m = static_cast<int>(batch_gate_g.dims()[0]);
      int n = static_cast<int>(batch_gate_g.dims()[1]);

      Tensor ones;
      ones.mutable_data<T>({1, m}, ctx.GetPlace());
      math::SetConstant<Place, T> set;
      set(device_ctx, &ones, static_cast<T>(1.0));

      math::gemv<Place, T>(device_ctx, true, m, n, 1., batch_gate_g.data<T>(),
                           ones.data<T>(), 0., bias_g->data<T>());
D
dangqingqing 已提交
313 314
    }
  }
D
dangqingqing 已提交
315 316 317 318
};

}  // namespace operators
}  // namespace paddle