parent
2b6c0c09d6
commit
312b7786d9
@ -1,114 +0,0 @@
|
|||||||
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
|
|
||||||
|
|
||||||
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/fluid/operators/distributed_ops/lookup_remote_table_op.h"
|
|
||||||
#include "paddle/fluid/framework/var_type_inference.h"
|
|
||||||
|
|
||||||
namespace paddle {
|
|
||||||
namespace operators {
|
|
||||||
|
|
||||||
class LookupRemoteTableOp : public framework::OperatorWithKernel {
|
|
||||||
public:
|
|
||||||
using framework::OperatorWithKernel::OperatorWithKernel;
|
|
||||||
|
|
||||||
void InferShape(framework::InferShapeContext* ctx) const override {
|
|
||||||
PADDLE_ENFORCE(ctx->HasInput("W"),
|
|
||||||
"Input(W) of LookupRemoteTableOp should not be null.");
|
|
||||||
PADDLE_ENFORCE(ctx->HasInput("Ids"),
|
|
||||||
"Input(Ids) of LookupRemoteTableOp should not be null.");
|
|
||||||
PADDLE_ENFORCE(ctx->HasOutput("Out"),
|
|
||||||
"Output(Out) of LookupRemoteTableOp should not be null.");
|
|
||||||
|
|
||||||
auto table_dims = ctx->GetInputDim("W");
|
|
||||||
auto ids_dims = ctx->GetInputDim("Ids");
|
|
||||||
int ids_rank = ids_dims.size();
|
|
||||||
|
|
||||||
PADDLE_ENFORCE_EQ(table_dims.size(), 2);
|
|
||||||
PADDLE_ENFORCE_EQ(ids_dims[ids_rank - 1], 1,
|
|
||||||
"The last dimension of the 'Ids' tensor must be 1.");
|
|
||||||
|
|
||||||
auto output_dims =
|
|
||||||
framework::vectorize(framework::slice_ddim(ids_dims, 0, ids_rank - 1));
|
|
||||||
output_dims.push_back(table_dims[1]);
|
|
||||||
ctx->SetOutputDim("Out", framework::make_ddim(output_dims));
|
|
||||||
|
|
||||||
if (ctx->GetOutputsVarType("Out")[0] ==
|
|
||||||
framework::proto::VarType::LOD_TENSOR) {
|
|
||||||
ctx->ShareLoD("Ids", /*->*/ "Out");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
protected:
|
|
||||||
framework::OpKernelType GetExpectedKernelType(
|
|
||||||
const framework::ExecutionContext& ctx) const override {
|
|
||||||
auto data_type = framework::GetDataTypeOfVar(ctx.InputVar("W"));
|
|
||||||
return framework::OpKernelType(data_type, ctx.device_context());
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
class LookupRemoteTableOpMaker : public framework::OpProtoAndCheckerMaker {
|
|
||||||
public:
|
|
||||||
void Make() override {
|
|
||||||
AddInput("W",
|
|
||||||
"(Tensor) The input represents embedding tensors, "
|
|
||||||
"which is a learnable parameter.");
|
|
||||||
AddInput("Ids",
|
|
||||||
"An input with type int32 or int64 "
|
|
||||||
"contains the ids to be looked up in W. "
|
|
||||||
"The last dimension size must be 1.");
|
|
||||||
AddOutput("Out", "The lookup results, which have the same type as W.");
|
|
||||||
AddAttr<std::vector<int64_t>>("height_sections",
|
|
||||||
"Height for each output SelectedRows.")
|
|
||||||
.SetDefault(std::vector<int64_t>({}));
|
|
||||||
AddAttr<int>("trainer_id", "trainer id from 0 ~ worker_num.").SetDefault(0);
|
|
||||||
AddAttr<std::vector<std::string>>(
|
|
||||||
"epmap",
|
|
||||||
"(string vector, default 127.0.0.1:6164)"
|
|
||||||
"Server endpoints in the order of input variables for mapping")
|
|
||||||
.SetDefault({"127.0.0.1:6164"});
|
|
||||||
AddAttr<int64_t>("padding_idx",
|
|
||||||
"(int64, default -1) "
|
|
||||||
"If the value is -1, it makes no effect to lookup. "
|
|
||||||
"Otherwise the given value indicates padding the output "
|
|
||||||
"with zeros whenever lookup encounters it in Ids.")
|
|
||||||
.SetDefault(kNoPadding);
|
|
||||||
// NOTE(minqiyang): grad_inplace is an temporal attribute,
|
|
||||||
// please do NOT set this attribute in python layer.
|
|
||||||
AddAttr<bool>("grad_inplace",
|
|
||||||
"(boolean, default false) "
|
|
||||||
"If the grad op reuse the input's variable.")
|
|
||||||
.SetDefault(false);
|
|
||||||
AddComment(R"DOC(
|
|
||||||
Lookup Remote Table Operator.
|
|
||||||
|
|
||||||
This operator is used to perform lookups on the parameter W,
|
|
||||||
then concatenated into a dense tensor.
|
|
||||||
|
|
||||||
The input Ids can carry the LoD (Level of Details) information,
|
|
||||||
or not. And the output only shares the LoD information with input Ids.
|
|
||||||
|
|
||||||
)DOC");
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
} // namespace operators
|
|
||||||
} // namespace paddle
|
|
||||||
|
|
||||||
namespace ops = paddle::operators;
|
|
||||||
REGISTER_OPERATOR(lookup_remote_table, ops::LookupRemoteTableOp,
|
|
||||||
paddle::framework::EmptyGradOpMaker,
|
|
||||||
ops::LookupRemoteTableOpMaker);
|
|
||||||
|
|
||||||
REGISTER_OP_CPU_KERNEL(lookup_remote_table, ops::LookupRemoteTableKernel<float>,
|
|
||||||
ops::LookupRemoteTableKernel<double>);
|
|
File diff suppressed because it is too large
Load Diff
Loading…
Reference in new issue