beam_search_op.cc 5.9 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2

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
6

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

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

15 16
#include "paddle/fluid/operators/beam_search_op.h"

17 18
#include <string>
#include <vector>
Y
Yi Wang 已提交
19
#include "paddle/fluid/framework/op_registry.h"
20 21 22 23

namespace paddle {
namespace operators {

K
ktlichkid 已提交
24
class BeamSearchOpMaker : public framework::OpProtoAndCheckerMaker {
25
 public:
Y
Yu Yang 已提交
26
  void Make() override {
27
    // inputs and outputs stored in proto
28 29 30 31
    AddInput("pre_ids",
             "(LoDTensor) The LoDTensor containing the selected ids at the "
             "previous step. It should be a tensor with shape (batch_size, 1) "
             "and lod `[[0, 1, ... , batch_size], [0, 1, ..., batch_size]]` at "
32
             "the first step.");
33 34 35 36 37
    AddInput("pre_scores",
             "(LoDTensor) The LoDTensor containing the accumulated "
             "scores corresponding to the selected ids at the previous step.");
    AddInput("ids",
             "(LoDTensor) The LoDTensor containing the candidates ids. Its "
38 39 40
             "shape should be (batch_size * beam_size, W). If not set, it will "
             "be calculated out according to Input(scores) in this operator.")
        .AsDispensable();
41
    AddInput("scores",
42 43 44 45 46 47 48
             "(LoDTensor) The LoDTensor containing the current scores "
             "corresponding to Input(ids). If Input(ids) is not nullptr, its "
             "shape is the same as that of Input(ids)."
             "If is_accumulated is true, Input(scores) is accumulated scores "
             "and will be used derectedly. Else, each score will be "
             "transformed to the log field and accumulate Input(pre_sores) "
             "first.");
49
    AddOutput("selected_ids",
50 51 52 53
              "A LodTensor that stores the IDs selected by beam search.");
    AddOutput("selected_scores",
              "A LoDTensor containing the accumulated scores corresponding to "
              "Output(selected_ids).");
54 55 56
    AddOutput("parent_idx",
              "A Tensor preserving the selected_ids' parent indice in pre_ids.")
        .AsDispensable();
57 58 59 60 61 62

    // Attributes stored in AttributeMap
    AddAttr<int>("level", "the level of LoDTensor");
    AddAttr<int>("beam_size", "beam size for beam search");
    AddAttr<int>("end_id",
                 "the token id which indicates the end of a sequence");
63 64 65
    AddAttr<bool>("is_accumulated",
                  "Whether the Input(scores) is accumulated scores.")
        .SetDefault(true);
66

67
    AddComment(R"DOC(
68
This operator does the search in beams for one time step.
69 70 71 72 73 74 75 76 77 78 79 80 81
Specifically, it selects the top-K candidate word ids of current step from
Input(ids) according to their Input(scores) for all source sentences,
where K is Attr(beam_size) and Input(ids), Input(scores) are predicted results
from the computation cell. Additionally, Input(pre_ids) and Input(pre_scores)
are the output of beam_search at previous step, they are needed for special use
to handle ended candidate translations. The paths linking prefixes and selected
candidates are organized and reserved in lod.

Note that the Input(scores) passed in should be accumulated scores, and
length penalty should be done with extra operators before calculating the
accumulated scores if needed, also suggest finding top-K before it and
using the top-K candidates following.
)DOC");
82 83 84
  }
};

85
class BeamSearchOp : public framework::OperatorWithKernel {
K
ktlichkid 已提交
86 87
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;
88

K
ktlichkid 已提交
89
  void InferShape(framework::InferShapeContext *ctx) const override {
90
    for (const std::string &arg :
91
         std::vector<std::string>({"pre_ids", "scores"})) {
K
ktlichkid 已提交
92 93
      PADDLE_ENFORCE(ctx->HasInput(arg), "BeamSearch need input argument '%s'",
                     arg);
94 95 96
    }
    for (const std::string &arg :
         std::vector<std::string>({"selected_ids", "selected_scores"})) {
K
ktlichkid 已提交
97
      PADDLE_ENFORCE(ctx->HasOutput(arg),
98 99
                     "BeamSearch need output argument '%s'", arg);
    }
100 101
  }

102
 protected:
103 104
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
105 106 107 108 109 110 111 112 113 114 115 116 117
    auto *scores = ctx.Input<framework::LoDTensor>("scores");
    size_t level = ctx.Attr<int>("level");
    size_t batch_size = scores->lod()[level].size() - 1;
    // The current CUDA kernel only support cases with batch_size < 4.
    // Compute on CPU for cases with batch_size > 4.
    if (batch_size <= 4) {
      return framework::OpKernelType(
          ctx.Input<framework::LoDTensor>("pre_ids")->type(), ctx.GetPlace());
    } else {
      return framework::OpKernelType(
          ctx.Input<framework::LoDTensor>("pre_ids")->type(),
          platform::CPUPlace());
    }
118 119 120
  }
};

Q
Qiao Longfei 已提交
121 122
class BeamSearchInferVarType : public framework::VarTypeInference {
 public:
M
minqiyang 已提交
123 124 125
  void operator()(framework::InferVarTypeContext *ctx) const override {
    for (auto &o : ctx->Output("selected_ids")) {
      ctx->SetType(o, framework::proto::VarType::LOD_TENSOR);
Q
Qiao Longfei 已提交
126
    }
M
minqiyang 已提交
127 128
    for (auto &o : ctx->Output("selected_scores")) {
      ctx->SetType(o, framework::proto::VarType::LOD_TENSOR);
Q
Qiao Longfei 已提交
129 130 131
    }
  }
};
K
ktlichkid 已提交
132

133 134
}  // namespace operators
}  // namespace paddle
K
ktlichkid 已提交
135

K
ktlichkid 已提交
136
namespace ops = paddle::operators;
K
ktlichkid 已提交
137 138 139

REGISTER_OPERATOR(beam_search, ops::BeamSearchOp, ops::BeamSearchOpMaker,
                  ops::BeamSearchInferVarType);
K
ktlichkid 已提交
140 141 142
REGISTER_OP_CPU_KERNEL(
    beam_search,
    ops::BeamSearchOpKernel<paddle::platform::CPUDeviceContext, float>,
K
ktlichkid 已提交
143 144 145
    ops::BeamSearchOpKernel<paddle::platform::CPUDeviceContext, double>,
    ops::BeamSearchOpKernel<paddle::platform::CPUDeviceContext, int>,
    ops::BeamSearchOpKernel<paddle::platform::CPUDeviceContext, int64_t>);
新手
引导
客服 返回
顶部