top_k_op.cc 2.9 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
武毅 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14

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. */

Y
Yi Wang 已提交
15
#include "paddle/fluid/operators/top_k_op.h"
武毅 已提交
16 17 18 19 20 21 22 23

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);
Q
Qiao Longfei 已提交
44 45
    ctx->ShareLoD("X", "Out");
    ctx->ShareLoD("X", "Indices");
武毅 已提交
46 47 48 49 50
  }
};

class TopkOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
51
  void Make() override {
52 53 54 55 56
    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
武毅 已提交
57

58 59 60 61 62
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");
武毅 已提交
63
    AddAttr<int>("k",
64 65
                 "(int, default 1) Number of top elements to look for along "
                 "the last dimension (along each row for matrices).")
武毅 已提交
66 67 68 69 70 71 72 73
        .SetDefault(1);
  }
};

}  // namespace operators
}  // namespace paddle

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