linear_chain_crf_op.h 14.3 KB
Newer Older
C
caoying03 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

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

    http://www.apache.org/licenses/LICENSE-2.0

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. */

#pragma once
#include "paddle/framework/eigen.h"
#include "paddle/framework/op_registry.h"
18
#include "paddle/operators/math/math_function.h"
C
caoying03 已提交
19 20 21 22

namespace paddle {
namespace operators {

C
caoying03 已提交
23
template <typename T>
C
caoying03 已提交
24
static inline T NormalizeL1(T* x, size_t len) {
C
caoying03 已提交
25 26 27 28 29 30 31 32 33 34 35 36 37
  T sum = 0.;
  for (size_t i = 0; i < len; ++i) sum += x[i];
  // (This comment is from the old LinearChainCRFLayer.)
  // Right now, we just bet that sum won't be zero. If this really happens, we
  // will figure out what should be done then.
  PADDLE_ENFORCE(sum,
                 "The unnormalized probabilities of all possible unfinished "
                 "sequences must be greater than 0.");
  T s = 1. / sum;
  for (size_t i = 0; i < len; ++i) x[i] *= s;
  return sum;
}

38 39 40 41 42 43 44 45
template <typename T>
struct ScalarMul {
  explicit ScalarMul(const T& scalar) : scalar(scalar) {}
  T operator()(const T& val) const { return val * scalar; }

  T scalar;
};

C
caoying03 已提交
46 47
using framework::LoDTensor;
using framework::LoD;
48
using framework::Tensor;
C
caoying03 已提交
49 50 51 52
template <typename T, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
using EigenMatrix = framework::EigenMatrix<T, MajorType, IndexType>;

Q
QI JUN 已提交
53
template <typename DeviceContext, typename T>
C
caoying03 已提交
54
class LinearChainCRFOpKernel : public framework::OpKernel<T> {
C
caoying03 已提交
55
 public:
C
caoying03 已提交
56 57 58
  void Compute(const framework::ExecutionContext& ctx) const override {
    // TODO(caoying) The checks related to LoD information should be
    // moved into InferShape once after the InferShape is refactored.
59
    PADDLE_ENFORCE_EQ(ctx.Input<LoDTensor>("Emission")->NumLevels(), 1UL,
C
caoying03 已提交
60
                      "The Input(Emission) should be a sequence.");
61
    PADDLE_ENFORCE_EQ(ctx.Input<LoDTensor>("Label")->NumLevels(), 1UL,
C
caoying03 已提交
62
                      "The Input(Label) should be a sequence.");
63 64
    auto in_lod = ctx.Input<LoDTensor>("Label")->lod();
    PADDLE_ENFORCE(in_lod.size(), "Input(Label) must be a sequence.");
C
caoying03 已提交
65
    const size_t level = 0;
66 67
    const size_t seq_num = in_lod[level].size() - 1;

68 69 70 71 72 73 74 75
    const LoDTensor* emission_weights = ctx.Input<LoDTensor>("Emission");
    const Tensor* transition_weights = ctx.Input<Tensor>("Transition");
    const LoDTensor* label = ctx.Input<LoDTensor>("Label");

    Tensor* emission_exps = ctx.Output<Tensor>("EmissionExps");
    Tensor* transition_exps = ctx.Output<Tensor>("TransitionExps");
    Tensor* alpha = ctx.Output<Tensor>("Alpha");
    Tensor* ll = ctx.Output<Tensor>("LogLikelihood");
C
caoying03 已提交
76

77 78 79 80 81 82 83 84 85 86 87
    // Because the computation codes only runs on CPU, here the memory for all
    // the outputs is FIXED to be allocated on the CPU memory.
    emission_exps->mutable_data<T>(platform::CPUPlace());
    transition_exps->mutable_data<T>(platform::CPUPlace());
    alpha->mutable_data<T>(platform::CPUPlace());

    // Resize the output tensor to its correct dimension.
    ll->Resize({static_cast<int>(seq_num), 1});
    ll->mutable_data<T>(platform::CPUPlace());

    // Now, all the inputs and outputs should be on the CPU memory.
C
caoying03 已提交
88 89 90 91 92 93
    auto emission_dims = emission_weights->dims();
    const size_t batch_size = emission_dims[0];
    const size_t tag_num = emission_dims[1];

    Tensor emission_row_max;
    emission_row_max.mutable_data<T>(
C
Cao Ying 已提交
94
        framework::make_ddim({static_cast<int64_t>(batch_size), 1}),
95
        platform::CPUPlace());
C
caoying03 已提交
96

Q
QI JUN 已提交
97 98
    auto& place = *ctx.template device_context<platform::CPUDeviceContext>()
                       .eigen_device();
C
caoying03 已提交
99 100 101 102 103 104 105 106 107 108 109 110 111 112
    auto x = EigenMatrix<T>::From(*emission_weights);
    auto x_row_max = EigenMatrix<T>::From(emission_row_max);
    x_row_max.device(place) =
        x.maximum(Eigen::DSizes<int, 1>(1))
            .reshape(Eigen::DSizes<int, 2>(int(batch_size), 1));

    auto x_exps = EigenMatrix<T>::From(*emission_exps);
    x_exps.device(place) =
        (x - x_row_max.broadcast(Eigen::DSizes<int, 2>(1, tag_num))).exp();

    auto w = EigenMatrix<T>::From(*transition_weights);
    auto w_exps = EigenMatrix<T>::From(*transition_exps);
    w_exps.device(place) = w.exp();

113
    T* log_likelihood = ll->data<T>();
C
caoying03 已提交
114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132
    for (size_t i = 0; i < seq_num; ++i) {
      int start_pos = static_cast<int>(in_lod[level][i]);
      int end_pos = static_cast<int>(in_lod[level][i + 1]);
      if (end_pos == start_pos) {
        // If an empty input sequence is given, pad 0 for its cost.
        log_likelihood[i] = 0.;
        continue;
      }

      const Tensor one_seq = emission_weights->Slice(start_pos, end_pos);
      Tensor one_seq_row_max = emission_row_max.Slice(start_pos, end_pos);
      Tensor one_seq_exps = emission_exps->Slice(start_pos, end_pos);
      const Tensor one_seq_label = label->Slice(start_pos, end_pos);
      Tensor one_seq_alpha = alpha->Slice(start_pos, end_pos);

      log_likelihood[i] = ForwardOneSequence(
          one_seq, one_seq_row_max, one_seq_exps, *transition_weights,
          *transition_exps, one_seq_label, &one_seq_alpha);
    }
133 134 135
  };

