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

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

39 40 41
    auto* hidden_t0 = ctx.Input<Tensor>("H0");
    auto* cell_t0 = ctx.Input<Tensor>("C0");

D
dangqingqing 已提交
42
    auto* batch_gate = ctx.Output<LoDTensor>("BatchGate");
43
    batch_gate->mutable_data<T>(ctx.GetPlace());
D
dangqingqing 已提交
44
    auto* hidden_out = ctx.Output<LoDTensor>("Hidden");
45
    hidden_out->mutable_data<T>(ctx.GetPlace());
D
dangqingqing 已提交
46
    auto* cell_out = ctx.Output<LoDTensor>("Cell");
47 48
    cell_out->mutable_data<T>(ctx.GetPlace());

49
    bool is_reverse = ctx.Attr<bool>("is_reverse");
50
    math::LoDTensor2BatchFunctor<Place, T> to_batch;
D
dangqingqing 已提交
51 52
    auto& device_ctx = ctx.device_context();
    to_batch(device_ctx, *input, *batch_gate, true, is_reverse);
53 54

    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
      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 已提交
72
    if (bias && ctx.Attr<bool>("use_peepholes")) {
D
dangqingqing 已提交
73 74
      T* bias_data = const_cast<T*>(bias->data<T>());
      // the code style in LstmMetaValue will be updated later.
75

D
dangqingqing 已提交
76 77 78 79 80 81 82 83
      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;
    }
84
    lstm_value.prevStateValue = nullptr;
85 86 87
    Tensor ordered_c0;
    if (cell_t0) {
      math::CopyMatrixRowsFunctor<Place, T> row_shuffle;
D
dangqingqing 已提交
88
      ordered_c0.mutable_data<T>(cell_t0->dims(), ctx.GetPlace());
89 90 91 92
      const size_t* order = batch_gate->lod()[2].data();
      row_shuffle(device_ctx, *cell_t0, order, ordered_c0, true);
      lstm_value.prevStateValue = ordered_c0.data<T>();
    }
93

D
dangqingqing 已提交
94 95
    // Use the local variable as here.
    LoDTensor batch_hidden, batch_cell;
96
    auto* batch_cell_pre_act = ctx.Output<LoDTensor>("BatchCellPreAct");
D
dangqingqing 已提交
97
    batch_hidden.mutable_data<T>(dims, ctx.GetPlace());
98
    batch_cell.mutable_data<T>(dims, ctx.GetPlace());
99
    batch_cell_pre_act->mutable_data<T>(dims, ctx.GetPlace());
100

D
dangqingqing 已提交
101
    auto batch_starts = batch_gate->lod()[0];
Y
Yu Yang 已提交
102
    size_t num_batch = batch_starts.size() - 1;
103 104 105
    auto gate_act = ctx.Attr<std::string>("gate_activation");
    auto cell_act = ctx.Attr<std::string>("cell_activation");
    auto cand_act = ctx.Attr<std::string>("candidate_activation");
106

Y
Yu Yang 已提交
107 108 109
    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]);
110

D
dangqingqing 已提交
111
      Tensor gate_t = batch_gate->Slice(bstart, bend);
D
dangqingqing 已提交
112
      Tensor out_t = batch_hidden.Slice(bstart, bend);
D
dangqingqing 已提交
113
      Tensor cell_t = batch_cell.Slice(bstart, bend);
114
      Tensor cell_pre_act_t = batch_cell_pre_act->Slice(bstart, bend);
115 116 117

      int cur_batch_size = bend - bstart;

118
      if (n > 0) {
Y
Yu Yang 已提交
119
        int pre_h_start = static_cast<int>(batch_starts[n - 1]);
D
dangqingqing 已提交
120
        int pre_h_end = pre_h_start + cur_batch_size;
D
dangqingqing 已提交
121 122 123
        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 已提交
124
                               static_cast<T>(1.0));
125 126 127
      } else if (hidden_t0) {
        math::CopyMatrixRowsFunctor<Place, T> row_shuffle;
        Tensor ordered_h0;
D
dangqingqing 已提交
128
        ordered_h0.mutable_data<T>(hidden_t0->dims(), ctx.GetPlace());
129 130 131 132 133
        const size_t* order = batch_gate->lod()[2].data();
        row_shuffle(device_ctx, *hidden_t0, order, ordered_h0, true);
        math::matmul<Place, T>(device_ctx, ordered_h0, false, *weight, false,
                               static_cast<T>(1.0), &gate_t,
                               static_cast<T>(1.0));
134 135 136 137 138 139
      }

      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 已提交
140
      math::LstmUnitFunctor<Place, T>::compute(device_ctx, lstm_value,
141 142 143
                                               frame_size, cur_batch_size,
                                               gate_act, cell_act, cand_act);
      lstm_value.prevStateValue = lstm_value.stateValue;
D
dangqingqing 已提交
144
    }
145 146

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

151
    batch_cell.set_lod(batch_gate->lod());
152
    // restore the output cell state in LoDTensor from the batch cell
D
dangqingqing 已提交
153
    to_seq(device_ctx, batch_cell, *cell_out);
D
dangqingqing 已提交
154
  }
D
dangqingqing 已提交
155 156 157 158 159
};

template <typename Place, typename T>
class LSTMGradKernel : public framework::OpKernel<T> {
 public:
D
dangqingqing 已提交
160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176
  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* 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"));

177 178 179 180 181 182
    auto* h0 = ctx.Input<Tensor>("H0");
    auto* c0 = ctx.Input<Tensor>("C0");

