top_k_op.cc 2.8 KB
Newer Older
武毅 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
/* 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/top_k_op.h"

namespace paddle {
namespace operators {

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

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

Q
Qiao Longfei 已提交
32 33
    auto input_dims = ctx->GetInputDim("X");
    const int k = static_cast<int>(ctx->Attrs().Get<int>("k"));
武毅 已提交
34 35

    PADDLE_ENFORCE_GE(k, 1, "k must >= 1");
Q
Qiao Longfei 已提交
36 37
    PADDLE_ENFORCE_GE(input_dims.size(), 1, "input must have >= 1d shape");
    PADDLE_ENFORCE_GE(input_dims[input_dims.size() - 1], k,
武毅 已提交
38 39
                      "input must have >= k columns");

Q
Qiao Longfei 已提交
40
    framework::DDim dims = input_dims;
武毅 已提交
41
    dims[dims.size() - 1] = k;
Q
Qiao Longfei 已提交
42 43
    ctx->SetOutputDim("Out", dims);
    ctx->SetOutputDim("Indices", dims);
武毅 已提交
44 45 46 47 48
  }
};

class TopkOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
49
  TopkOpMaker(OpProto *proto, OpAttrChecker *op_checker)
武毅 已提交
50
      : OpProtoAndCheckerMaker(proto, op_checker) {
51 52 53 54 55
    AddInput("X", "(Tensor) The input of Topk op");
    AddOutput("Out", "(Tensor) The output tensor of Topk op");
    AddOutput("Indices", "(Tensor) The indices of Topk elements of input");
    AddComment(R"DOC(
Top K operator
武毅 已提交
56

57 58 59 60 61
If the input is a vector (1d tensor), this operator finds the k largest 
entries in the vector and outputs their values and indices as vectors. 
Thus values[j] is the j-th largest entry in input, and its index is indices[j].

For matrices, this operator computes the top k entries in each row. )DOC");
武毅 已提交
62
    AddAttr<int>("k",
63 64
                 "(int, default 1) Number of top elements to look for along "
                 "the last dimension (along each row for matrices).")
武毅 已提交
65 66 67 68 69 70 71 72
        .SetDefault(1);
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
F
fengjiayi 已提交
73 74
REGISTER_OPERATOR(top_k, ops::TopkOp, ops::TopkOpMaker,
                  paddle::framework::EmptyGradOpMaker);
武毅 已提交
75 76
REGISTER_OP_CPU_KERNEL(top_k,
                       ops::TopkKernel<paddle::platform::CPUPlace, float>);