gather_op.cc 3.3 KB
Newer Older
Z
zchen0211 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21
/* 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/gather_op.h"
#include "paddle/framework/ddim.h"

namespace paddle {
namespace operators {

class GatherOp : public framework::OperatorWithKernel {
Z
zchen0211 已提交
22 23 24
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

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

Z
zchen0211 已提交
33 34
    auto index_dims = ctx->GetInputDim("Index");
    PADDLE_ENFORCE(index_dims.size() == 1);
Q
Qiao Longfei 已提交
35
    int batch_size = ctx->GetInputDim("Index")[0];
Z
zchen0211 已提交
36
    PADDLE_ENFORCE_GE(batch_size, 0, "Batch size must be >0");
Q
Qiao Longfei 已提交
37
    framework::DDim output_dims(ctx->GetInputDim("X"));
Z
zchen0211 已提交
38
    output_dims[0] = batch_size;
Q
Qiao Longfei 已提交
39
    ctx->SetOutputDim("Out", output_dims);
Z
zchen0211 已提交
40
  }
Y
Yu Yang 已提交
41

42
 protected:
Y
Yu Yang 已提交
43
  framework::OpKernelType GetKernelType(
Y
Yu Yang 已提交
44
      const framework::ExecutionContext& ctx) const override {
Y
Yu Yang 已提交
45 46 47
    return framework::OpKernelType(
        framework::ToDataType(ctx.Input<Tensor>("X")->type()),
        ctx.device_context());
Y
Yu Yang 已提交
48
  }
Z
zchen0211 已提交
49 50 51 52 53 54
};

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

55
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
56
    ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
Z
zchen0211 已提交
57
  }
Y
Yu Yang 已提交
58

59
 protected:
Y
Yu Yang 已提交
60
  framework::OpKernelType GetKernelType(
Y
Yu Yang 已提交
61
      const framework::ExecutionContext& ctx) const override {
Y
Yu Yang 已提交
62 63 64
    return framework::OpKernelType(
        framework::ToDataType(ctx.Input<Tensor>("X")->type()),
        ctx.device_context());
Y
Yu Yang 已提交
65
  }
Z
zchen0211 已提交
66 67 68 69
};

class GatherOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Q
Qiao Longfei 已提交
70
  GatherOpMaker(framework::OpProto* proto, framework::OpAttrChecker* op_checker)
Z
zchen0211 已提交
71 72 73
      : OpProtoAndCheckerMaker(proto, op_checker) {
    AddInput("X", "The source input of gather op");
    AddInput("Index", "The index input of gather op");
K
kexinzhao 已提交
74
    AddOutput("Out", "The output of gather op");
Z
zchen0211 已提交
75
    AddComment(R"DOC(
K
kexinzhao 已提交
76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94
Gather Operator.

$Out = X[Index]$

Out is obtained by gathering entries of the outer-most dimension 
of X indexed by Index and concatenate them together.

Example:

X = [[1, 2],
     [3, 4],
     [5, 6]]

Index = [[1, 2]]

Then:

Out = [[3, 4],
       [5, 6]]
Z
zchen0211 已提交
95 96 97 98 99 100 101 102

)DOC");
  }
};
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
103 104
REGISTER_OP(gather, ops::GatherOp, ops::GatherOpMaker, gather_grad,
            ops::GatherGradOp);
Z
zchen0211 已提交
105 106
REGISTER_OP_CPU_KERNEL(gather, ops::GatherOpKernel<float>);
REGISTER_OP_CPU_KERNEL(gather_grad, ops::GatherGradientOpKernel<float>);