linear_chain_crf_op.cc 23.9 KB
Newer Older
C
caoying03 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
/* 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. */

#include "paddle/operators/linear_chain_crf_op.h"

namespace paddle {
namespace operators {

C
caoying03 已提交
20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35
namespace {
template <typename T>
T NormalizeL1(T* x, size_t len) {
  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 probabilites of all possible unfinished "
                 "sequences must be greater than 0.");
  for (size_t i = 0; i < len; ++i) x[i] /= sum;
  return sum;
}
}  // namespace

36 37 38
using framework::LoDTensor;
using framework::LoD;

C
caoying03 已提交
39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72
class LinearChainCrfOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
  LinearChainCrfOpMaker(framework::OpProto* proto,
                        framework::OpAttrChecker* op_checker)
      : OpProtoAndCheckerMaker(proto, op_checker) {
    AddInput(
        "Emission",
        "(LoDTensor, default: LoDTensor<float>). "
        "The unscaled emission weight matrix for the linear chain CRF. "
        "This input is a LoDTensor with shape [N x D] where N is the total "
        "element number of all input squences in a mini-batch, "
        "and D is the total tag number.");
    AddInput(
        "Transition",
        "(Tensor, default: Tensor<float>). A Tensor with shape [(D + 2) x D]. "
        "The learnable parameter for linear_chain_crf operator. "
        "See more details in the operator's comments.");
    AddInput(
        "Label",
        "(LoDTensor, default: LoDTensor<int>). The ground truth which is a 2-D "
        "LoDTensor with shape [N x 1], where N is the total element number in "
        "a mini-batch.");
    AddOutput(
        "Alpha",
        "Tensor, default: Tensor<float>. The forward vectors for the entire "
        "batch. A two dimensional tensor with shape [N x D], "
        "denoted as \f$\alpha\f$. \f$\alpha$\f is a memo table used to "
        "calculate the normalization factor in CRF. \f$\alpha[k, v]$\f stores "
        "the unnormalized probabilites of all possible unfinished sequences of "
        "tags that end at position \f$k$\f with tag \f$v$\f. For each \f$k$\f, "
        "\f$\alpha[k, v]$\f is a vector of length \f$D$\f with a component for "
        "each tag value \f$v$\f. This vector is called a forward vecotr and "
        "will also be used in backward computations.")
        .AsIntermediate();
C
caoying03 已提交
73 74 75 76 77 78 79 80 81 82
    AddOutput("EmissionExps",
              "The exponentials of Input(Emission). This is an intermediate "
              "computational result in forward computation, and will be reused "
              "in backward computation.")
        .AsIntermediate();
    AddOutput("TransitionExps",
              "The exponentials of Input(Transition). This is an intermediate "
              "computational result in forward computation, and will be reused "
              "in backward computation.")
        .AsIntermediate();
C
caoying03 已提交
83 84
    AddOutput(
        "LogLikelihood",
C
caoying03 已提交
85 86
        "(Tensor, default: Tensor<float>). The logarithm of the "
        "conditional "
C
caoying03 已提交
87 88 89
        "likelihood of each training sample in a mini-batch. This is a 2-D "
        "tensor with shape [S x 1], where S is the sequence number in a "
        "mini-batch. "
C
caoying03 已提交
90 91
        "Note: S is equal to the sequence number in a mini-batch. The "
        "output "
C
caoying03 已提交
92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110
        "is no longer a LoDTensor.");
    AddComment(R"DOC(
Conditional Random Field defines an undirected probabilistic graph with nodes
denoting random variables and edges denoting dependencies between these
variables. CRF learns the conditional probability \f$P(Y|X)\f$, where
\f$X = (x_1, x_2, ... , x_n)\f$ are structured inputs and
\f$Y = (y_1, y_2, ... , y_n)\f$ are labels for the inputs.

Linear chain CRF is a special case of CRF that is useful for sequence labeling
task. Sequence labeling tasks do not assume a lot of conditional
independences among inputs. They only concern about the input and the output
being linear sequences. Thus, the graph model of CRF is a simple chain or
a line, which results in a linear chain CRF.

This operator implements the Forward-Backward algorithm for linear chain CRF.
Please see http://www.cs.columbia.edu/~mcollins/fb.pdf for reference.

Equation:

111 112 113 114 115 116 117 118
- Denote Input(Emission) to this operator as \f$x\f$ here.
- The first D values of Input(Transition) to this operator are for starting
weights, denoted as \f$a\f$ here.
- The next D values of Input(Transition) of this operator are for ending
weights, denoted as \f$b\f$ here.
- The remaning values of Input(Transition) are for transition weights,
denoted as \f$w\f$ here.
- Denote Input(Label) as \f$s\f$ here.
C
caoying03 已提交
119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140

The probability of a sequence \f$s\f$ of length \f$L\f$ is defined as:
\f$P(s) = (1/Z) exp(a_{s_1} + b_{s_L}
                 + \sum_{l=1}^L x_{s_l}
                 + \sum_{l=2}^L w_{s_{l-1},s_l})\f$
where \f$Z\f$ is a normalization value so that the sum of \f$P(s)\f$ over
all possible sequences is \f$1\f$, and \f$x\f$ is the emission feature weight
to the linear chain CRF.

Finaly, the linear chain CRF operator outputs the logarithm of the conditional
likelihood of each training sample in a mini-batch.

NOTE:
1. The feature function for a CRF is made up of the emission features and the
transition features. The emission feature weights are NOT computed in
this operator. They MUST be computed first before this operator is called.

2. Because this operator performs globally normaliztion over all possible
sequences internally, it expects UNSCALED emission feature weights.
Please do not call this op with the emission feature being output of any
nonlinear activation.

141
3. The 2nd dimension of Input(Emission) MUST be equal to the tag number.
C
caoying03 已提交
142 143 144 145 146 147 148 149 150

)DOC");
  }
};

class LinearChainCrfOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

C
caoying03 已提交
151 152 153 154 155 156 157 158 159
  void InferShape(framework::InferShapeContext* ctx) const override {
    PADDLE_ENFORCE(ctx->HasInput("Emission"),
                   "Input(Emission) should be not null.");
    PADDLE_ENFORCE(ctx->HasInput("Transition"),
                   "Input(Transition) should be not null.");
    PADDLE_ENFORCE(ctx->HasInput("Label"), "Input(Label) should be not null.");

    PADDLE_ENFORCE(ctx->HasOutput("Alpha"),
                   "Output(Alpha) should be not null.");
C
caoying03 已提交
160 161 162 163
    PADDLE_ENFORCE(ctx->HasOutput("EmissionExps"),
                   "Output(EmissionExps) should be not null.");
    PADDLE_ENFORCE(ctx->HasOutput("TransitionExps"),
                   "Output(TransitionExps) should be not null.");
C
caoying03 已提交
164 165 166 167 168
    PADDLE_ENFORCE(ctx->HasOutput("LogLikelihood"),
                   "Output(LogLikelihood) should be not null.");

    auto emission_dims = ctx->GetInputDim("Emission");
    PADDLE_ENFORCE_EQ(emission_dims.size(), 2UL,
169
                      "The Input(Emission) should be a 2-D tensor.");
C
caoying03 已提交
170 171 172
    PADDLE_ENFORCE(emission_dims[0], "An empty mini-batch is not allowed.");

    auto transition_dims = ctx->GetInputDim("Transition");
C
caoying03 已提交
173
    PADDLE_ENFORCE_EQ(transition_dims.size(), 2UL,
174
                      "The Input(Transition) should be a 2-D tensor.");
C
caoying03 已提交
175
    PADDLE_ENFORCE_EQ(
176 177
        transition_dims[0] - 2, transition_dims[1],
        "An invalid dimension for the Input(Transition), which should "
C
caoying03 已提交
178
        "be a 2-D tensor with shape [(D + 2) x D].");
C
caoying03 已提交
179 180
    PADDLE_ENFORCE_EQ(
        emission_dims[1], transition_dims[1],
181
        "The 2nd dimension of the Input(Emission) and the Input(Transition) "
C
caoying03 已提交
182
        "should be equal to the tag number.");
C
caoying03 已提交
183 184

    auto label_dims = ctx->GetInputDim("Label");
C
caoying03 已提交
185
    PADDLE_ENFORCE(label_dims.size() == 2UL && label_dims[1] == 1UL,
186 187 188 189 190 191
                   "The Input(Label) should be a 2-D tensor with the 2nd "
                   "dimensions fixed to 1.");
    PADDLE_ENFORCE_EQ(
        emission_dims[0], label_dims[0],
        "The height of Input(Emission) and the height of Input(Label) "
        "should be the same.");
C
caoying03 已提交
192 193

    ctx->SetOutputDim("Alpha", emission_dims);
C
caoying03 已提交
194 195
    ctx->SetOutputDim("EmissionExps", emission_dims);
    ctx->SetOutputDim("TransitionExps", transition_dims);
196 197 198
    // (TODO caoying) This is tricky. The 1st dimension of Output(LogLikelihood)
    // is the sequence number in a mini-batch. The dimension set here should be
    // resized to its correct size in the function Compute.
C
caoying03 已提交
199
    ctx->SetOutputDim("LogLikelihood", {emission_dims[0], 1});
C
caoying03 已提交
200 201

    ctx->ShareLoD("Emission", /*->*/ "EmissionExps");
C
caoying03 已提交
202 203
  }