    auto* h0_g = ctx.Output<Tensor>(framework::GradVarName("H0"));
    auto* c0_g = ctx.Output<Tensor>(framework::GradVarName("C0"));

D
dangqingqing 已提交
183
    auto& device_ctx = ctx.device_context();
184
    math::SetConstant<Place, T> zero;
D
dangqingqing 已提交
185
    if (weight_g) {
186
      weight_g->mutable_data<T>(ctx.GetPlace());
D
dangqingqing 已提交
187 188 189
      zero(device_ctx, weight_g, static_cast<T>(0.0));
    }

190 191 192 193 194 195 196 197
    Tensor ordered_h0, ordered_c0, ordered_h0_g, ordered_c0_g;
    math::CopyMatrixRowsFunctor<Place, T> row_shuffle;
    const size_t* order = batch_gate->lod()[2].data();
    if (c0) {
      ordered_c0.mutable_data<T>(c0->dims(), ctx.GetPlace());
      row_shuffle(device_ctx, *c0, order, ordered_c0, true);
    }

D
dangqingqing 已提交
198 199 200 201 202 203
    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;
D
dangqingqing 已提交
204
    if (bias && ctx.Attr<bool>("use_peepholes")) {
D
dangqingqing 已提交
205 206 207 208 209 210 211 212 213 214 215
      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;
D
dangqingqing 已提交
216

D
dangqingqing 已提交
217
    if (bias && bias_g) {
D
dangqingqing 已提交
218
      bias_g->mutable_data<T>(ctx.GetPlace());
219
      zero(device_ctx, bias_g, static_cast<T>(0.0));
D
dangqingqing 已提交
220 221 222
    }
    if (bias && bias_g && ctx.Attr<bool>("use_peepholes")) {
      T* bias_g_data = bias_g->data<T>();
D
dangqingqing 已提交
223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252
      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());
253
    // TODO(qingqing) support the case output cell has gradient.
D
dangqingqing 已提交
254
    // to_batch(device_ctx, *cell_g, batch_cell_g, false);
255
    zero(device_ctx, &batch_cell_g, static_cast<T>(0.0));
D
dangqingqing 已提交
256 257 258 259 260

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

261 262 263
    auto gate_act = ctx.Attr<std::string>("gate_activation");
    auto cell_act = ctx.Attr<std::string>("cell_activation");
    auto cand_act = ctx.Attr<std::string>("candidate_activation");
D
dangqingqing 已提交
264 265 266

    auto batch_starts = batch_gate->lod()[0];
    size_t num_batch = batch_starts.size() - 1;
267
    for (int n = static_cast<int>(num_batch) - 1; n >= 0; n--) {
D
dangqingqing 已提交
268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284
      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>();

285
      if (n > 0) {
D
dangqingqing 已提交
286 287 288 289 290 291
        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 {
292 293 294 295 296 297 298 299 300 301 302
        if (c0) {
          lstm_value.prevStateValue = ordered_c0.data<T>();
        } else {
          lstm_value.prevStateValue = nullptr;
        }
        if (c0 && c0_g) {
          ordered_c0_g.mutable_data<T>(c0_g->dims(), ctx.GetPlace());
          lstm_grad.prevStateGrad = ordered_c0_g.data<T>();
        } else {
          lstm_grad.prevStateGrad = nullptr;
        }
D
dangqingqing 已提交
303 304 305 306 307 308 309
      }

      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);

310
      if (n > 0) {
D
dangqingqing 已提交
311 312 313 314 315 316 317 318 319 320 321 322 323
        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));
        }
324 325 326 327 328 329 330 331 332 333 334 335 336 337
      } else {
        if (h0 && weight_g) {
          ordered_h0.mutable_data<T>(h0->dims(), ctx.GetPlace());
          row_shuffle(device_ctx, *h0, order, ordered_h0, true);
          math::matmul<Place, T>(device_ctx, ordered_h0, true, gate_g, false,
                                 static_cast<T>(1.0), weight_g,
                                 static_cast<T>(1.0));
        }
        if (h0 && h0_g) {
          ordered_h0_g.mutable_data<T>(h0_g->dims(), ctx.GetPlace());
          math::matmul<Place, T>(device_ctx, gate_g, false, *weight, true,
                                 static_cast<T>(1.0), &ordered_h0_g,
                                 static_cast<T>(0.0));
        }
D
dangqingqing 已提交
338 339 340 341 342 343
      }
    }

    math::Batch2LoDTensorFunctor<Place, T> to_seq;
    if (in_g) {
      /* backward data */
344
      in_g->mutable_data<T>(ctx.GetPlace());
D
dangqingqing 已提交
345 346 347 348
      to_seq(device_ctx, batch_gate_g, *in_g);
    }
    if (bias && bias_g) {
      /* backward bias */
349 350 351 352
      int m = static_cast<int>(batch_gate_g.dims()[0]);
      int n = static_cast<int>(batch_gate_g.dims()[1]);

      Tensor ones;
353
      ones.mutable_data<T>({m}, ctx.GetPlace());
354 355 356 357 358
      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 已提交
359
    }
360 361 362 363 364 365 366 367 368

    if (h0 && h0_g) {
      h0_g->mutable_data<T>(ctx.GetPlace());
      row_shuffle(device_ctx, ordered_h0_g, order, *h0_g, false);
    }
    if (c0 && c0_g) {
      c0_g->mutable_data<T>(ctx.GetPlace());
      row_shuffle(device_ctx, ordered_c0_g, order, *c0_g, false);
    }
D
dangqingqing 已提交
369
  }
D
dangqingqing 已提交
370 371 372 373
};

}  // namespace operators
}  // namespace paddle