 private:
C
caoying03 已提交
136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163
  T ForwardOneSequence(const Tensor& emission, const Tensor& emission_row_max,
                       const Tensor& emission_exps, const Tensor& trans_weights,
                       const Tensor& trans_weight_exps, const Tensor& label,
                       Tensor* alpha) const {
    const T* x = emission.data<T>();
    const T* x_row_max = emission_row_max.data<T>();
    const T* x_exps = emission_exps.data<T>();
    const T* w = trans_weights.data<T>();
    const T* w_exps = trans_weight_exps.data<T>();
    T* alpha_value = alpha->data<T>();

    auto x_dims = emission.dims();
    const size_t seq_length = x_dims[0];
    const size_t tag_num = x_dims[1];
    // The 1st row of w are transition weights for start mask.
    // The 2nd row of w are transition weights for end mask.
    // Transition weights between other tags begin from the 3rd row of w.
    const size_t state_trans_base_idx = 2;

    for (size_t i = 0; i < tag_num; ++i) {
      alpha_value[i] = w_exps[i] * x_exps[i];
    }
    T ll = -x_row_max[0] - std::log(NormalizeL1<T>(alpha_value, tag_num));

    for (size_t k = 1; k < seq_length; ++k) {
      for (size_t i = 0; i < tag_num; ++i) {
        T sum = 0.;
        for (size_t j = 0; j < tag_num; ++j) {
C
caoying03 已提交
164
          sum += alpha_value[(k - 1) * tag_num + j] *  // (*)
C
caoying03 已提交
165 166 167 168 169 170 171 172 173 174 175 176 177 178 179
                 w_exps[(j + state_trans_base_idx) * tag_num + i];
        }
        alpha_value[k * tag_num + i] = x_exps[k * tag_num + i] * sum;
      }
      // NormalizeL1 is to avoid underflow or overflow at (*).
      ll -= x_row_max[k] +
            std::log(NormalizeL1<T>(alpha_value + k * tag_num, tag_num));
    }
    T sum = 0.;
    for (size_t i = 0; i < tag_num; ++i) {
      sum += alpha_value[(seq_length - 1) * tag_num + i] * w_exps[tag_num + i];
    }
    ll -= std::log(sum);
    // Now ll is equal to -log(Z).

Q
Qiao Longfei 已提交
180
    const int64_t* lbl = label.data<int64_t>();
C
caoying03 已提交
181
    PADDLE_ENFORCE_LT(
C
Cao Ying 已提交
182
        static_cast<size_t>(*std::max_element(lbl, lbl + seq_length)), tag_num,
C
caoying03 已提交
183 184 185 186 187 188 189 190 191 192
        "An invalid tag label that execesses the largest tag number.");

    // Calculate the nominator part, which depends on the label sequence.
    ll += w[lbl[0]] /*start transition*/ + x[lbl[0]] +
          w[tag_num + lbl[seq_length - 1]] /*end transition*/;
    for (size_t k = 1; k < seq_length; ++k) {
      ll += x[k * tag_num + lbl[k]] +
            w[(lbl[k - 1] + state_trans_base_idx) * tag_num + lbl[k]];
    }
    return -ll;
C
caoying03 已提交
193
  }
C
caoying03 已提交
194 195
};

Q
QI JUN 已提交
196
template <typename DeviceContext, typename T>
C
caoying03 已提交
197
class LinearChainCRFGradOpKernel : public framework::OpKernel<T> {
C
caoying03 已提交
198
 public:
C
caoying03 已提交
199
  void Compute(const framework::ExecutionContext& ctx) const override {
200 201 202 203
    const size_t level = 0;  // currently, only support sequence.
    auto lod = ctx.Input<LoDTensor>("Label")->lod();
    PADDLE_ENFORCE(lod.size(), "Input(Label) must be a sequence.");

204 205 206 207 208 209
    const Tensor* label = ctx.Input<LoDTensor>("Label");
    const Tensor* emission_exps = ctx.Input<Tensor>("EmissionExps");
    const Tensor* transition_exps = ctx.Input<Tensor>("TransitionExps");
    const Tensor* alpha = ctx.Input<Tensor>("Alpha");
    const T* ll_grad =
        ctx.Input<Tensor>(framework::GradVarName("LogLikelihood"))->data<T>();
210

211 212 213 214
    Tensor* emission_grad =
        ctx.Output<Tensor>(framework::GradVarName("Emission"));
    Tensor* transition_grad =
        ctx.Output<Tensor>(framework::GradVarName("Transition"));
C
caoying03 已提交
215 216 217

    // TODO(caoying) Fix this constraint. When the Input(Emission) is from the
    // data reader operator, it can have no gradients.
218 219 220 221
    PADDLE_ENFORCE(emission_grad, "Output(Emission@Grad) should not be null.");
    emission_grad->mutable_data<T>(platform::CPUPlace());
    if (transition_grad) {
      transition_grad->mutable_data<T>(platform::CPUPlace());
Q
QI JUN 已提交
222
      math::set_constant(ctx.device_context(), transition_grad, 0.);
C
caoying03 已提交
223
    }
224
    // Now, all the inputs and outputs should be on the CPU memory.
C
caoying03 已提交
225 226 227 228

    auto emission_dims = emission_exps->dims();
    // Beta is the memo table used in dynamic programming to calculate the
    // backwark vectors. For a backward vector i (the i-th row of beta), it
229 230
    // captures the unnormalized probabilities of partial sequences starting
    // at position i.
C
caoying03 已提交
231
    Tensor beta;
232
    beta.mutable_data<T>(emission_dims, platform::CPUPlace());
C
caoying03 已提交
233 234 235 236 237 238 239 240 241 242 243 244 245

    for (size_t i = 0; i < lod[level].size() - 1; ++i) {
      int start_pos = static_cast<int>(lod[level][i]);
      int end_pos = static_cast<int>(lod[level][i + 1]);
      if (end_pos == start_pos) continue;

      const Tensor one_seq_emission_exps =
          emission_exps->Slice(start_pos, end_pos);
      const Tensor one_seq_label = label->Slice(start_pos, end_pos);
      const Tensor one_seq_alpha = alpha->Slice(start_pos, end_pos);
      Tensor one_seq_beta = beta.Slice(start_pos, end_pos);
      Tensor one_seq_emission_grad = emission_grad->Slice(start_pos, end_pos);

Q
QI JUN 已提交
246 247 248 249
      BackwardOneSequence(
          ctx.template device_context<platform::CPUDeviceContext>(), ll_grad[i],
          one_seq_emission_exps, *transition_exps, one_seq_alpha, one_seq_label,
          &one_seq_beta, transition_grad, &one_seq_emission_grad);
250
    }
C
caoying03 已提交
251
  };
C
caoying03 已提交
252

253
 private:
Q
QI JUN 已提交
254 255
  void BackwardOneSequence(const platform::CPUDeviceContext& ctx,
                           const T ll_grad, const Tensor& emission_exps,
C
caoying03 已提交
256 257
                           const Tensor& transition_exps, const Tensor& alpha,
                           const Tensor& label, Tensor* beta,
C
caoying03 已提交
258
                           Tensor* transition_grad,
C
caoying03 已提交
259 260 261
                           Tensor* emission_grad) const {
    const T* w_exps = transition_exps.data<T>();
    const T* x_exps = emission_exps.data<T>();
Q
Qiao Longfei 已提交
262
    const int64_t* label_value = label.data<int64_t>();
C
caoying03 已提交
263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279
    T* beta_value = beta->data<T>();

    auto x_dims = emission_exps.dims();
    const size_t seq_length = x_dims[0];
    const size_t tag_num = x_dims[1];
    const size_t state_trans_base_idx = 2;

    // Calculate the backward vectors: beta.
    // First, calculate the initialition state.
    for (size_t i = 0; i < tag_num; ++i) {
      beta_value[(seq_length - 1) * tag_num + i] = w_exps[tag_num + i];
    }
    NormalizeL1<T>(beta_value + (seq_length - 1) * tag_num, tag_num);
    for (int k = static_cast<int>(seq_length) - 2; k >= 0; --k) {
      for (size_t i = 0; i < tag_num; ++i) {
        T sum = 0.;
        for (size_t j = 0; j < tag_num; ++j) {
C
caoying03 已提交
280
          sum += w_exps[(i + state_trans_base_idx) * tag_num + j] *  // (**)
C
caoying03 已提交
281 282 283 284 285 286 287 288 289
                 x_exps[(k + 1) * tag_num + j] *
                 beta_value[(k + 1) * tag_num + j];
        }
        beta_value[k * tag_num + i] = sum;
      }
      // NormalizeL1 is to avoid underflow or overflow at (**).
      NormalizeL1<T>(beta_value + k * tag_num, tag_num);
    }

290
    auto x_grad_mat = EigenMatrix<T>::From(*emission_grad);
C
caoying03 已提交
291 292
    auto alpha_mat = EigenMatrix<T>::From(alpha);
    auto beta_mat = EigenMatrix<T>::From(*beta);
293

Q
QI JUN 已提交
294
    auto* place = ctx.eigen_device();
C
caoying03 已提交
295 296 297 298
    auto prob = alpha_mat * beta_mat;
    auto row_sum = prob.sum(Eigen::DSizes<int, 1>(1))
                       .reshape(Eigen::DSizes<int, 2>(seq_length, 1))
                       .broadcast(Eigen::DSizes<int, 2>(1, tag_num));
299 300
    x_grad_mat.device(*place) =
        (prob / row_sum).unaryExpr(ScalarMul<T>(ll_grad));
C
caoying03 已提交
301 302

    for (size_t k = 0; k < seq_length; ++k) {
303
      x_grad_mat(k, label_value[k]) -= static_cast<T>(ll_grad);
C
caoying03 已提交
304 305 306 307 308
    }

    if (transition_grad) {
      T* trans_grad = transition_grad->data<T>();
      for (size_t k = 0; k < tag_num; ++k) {
309 310
        // Do not multiply by the output gradient here, because x_grad_mat has
        // alrealy done this.
C
caoying03 已提交
311 312 313 314 315 316 317
        trans_grad[k] += x_grad_mat(/*from start state*/ 0, k);
        trans_grad[tag_num + k] +=
            x_grad_mat(/*to end state*/ seq_length - 1, k);
      }

      auto x_exps_mat = EigenMatrix<T>::From(emission_exps);

318 319
      // TODO(caoying): Fix this to avoid using this local variable if we can
      // profile the training process.
C
caoying03 已提交
320
      Tensor tmp;
321
      tmp.mutable_data<T>(beta->dims(), platform::CPUPlace());
C
caoying03 已提交
322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341
      auto tmp_mat = EigenMatrix<T>::From(tmp);
      auto prob = beta_mat * x_exps_mat;
      auto row_sum = prob.sum(Eigen::DSizes<int, 1>(1))
                         .reshape(Eigen::DSizes<int, 2>(seq_length, 1))
                         .broadcast(Eigen::DSizes<int, 2>(1, tag_num));
      tmp_mat.device(*place) = prob / row_sum;

      for (size_t k = 1; k < seq_length; ++k) {
        T sum = 0.;
        for (size_t i = 0; i < tag_num; ++i) {
          for (size_t j = 0; j < tag_num; ++j) {
            sum += w_exps[(i + state_trans_base_idx) * tag_num + j] *  // (**)
                   alpha_mat(k - 1, i) * tmp_mat(k, j);
          }
        }
        sum = 1. / sum;
        for (size_t i = 0; i < tag_num; ++i) {
          for (size_t j = 0; j < tag_num; ++j) {
            trans_grad[(i + state_trans_base_idx) * tag_num + j] +=
                sum * w_exps[(i + state_trans_base_idx) * tag_num + j] *
342
                alpha_mat(k - 1, i) * tmp_mat(k, j) * ll_grad;
C
caoying03 已提交
343 344 345
          }
        }
        trans_grad[(label_value[k - 1] + state_trans_base_idx) * tag_num +
346
                   label_value[k]] -= static_cast<T>(ll_grad);
C
caoying03 已提交
347 348
      }
    }
C
caoying03 已提交
349
  }
C
caoying03 已提交
350 351 352 353
};

}  // namespace operators
}  // namespace paddle