C
caoying03 已提交
204
 protected:
205 206
  // Explicitly set that the data type of output of the linear_chain_crf
  // operator is determined by its input "Emission".
C
caoying03 已提交
207 208
  framework::DataType IndicateDataType(
      const framework::ExecutionContext& ctx) const override {
C
caoying03 已提交
209
    return framework::ToDataType(ctx.Input<LoDTensor>("Emission")->type());
C
caoying03 已提交
210
  }
C
caoying03 已提交
211 212
};

213 214 215 216 217 218 219 220 221
template <typename T>
class LinearChainCrfOpKernel<platform::CPUPlace, T>
    : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
    PADDLE_ENFORCE(platform::is_cpu_place(ctx.GetPlace()),
                   "This kernel only runs on CPU.");
    auto* emission_weights = ctx.Input<LoDTensor>("Emission");
    auto* transition_weights = ctx.Input<Tensor>("Transition");
C
caoying03 已提交
222 223 224 225
    auto* emission_exps = ctx.Output<LoDTensor>("EmissionExps");
    emission_exps->mutable_data<T>(platform::CPUPlace());
    auto* transition_exps = ctx.Output<Tensor>("TransitionExps");
    transition_exps->mutable_data<T>(platform::CPUPlace());
226 227 228
    auto* label = ctx.Input<LoDTensor>("Label");

    auto in_lod = emission_weights->lod();
C
caoying03 已提交
229 230
    PADDLE_ENFORCE(in_lod.size(), "Input(Emission) is not a sequence.");

231 232 233 234 235 236 237 238 239
    // TODO(caoying) The checks related to LoD information should be
    // moved into InferShape once after the InferShape is refactored.
    PADDLE_ENFORCE_EQ(emission_weights->NumLevels(), 1UL,
                      "The Input(Emission) should be a sequence.");
    PADDLE_ENFORCE_EQ(label->NumLevels(), 1UL,
                      "The Input(Label) should be a sequence.");
    const size_t level = 0;

    auto emission_dims = emission_weights->dims();
C
caoying03 已提交
240 241
    const size_t batch_size = emission_dims[0];
    const size_t tag_num = emission_dims[1];
242 243 244 245
    const size_t seq_num = in_lod[level].size() - 1;

    Tensor emission_row_max;
    emission_row_max.mutable_data<T>(
C
caoying03 已提交
246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262
        framework::make_ddim({static_cast<int>(batch_size), 1}),
        platform::CPUPlace());

    auto place = ctx.GetEigenDevice<platform::CPUPlace>();
    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();
263

C
caoying03 已提交
264
    auto* alpha = ctx.Output<LoDTensor>("Alpha");
265
    alpha->mutable_data<T>(ctx.GetPlace());
C
caoying03 已提交
266
    auto* ll = ctx.Output<LoDTensor>("LogLikelihood");
267 268 269 270 271 272
    // resize the output tensor to the correct dimension.
    ll->Resize({static_cast<int>(seq_num), 1});
    T* log_likelihood = ll->mutable_data<T>(ctx.GetPlace());
    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]);
C
caoying03 已提交
273 274
      if (end_pos == start_pos) {
        // If an empty input sequence is given, pad 0 for its cost.
C
caoying03 已提交
275
        log_likelihood[i] = 0.;
C
caoying03 已提交
276 277
        continue;
      }
