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

L
Luo Tao 已提交
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

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
8

L
Luo Tao 已提交
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

Y
Yi Wang 已提交
15 16
#include "paddle/fluid/operators/lookup_table_op.h"
#include "paddle/fluid/framework/var_type_inference.h"
17 18 19 20 21 22 23 24

namespace paddle {
namespace operators {

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

25
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
26 27 28 29 30 31 32 33 34 35
    PADDLE_ENFORCE(ctx->HasInput("W"),
                   "Input(W) of LookupTableOp should not be null.");
    PADDLE_ENFORCE(ctx->HasInput("Ids"),
                   "Input(Ids) of LookupTableOp should not be null.");
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of LookupTableOp should not be null.");

    auto table_dims = ctx->GetInputDim("W");
    auto ids_dims = ctx->GetInputDim("Ids");

C
chengduoZH 已提交
36
    auto ids_var_type = ctx->GetInputsVarType("Ids").front();
C
chengduoZH 已提交
37 38 39 40 41
    // The type of Ids(Input) is SelectedRows or LoDTensor, when Ids's type
    // is LoDTensor, this tensor contains the ids to be looked up in W
    // and it must be a column vector with rank = 2 while the 2nd dimension
    // size must be 1, when Ids's type is SelectedRows, the rows of Ids
    // contains the ids to be looked up in W;
C
chengduoZH 已提交
42 43 44 45
    if (ids_var_type == framework::proto::VarType::LOD_TENSOR) {
      PADDLE_ENFORCE_EQ(ids_dims.size(), 2);
      PADDLE_ENFORCE_EQ(ids_dims[1], 1);
    }
46

Q
Qiao Longfei 已提交
47 48
    ctx->SetOutputDim("Out", {ids_dims[0], table_dims[1]});
    ctx->ShareLoD("Ids", /*->*/ "Out");
49
  }
Y
Yu Yang 已提交
50

51
 protected:
52
  framework::OpKernelType GetExpectedKernelType(
Y
Yu Yang 已提交
53
      const framework::ExecutionContext& ctx) const override {
Q
qiaolongfei 已提交
54 55
    auto data_type = framework::GetDataTypeOfVar(ctx.InputVar("W"));
    return framework::OpKernelType(data_type, ctx.device_context());
Y
Yu Yang 已提交
56
  }
57 58 59 60
};

class LookupTableOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
61
  LookupTableOpMaker(OpProto* proto, OpAttrChecker* op_checker)
62 63
      : OpProtoAndCheckerMaker(proto, op_checker) {
    AddInput("W",
C
chengduoZH 已提交
64
             "(Tensor) The input represents embedding tensors, "
K
kexinzhao 已提交
65
             "which is a learnable parameter.");
C
chengduoZH 已提交
66 67 68 69 70 71 72 73 74 75 76
    AddInput(
        "Ids",
        "(Tensor or SelectedRows) Ids's type can be Tensor or "
        "SelectedRows, when Ids's type is Tensor, this tensor contains "
        "the ids to be looked up in W and it must be a column vector with "
        "rank = 2 while the 2nd dimension size must be 1; when Ids's type is "
        "SelectedRows, the rows of Ids contains the ids to be looked up "
        "in W.");
    AddOutput("Out",
              "(Tensor or SelectedRows) The lookup results, which have the "
              "same type as W.");
C
chengduoZH 已提交
77
    AddAttr<bool>("is_sparse",
C
chengduoZH 已提交
78
                  "(boolean, default false) "
C
chengduoZH 已提交
79
                  "Sparse update.")
C
chengduoZH 已提交
80
        .SetDefault(false);
C
chengduoZH 已提交
81 82 83 84 85
    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.")
86
        .SetDefault(kNoPadding);
C
chengduoZH 已提交
87
    AddComment(R"DOC(
C
chengduoZH 已提交
88
Lookup Table Operator.
C
chengduoZH 已提交
89

C
chengduoZH 已提交
90
This operator is used to perform lookups on the parameter W,
C
chengduoZH 已提交
91 92 93 94 95 96 97 98
then concatenated into a dense or sparse tensor.

The type of Ids(Input) is SelectedRows, Tensor or LoDTensor, when Ids's
type is SelectedRows, the rows of Ids contains the ids to be looked up in W;
when Ids's type is Tensor, this tensor contains the ids to be looked up in W
and it must be a column vector with rank = 2 while the 2nd dimension size must be 1,
at this time, Ids can carry the LoD (Level of Details) information, or not, and
the output only shares the LoD information with input Ids.
99

C
chengduoZH 已提交
100 101 102 103 104

)DOC");
  }
};

105 106 107 108 109 110 111 112 113
class LookupTableOpGradDescMaker
    : public framework::DefaultGradOpDescMaker<true> {
  using ::paddle::framework::DefaultGradOpDescMaker<
      true>::DefaultGradOpDescMaker;

 protected:
  virtual std::string GradOpType() const { return "lookup_table_grad"; }
};

114 115 116 117
class LookupTableOpGrad : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

118
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
119 120
    auto table_dims = ctx->GetInputDim("W");
    ctx->SetOutputDim(framework::GradVarName("W"), table_dims);
121
  }
Y
Yu Yang 已提交
122

123
 protected:
124
  framework::OpKernelType GetExpectedKernelType(
Y
Yu Yang 已提交
125
      const framework::ExecutionContext& ctx) const override {
Q
qiaolongfei 已提交
126 127
    auto data_type = framework::GetDataTypeOfVar(ctx.InputVar("W"));
    return framework::OpKernelType(data_type, ctx.device_context());
Y
Yu Yang 已提交
128
  }
129 130
};

131 132
class LookupTableOpGradVarTypeInference : public framework::VarTypeInference {
 public:
Y
Yu Yang 已提交
133 134
  void operator()(const framework::OpDesc& op_desc,
                  framework::BlockDesc* block) const override {
135 136 137 138 139 140
    auto out_var_name = op_desc.Output(framework::GradVarName("W")).front();
    auto attr = op_desc.GetAttr("is_sparse");
    bool is_sparse = boost::get<bool>(attr);
    if (is_sparse) {
      VLOG(3) << "lookup_table_grad op " << framework::GradVarName("W")
              << " is set to SelectedRows";
141
      block->Var(out_var_name)
142
          ->SetType(framework::proto::VarType::SELECTED_ROWS);
143 144 145
    } else {
      VLOG(3) << "lookup_table_grad op " << framework::GradVarName("W")
              << " is set to LoDTensor";
146
      block->Var(out_var_name)->SetType(framework::proto::VarType::LOD_TENSOR);
147 148 149 150
    }
  }
};

151 152 153 154
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
155 156 157 158 159 160 161 162 163
REGISTER_OPERATOR(lookup_table, ops::LookupTableOp,
                  ops::LookupTableOpGradDescMaker, ops::LookupTableOpMaker);
REGISTER_OPERATOR(lookup_table_grad, ops::LookupTableOpGrad,
                  ops::LookupTableOpGradVarTypeInference);

REGISTER_OP_CPU_KERNEL(lookup_table, ops::LookupTableKernel<float>,
                       ops::LookupTableKernel<double>);
REGISTER_OP_CPU_KERNEL(lookup_table_grad, ops::LookupTableGradKernel<float>,
                       ops::LookupTableGradKernel<double>);