wanghaox
7 years ago
commit
0968c7cd6b
@ -0,0 +1,185 @@
|
||||
/* 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/beam_search_op.h"
|
||||
|
||||
#include <map>
|
||||
#include "paddle/framework/lod_tensor.h"
|
||||
#include "paddle/framework/op_registry.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace operators {
|
||||
|
||||
void BeamSearch::operator()(const framework::LoDTensor &pre_ids,
|
||||
framework::LoDTensor *selected_ids,
|
||||
framework::LoDTensor *selected_scores) {
|
||||
auto items = SelectTopBeamSizeItems();
|
||||
auto selected_items = ToMap(items);
|
||||
PruneEndidCandidates(pre_ids, &selected_items);
|
||||
// calculate the output tensor's height
|
||||
size_t num_instances = std::accumulate(
|
||||
std::begin(items), std::end(items), 0,
|
||||
[](size_t a, std::vector<Item> &b) { return a + b.size(); });
|
||||
// the output tensor shape should be [num_instances, 1]
|
||||
auto dims = framework::make_ddim(
|
||||
std::vector<int64_t>({static_cast<int>(num_instances), 1}));
|
||||
selected_ids->Resize(dims);
|
||||
selected_scores->Resize(dims);
|
||||
|
||||
std::map<size_t /*offset*/, std::vector<Item>> hash;
|
||||
framework::LoD new_lod;
|
||||
auto *ids_data = selected_ids->mutable_data<int>(platform::CPUPlace());
|
||||
auto *scores_data =
|
||||
selected_scores->mutable_data<float>(platform::CPUPlace());
|
||||
|
||||
// fill in data
|
||||
std::vector<size_t> low_level;
|
||||
size_t low_offset = 0;
|
||||
for (auto &items : selected_items) {
|
||||
low_level.push_back(low_offset);
|
||||
for (auto &item : items) {
|
||||
ids_data[low_offset] = item.id;
|
||||
scores_data[low_offset] = item.score;
|
||||
low_offset++;
|
||||
}
|
||||
}
|
||||
// fill lod
|
||||
auto abs_lod = framework::ToAbsOffset(ids_->lod());
|
||||
auto &high_level = abs_lod[lod_level_];
|
||||
framework::LoD lod(2);
|
||||
lod[0].assign(high_level.begin(), high_level.end());
|
||||
lod[1].assign(low_level.begin(), low_level.end());
|
||||
selected_ids->set_lod(lod);
|
||||
selected_scores->set_lod(lod);
|
||||
}
|
||||
|
||||
void BeamSearch::PruneEndidCandidates(const framework::LoDTensor &pre_ids,
|
||||
std::vector<std::vector<Item>> *items) {
|
||||
auto *pre_ids_data = pre_ids.data<int>();
|
||||
|
||||
for (size_t offset = 0; offset < items->size(); offset++) {
|
||||
auto prefix_id = pre_ids_data[offset];
|
||||
if (prefix_id == end_id_) {
|
||||
items->at(offset).clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::vector<BeamSearch::Item>> BeamSearch::ToMap(
|
||||
const std::vector<std::vector<Item>> &items) {
|
||||
std::vector<std::vector<Item>> result;
|
||||
for (auto &entries : items) {
|
||||
for (const auto &item : entries) {
|
||||
if (item.offset >= result.size()) {
|
||||
result.resize(item.offset + 1);
|
||||
}
|
||||
result[item.offset].push_back(item);
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<std::vector<BeamSearch::Item>>
|
||||
BeamSearch::SelectTopBeamSizeItems() {
|
||||
std::vector<std::vector<Item>> result;
|
||||
std::vector<Item> items;
|
||||
// for each source sentence, select the top beam_size items across all
|
||||
// candidate sets.
|
||||
while (NextItemSet(&items)) {
|
||||
std::nth_element(std::begin(items), std::begin(items) + beam_size_,
|
||||
std::end(items), [](const Item &a, const Item &b) {
|
||||
// TODO(superjom) make score's comparation customizable.
|
||||
// partial sort in descending order
|
||||
return a.score > b.score;
|
||||
});
|
||||
// prune the top beam_size items.
|
||||
if (items.size() > beam_size_) {
|
||||
items.resize(beam_size_);
|
||||
}
|
||||
result.emplace_back(items);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// the candidates of a source
|
||||
bool BeamSearch::NextItemSet(std::vector<BeamSearch::Item> *items) {
|
||||
if (sent_offset_ >= ids_->NumElements(lod_level_)) {
|
||||
return false;
|
||||
}
|
||||
// find the current candidates
|
||||
auto ids = *ids_;
|
||||
auto scores = *scores_;
|
||||
|
||||
auto source_abs_two_level_lod = framework::SliceInLevel(
|
||||
ids.lod(), lod_level_, sent_offset_, sent_offset_ + 1);
|
||||
source_abs_two_level_lod = framework::ToAbsOffset(source_abs_two_level_lod);
|
||||
auto abs_lod = framework::ToAbsOffset(ids.lod());
|
||||
PADDLE_ENFORCE_GE(source_abs_two_level_lod.size(), 2UL);
|
||||
|
||||
auto *ids_data = ids.data<int>();
|
||||
auto *scores_data = scores.data<float>();
|
||||
|
||||
size_t instance_dim = 1;
|
||||
for (int i = 1; i < ids.dims().size(); i++) {
|
||||
instance_dim *= ids.dims()[i];
|
||||
}
|
||||
|
||||
items->clear();
|
||||
items->reserve(framework::product(ids.dims()));
|
||||
for (size_t offset = abs_lod[lod_level_][sent_offset_];
|
||||
offset < abs_lod[lod_level_][sent_offset_ + 1]; offset++) {
|
||||
for (int d = 0; d < instance_dim; d++) {
|
||||
const size_t dim_offset = offset * instance_dim + d;
|
||||
items->emplace_back(offset, ids_data[dim_offset],
|
||||
scores_data[dim_offset]);
|
||||
}
|
||||
}
|
||||
|
||||
sent_offset_++;
|
||||
return true;
|
||||
}
|
||||
|
||||
class BeamSearchProtoAndCheckerMaker
|
||||
: public framework::OpProtoAndCheckerMaker {
|
||||
public:
|
||||
BeamSearchProtoAndCheckerMaker(framework::OpProto *proto,
|
||||
framework::OpAttrChecker *op_checker)
|
||||
: OpProtoAndCheckerMaker(proto, op_checker) {
|
||||
// inputs and outputs stored in proto
|
||||
AddInput("pre_ids", "ids in previous step");
|
||||
AddInput("ids", "a LoDTensor of shape of [None,k]");
|
||||
AddInput("scores",
|
||||
"a LoDTensor that has the same shape and LoD with `ids`");
|
||||
AddOutput("selected_ids",
|
||||
"a LoDTensor that stores the IDs selected by beam search");
|
||||
AddOutput(
|
||||
"selected_scores",
|
||||
"a LoDTensor that has the same shape and LoD with `selected_ids`");
|
||||
|
||||
// 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");
|
||||
|
||||
AddComment(
|
||||
"This is a beam search operator that help to generate sequences.");
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace operators
|
||||
} // namespace paddle
|
||||
|
||||
REGISTER_OP_WITHOUT_GRADIENT(beam_search, paddle::operators::BeamSearchOp,
|
||||
paddle::operators::BeamSearchProtoAndCheckerMaker);
|
||||
@ -0,0 +1,226 @@
|
||||
/* 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
|
||||
|
||||
#ifdef PADDLE_WITH_TESTING
|
||||
#include "gtest/gtest.h"
|
||||
#endif
|
||||
|
||||
#include "paddle/framework/lod_tensor.h"
|
||||
#include "paddle/framework/operator.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace operators {
|
||||
|
||||
/*
|
||||
* This is an implementation of beam search.
|
||||
*
|
||||
* To explain the details, lets take machine translation task for example, in
|
||||
* this task, one source sentence is translated to multiple target sentences,
|
||||
* during this period, one sentence will be translated to multiple translation
|
||||
* prefixes(target sentence that have not ended), in each time step a prefix
|
||||
* will have some candidates, input the candidate ids and their corresponding
|
||||
* scores (probabilities), it will sort and select the top beam_size candidates
|
||||
* for each source sentence, and store the selected candidates's score and their
|
||||
* corresponding ids to LoDTensors.
|
||||
*
|
||||
* A detailed example:
|
||||
*
|
||||
* Input
|
||||
*
|
||||
* ids:
|
||||
* LoD (should have 2 levels)
|
||||
* first level: [0, 1, 4]
|
||||
* second level: [0, 1, 2, 3, 4]
|
||||
*
|
||||
* tensor's data
|
||||
* [
|
||||
* [4, 2, 5]
|
||||
* [2, 1, 3]
|
||||
* [3, 5, 2]
|
||||
* [8, 2, 1]
|
||||
* ]
|
||||
*
|
||||
* scores:
|
||||
* LoD same as `ids`
|
||||
* tensor's data
|
||||
* [
|
||||
* [0.5, 0.3, 0.2]
|
||||
* [0.6, 0.3, 0.1]
|
||||
* [0.9, 0.5, 0.1]
|
||||
* [0.7, 0.5, 0.1]
|
||||
* ]
|
||||
*
|
||||
* the inputs means that there are 2 source sentences to translate, and the
|
||||
* first source has 1 prefix, the second source has 2 prefix.
|
||||
*
|
||||
* lets assume beam size is 2, and the beam search's output should be
|
||||
* LoD
|
||||
* first level:
|
||||
* [0, 1, 2]
|
||||
* second level:
|
||||
* [0, 2, 4]
|
||||
*
|
||||
* tensor's data
|
||||
* [[
|
||||
* 0.5,
|
||||
* 0.3,
|
||||
* 0.9,
|
||||
* 0.7
|
||||
* ]]
|
||||
*
|
||||
* TODO all the prune operations should be in the beam search, so it is better
|
||||
* to split the beam search algorithm into a sequence of smaller operators, and
|
||||
* the prune operators can be inserted in this sequence.
|
||||
*/
|
||||
class BeamSearch {
|
||||
public:
|
||||
// TODO(superjom) make type customizable
|
||||
using id_t = size_t;
|
||||
using score_t = float;
|
||||
/*
|
||||
* Input the arguments that needed by this class.
|
||||
*/
|
||||
BeamSearch(const framework::LoDTensor& ids,
|
||||
const framework::LoDTensor& scores, size_t level, size_t beam_size,
|
||||
int end_id)
|
||||
: beam_size_(beam_size),
|
||||
ids_(&ids),
|
||||
scores_(&scores),
|
||||
lod_level_(level),
|
||||
end_id_(end_id) {}
|
||||
|
||||
/*
|
||||
* The main function of beam search.
|
||||
*
|
||||
* @selected_ids: a [None, 1]-shaped tensor with LoD.
|
||||
* In a machine translation model, it might be the candidate term id sets,
|
||||
* each set stored as a varience-length sequence.
|
||||
* The format might be described with a two-level LoD
|
||||
* - [[0 1]
|
||||
* - [0 1 2]]
|
||||
* - [[]
|
||||
* - [0 1]]
|
||||
* the first level of LoD tells that there are two source sentences. The
|
||||
* second level describes the details of the candidate id set's offsets in
|
||||
* the
|
||||
* source sentences.
|
||||
*
|
||||
* @selected_scores: a LoD tensor with the same shape and LoD with
|
||||
* selected_ids.
|
||||
* It stores the corresponding scores of candidate ids in selected_ids.
|
||||
*
|
||||
* Return false if all the input tensor is empty, in machine translation task
|
||||
* that means no candidates is provided, and the task will stop running.
|
||||
*/
|
||||
void operator()(const framework::LoDTensor& pre_ids,
|
||||
framework::LoDTensor* selected_ids,
|
||||
framework::LoDTensor* selected_scores);
|
||||
|
||||
protected:
|
||||
/*
|
||||
* The basic items help to sort.
|
||||
*/
|
||||
struct Item {
|
||||
Item() {}
|
||||
Item(size_t offset, size_t id, float score)
|
||||
: offset(offset), id(id), score(score) {}
|
||||
// offset in the lod_level_+1
|
||||
size_t offset;
|
||||
// the candidate id
|
||||
id_t id;
|
||||
// the corresponding score
|
||||
score_t score;
|
||||
};
|
||||
|
||||
void PruneEndidCandidates(const framework::LoDTensor& pre_ids,
|
||||
std::vector<std::vector<Item>>* items);
|
||||
|
||||
/*
|
||||
* Transform the items into a map whose key is offset, value is the items.
|
||||
* NOTE low performance
|
||||
*/
|
||||
std::vector<std::vector<Item>> ToMap(
|
||||
const std::vector<std::vector<Item>>& inputs);
|
||||
|
||||
/*
|
||||
* For each source, select top beam_size records.
|
||||
*/
|
||||
std::vector<std::vector<Item>> SelectTopBeamSizeItems();
|
||||
|
||||
/*
|
||||
* Get the items of next source sequence, return false if no remaining items.
|
||||
*/
|
||||
bool NextItemSet(std::vector<Item>* items);
|
||||
|
||||
private:
|
||||
size_t beam_size_;
|
||||
const framework::LoDTensor* ids_;
|
||||
const framework::LoDTensor* scores_;
|
||||
size_t lod_level_{0};
|
||||
size_t sent_offset_{0};
|
||||
int end_id_{0};
|
||||
};
|
||||
|
||||
class BeamSearchOp : public framework::OperatorBase {
|
||||
public:
|
||||
BeamSearchOp(const std::string& type,
|
||||
const framework::VariableNameMap& inputs,
|
||||
const framework::VariableNameMap& outputs,
|
||||
const framework::AttributeMap& attrs)
|
||||
: OperatorBase(type, inputs, outputs, attrs) {}
|
||||
|
||||
BeamSearchOp(const BeamSearchOp& o)
|
||||
: framework::OperatorBase(
|
||||
static_cast<const framework::OperatorBase&>(o)) {
|
||||
PADDLE_THROW("Not Implemented");
|
||||
}
|
||||
|
||||
void Run(const framework::Scope& scope,
|
||||
const platform::DeviceContext& dev_ctx) const override {
|
||||
LOG(INFO) << "run beam search op";
|
||||
auto ids_var = scope.FindVar(Input("ids"));
|
||||
auto scores_var = scope.FindVar(Input("scores"));
|
||||
auto pre_ids_var = scope.FindVar(Input("pre_ids"));
|
||||
PADDLE_ENFORCE_NOT_NULL(ids_var);
|
||||
PADDLE_ENFORCE_NOT_NULL(scores_var);
|
||||
PADDLE_ENFORCE_NOT_NULL(pre_ids_var);
|
||||
|
||||
auto& ids = ids_var->Get<framework::LoDTensor>();
|
||||
auto& scores = scores_var->Get<framework::LoDTensor>();
|
||||
auto& pre_ids = pre_ids_var->Get<framework::LoDTensor>();
|
||||
size_t level = Attr<int>("level");
|
||||
size_t beam_size = Attr<int>("beam_size");
|
||||
int end_id = Attr<int>("end_id");
|
||||
LOG(INFO) << "init beam search";
|
||||
BeamSearch alg(ids, scores, level, beam_size, end_id);
|
||||
|
||||
LOG(INFO) << "after beam search";
|
||||
auto selected_ids_var = scope.FindVar(Output("selected_ids"));
|
||||
auto selected_scores_var = scope.FindVar(Output("selected_scores"));
|
||||
PADDLE_ENFORCE_NOT_NULL(selected_ids_var);
|
||||
PADDLE_ENFORCE_NOT_NULL(selected_scores_var);
|
||||
auto& selected_ids_tensor =
|
||||
*selected_ids_var->GetMutable<framework::LoDTensor>();
|
||||
auto& selected_scores_tensor =
|
||||
*selected_scores_var->GetMutable<framework::LoDTensor>();
|
||||
LOG(INFO) << "run beam search";
|
||||
alg(pre_ids, &selected_ids_tensor, &selected_scores_tensor);
|
||||
LOG(INFO) << "finish beam search";
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace operators
|
||||
} // namespace paddle
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in new issue