278

C
caoying03 已提交
279 280 281 282 283
      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);
284 285

      log_likelihood[i] = ForwardOneSequence(
C
caoying03 已提交
286 287
          &one_seq, &one_seq_row_max, &one_seq_exps, transition_weights,
          transition_exps, &one_seq_label, &one_seq_alpha);
288 289 290 291
    }
  }

 protected:
C
caoying03 已提交
292 293 294 295 296 297 298 299 300 301 302 303
  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();
304 305 306 307
    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.
C
caoying03 已提交
308
    // Transition weights among other tags begin from the 3rd row of w.
C
caoying03 已提交
309
    const size_t state_trans_base_idx = 2;
310 311

    for (size_t i = 0; i < tag_num; ++i) {
C
caoying03 已提交
312
      alpha_value[i] = w_exps[i] * x_exps[i];
313
    }
C
caoying03 已提交
314
    T ll = -x_row_max[0] - std::log(NormalizeL1<T>(alpha_value, tag_num));
315 316 317

    for (size_t k = 1; k < seq_length; ++k) {
      for (size_t i = 0; i < tag_num; ++i) {
C
caoying03 已提交
318
        T sum = 0.;
319 320
        for (size_t j = 0; j < tag_num; ++j) {
          sum += alpha_value[(k - 1) * tag_num + j] *
C
caoying03 已提交
321
                 w_exps[(j + state_trans_base_idx) * tag_num + i];
322
        }
C
caoying03 已提交
323
        alpha_value[k * tag_num + i] = x_exps[k * tag_num + i] * sum;
324
      }
C
caoying03 已提交
325 326
      ll -= x_row_max[k] +
            std::log(NormalizeL1<T>(alpha_value + k * tag_num, tag_num));
327 328 329
    }
    T sum = 0.;
    for (size_t i = 0; i < tag_num; ++i) {
C
caoying03 已提交
330
      sum += alpha_value[(seq_length - 1) * tag_num + i] * w_exps[tag_num + i];
331 332 333
    }
    ll -= std::log(sum);

C
caoying03 已提交
334
    const int* lbl = label->data<int>();
335 336 337
    PADDLE_ENFORCE_LT(
        *std::max_element(lbl, lbl + seq_length), tag_num,
        "An invalid tag label that execesses the largest tag number.");
C
caoying03 已提交
338

339
    // Calculate the nominator part, which depends on the label sequence.
C
caoying03 已提交
340 341
    ll += w[lbl[0]] /*start transition*/ + x[lbl[0]] +
          w[tag_num + lbl[seq_length - 1]] /*end transition*/;
C
caoying03 已提交
342 343 344 345
    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]];
    }
346 347 348 349
    return -ll;
  }
};

C
caoying03 已提交
350 351 352 353
class LinearChainCrfGradOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

C
caoying03 已提交
354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369
  void InferShape(framework::InferShapeContext* ctx) const override {
    PADDLE_ENFORCE(ctx->HasInput("EmissionExps"),
                   "Input(EmissionExps) should be not null.");
    PADDLE_ENFORCE(ctx->HasInput("TransitionExps"),
                   "Input(TransitionExps) should be not null.");
    PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("LogLikelihood")),
                   "Input(LogLikelihood@GRAD) shoudl be not null.");

    PADDLE_ENFORCE(ctx->HasOutput(framework::GradVarName("Emission")),
                   "Output(Emission@GRAD) should be not null.");
    PADDLE_ENFORCE(ctx->HasOutput(framework::GradVarName("Transition")),
                   "Output(Transition@GRAD) should be not null.");

    auto emission_exps_dims = ctx->GetInputDim("EmissionExps");
    PADDLE_ENFORCE_EQ(emission_exps_dims.size(), 2UL,
                      "The Input(EmissionExps) should be a 2-D tensor.");
C
caoying03 已提交
370 371 372 373 374
    PADDLE_ENFORCE(emission_exps_dims[0],
                   "An empty mini-batch is not allowed.");

    auto transition_exps_dims =
        ctx->GetInputDim(framework::GradVarName("TransitionExps"));
C
caoying03 已提交
375 376 377 378 379 380 381 382 383 384
    PADDLE_ENFORCE_EQ(transition_exps_dims.size(), 2UL,
                      "The Input(TransitionExps) should be a 2-D tensor.");
    PADDLE_ENFORCE_EQ(
        transition_exps_dims[0] - 2, transition_exps_dims[1],
        "An invalid dimension for the Input(TransitionExps), which should "
        "be a 2-D tensor with shape [(D + 2) x D].");
    PADDLE_ENFORCE_EQ(
        emission_exps_dims[1], transition_exps_dims[1],
        "The 2nd dimension of the Input(EmissionExps) and the "
        "Input(TransitionExps) should be equal to the tag number.");
C
caoying03 已提交
385 386

    auto label_dims = ctx->GetInputDim("Label");
C
caoying03 已提交
387 388 389 390 391 392 393 394 395 396 397 398
    PADDLE_ENFORCE(label_dims.size() == 2UL && label_dims[1] == 1UL,
                   "The Input(Label) should be a 2-D tensor with the 2nd "
                   "dimensions fixed to 1.");
    PADDLE_ENFORCE_EQ(
        emission_exps_dims[0], label_dims[0],
        "The height of Input(EmissionExps) and the height of Input(Label) "
        "should be the same.");

    ctx->SetOutputDim(framework::GradVarName("Emission"), emission_exps_dims);
    ctx->SetOutputDim(framework::GradVarName("Transition"),
                      transition_exps_dims);
  }
C
caoying03 已提交
399 400 401 402 403 404 405 406

 protected:
  // Explicitly set that the data type of output of the linear_chain_crf_grad
  // operator is determined by its input "EmissionExps".
  framework::DataType IndicateDataType(
      const framework::ExecutionContext& ctx) const override {
    return framework::ToDataType(ctx.Input<LoDTensor>("EmissionExps")->type());
  }
C
caoying03 已提交
407 408
};

409 410 411 412 413 414 415
template <typename T>
class LinearChainCrfGradOpKernel<platform::CPUPlace, T>
    : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
    PADDLE_ENFORCE(platform::is_cpu_place(ctx.GetPlace()),
                   "This kernel only runs on CPU.");
C
caoying03 已提交
416 417 418
    auto* label = ctx.Input<LoDTensor>("Label");
    auto* emission_exps = ctx.Input<LoDTensor>("EmissionExps");
    auto* transition_exps = ctx.Input<Tensor>("TransitionExps");
C
caoying03 已提交
419 420 421
    auto* alpha = ctx.Input<LoDTensor>("Alpha");
    const T* ll_grad =
        ctx.Input<Tensor>(framework::GradVarName("LogLikelihood"))->data<T>();
C
caoying03 已提交
422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439

    auto* emission_grad =
        ctx.Output<Tensor>(framework::GradVarName("Emission"));
    emission_grad->mutable_data<T>(platform::CPUPlace());

    auto* trans_grad = ctx.Output<Tensor>(framework::GradVarName("Transition"));
    if (trans_grad) trans_grad->mutable_data<T>(platform::CPUPlace());

    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
    // captures the unnormalized probabilities of partial sequences starting at
    // position i.
    Tensor beta;
    beta.mutable_data<T>(emission_dims, platform::CPUPlace());

    const size_t level = 0;  // currently, only support sequence.
C
caoying03 已提交
440 441 442
    auto lod = label->lod();
    PADDLE_ENFORCE(lod.size(), "Input(Label) is not a sequence.");

C
caoying03 已提交
443 444 445
    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]);
C
caoying03 已提交
446
      if (end_pos == start_pos) continue;
C
caoying03 已提交
447 448

      const Tensor one_seq_emission_exps =
C
caoying03 已提交
449 450 451 452 453 454 455 456 457 458
          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);

      BackwardOneSequence(ctx.device_context(), ll_grad[i],
                          &one_seq_emission_exps, transition_exps,
                          &one_seq_alpha, &one_seq_label, &one_seq_beta,
                          trans_grad, &one_seq_emission_grad);
C
caoying03 已提交
459 460 461 462
    }
  }

 protected:
C
caoying03 已提交
463
  void BackwardOneSequence(const platform::DeviceContext& ctx, const T ll_grad,
C
caoying03 已提交
464 465 466 467 468 469 470 471 472 473 474 475 476 477 478
                           const Tensor* emission_exps,
                           const Tensor* transition_exps, const Tensor* alpha,
                           const Tensor* label, Tensor* beta,
                           Tensor* transition_grad,
                           Tensor* emission_grad) const {
    const T* w_exps = transition_exps->data<T>();
    const T* x_exps = emission_exps->data<T>();
    const int* label_value = label->data<int>();
    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;

C
caoying03 已提交
479
    // Calculate the backward vectors: beta.
C
caoying03 已提交
480
    // First, calculate the initialition state.
C
caoying03 已提交
481
    for (size_t i = 0; i < tag_num; ++i) {
C
caoying03 已提交
482
      beta_value[(seq_length - 1) * tag_num + i] = w_exps[tag_num + i];
C
caoying03 已提交
483
    }
C
caoying03 已提交
484
    NormalizeL1<T>(beta_value + (seq_length - 1) * tag_num, tag_num);
C
caoying03 已提交
485

C
caoying03 已提交
486 487 488 489
    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 已提交
490 491 492
          sum += w_exps[(i + state_trans_base_idx) * tag_num + j] *
                 x_exps[(k + 1) * tag_num + j] *
                 beta_value[(k + 1) * tag_num + j];
C
caoying03 已提交
493 494 495 496 497 498 499 500 501 502
        }
        beta_value[k * tag_num + i] = sum;
      }
      NormalizeL1<T>(beta_value + k * tag_num, tag_num);
    }

    auto alpha_mat = EigenMatrix<T>::From(*alpha);
    auto beta_mat = EigenMatrix<T>::From(*beta);
    auto x_grad_mat = EigenMatrix<T>::From(*emission_grad);
    auto* place = ctx.GetEigenDevice<platform::CPUPlace>();
C
caoying03 已提交
503 504 505 506 507 508 509 510
    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));
    x_grad_mat.device(*place) = prob / row_sum;

    for (size_t k = 0; k < seq_length; ++k) {
      x_grad_mat(k, label_value[k]) -= static_cast<T>(1.);
C
caoying03 已提交
511
    }
C
caoying03 已提交
512 513 514 515 516 517 518 519 520 521

    if (transition_grad) {
      T* trans_grad = transition_grad->data<T>();
      for (size_t k = 0; k < tag_num; ++k) {
        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);
C
caoying03 已提交
522 523 524 525 526 527 528 529 530 531 532 533 534 535 536

      // TODO(caoying): Fix this to avoid using this local variable.
      Tensor tmp;
      tmp.mutable_data<T>(beta->dims(), platform::CPUPlace());
      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) {
C
caoying03 已提交
537
            sum += w_exps[(i + state_trans_base_idx) * tag_num + j] *
C
caoying03 已提交
538
                   alpha_mat(k - 1, i) * tmp_mat(k, j);
C
caoying03 已提交
539
          }
C
caoying03 已提交
540
        }
C
caoying03 已提交
541 542 543
        sum = 1. / sum;
        for (size_t i = 0; i < tag_num; ++i) {
          for (size_t j = 0; j < tag_num; ++j) {
C
caoying03 已提交
544 545
            trans_grad[(i + state_trans_base_idx) * tag_num + j] +=
                sum * w_exps[(i + state_trans_base_idx) * tag_num + j] *
C
caoying03 已提交
546
                alpha_mat(k - 1, i) * tmp_mat(k, j);
C
caoying03 已提交
547 548
          }
        }
C
caoying03 已提交
549 550
        trans_grad[(label_value[k - 1] + state_trans_base_idx) * tag_num +
                   label_value[k]] -= static_cast<T>(1.);
C
caoying03 已提交
551 552
      }
    }
553 554 555
  }
};

C
caoying03 已提交
556 557 558 559 560 561
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
REGISTER_OP(linear_chain_crf, ops::LinearChainCrfOp, ops::LinearChainCrfOpMaker,
            linear_chain_crf_grad, ops::LinearChainCrfGradOp);
562 563
REGISTER_OP_CPU_KERNEL(
    linear_chain_crf,
C
caoying03 已提交
564
    ops::LinearChainCrfOpKernel<paddle::platform::CPUPlace, float>);
565 566
REGISTER_OP_CPU_KERNEL(
    linear_chain_crf_grad,
C
caoying03 已提交
567
    ops::LinearChainCrfGradOpKernel<paddle::platform::CPUPlace, float>);