multiplex_op.cc 4.6 KB
Newer Older
Y
Yibing Liu 已提交
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/multiplex_op.h"

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;

class MultiplexOp : public framework::OperatorWithKernel {
 public:
24
  using framework::OperatorWithKernel::OperatorWithKernel;
Y
Yibing Liu 已提交
25

26
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
27 28
    PADDLE_ENFORCE(ctx->HasInput("Ids"), "Input(Ids) shouldn't be null.");
    PADDLE_ENFORCE(!ctx->Inputs("X").empty(),
29
                   "MultiInput(X) shouldn't be empty.");
Q
Qiao Longfei 已提交
30 31
    PADDLE_ENFORCE(ctx->HasOutput("Out"), "Output(Out) shouldn't be null.");
    auto ids_dim = ctx->GetInputDim("Ids");
32 33 34 35
    PADDLE_ENFORCE(
        ids_dim.size() == 2 && ids_dim[1] == 1,
        "The index tensor must be a vector with size batchSize x 1.");

Q
Qiao Longfei 已提交
36 37
    auto ins_dims = ctx->GetInputsDim("X");
    auto num_ins = ins_dims.size();
38 39 40
    PADDLE_ENFORCE(num_ins > 1,
                   "multiplex operator should have more than "
                   "one candidate input tensors.");
Y
Yibing Liu 已提交
41

Q
Qiao Longfei 已提交
42
    auto in_dim = ins_dims[0];
43 44
    PADDLE_ENFORCE(in_dim.size() >= 2,
                   "The rank of candidate tensors must be not less than 2.");
45
    for (size_t i = 1; i < num_ins; i++) {
Q
Qiao Longfei 已提交
46
      auto dim = ins_dims[i];
Y
Yibing Liu 已提交
47
      PADDLE_ENFORCE(in_dim == dim,
48
                     "All the candidate tensors must have the same size.");
Y
Yibing Liu 已提交
49
    }
Q
Qiao Longfei 已提交
50
    ctx->SetOutputDim("Out", in_dim);
Y
Yibing Liu 已提交
51
  }
Y
Yu Yang 已提交
52

53
 protected:
Y
Yu Yang 已提交
54
  framework::OpKernelType GetKernelType(
Y
Yu Yang 已提交
55
      const framework::ExecutionContext& ctx) const override {
Y
Yu Yang 已提交
56 57 58
    return framework::OpKernelType(
        framework::ToDataType(ctx.MultiInput<Tensor>("X")[0]->type()),
        ctx.device_context());
Y
Yu Yang 已提交
59
  }
Y
Yibing Liu 已提交
60 61 62 63
};

class MultiplexOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Q
Qiao Longfei 已提交
64 65
  MultiplexOpMaker(framework::OpProto* proto,
                   framework::OpAttrChecker* op_checker)
Y
Yibing Liu 已提交
66
      : OpProtoAndCheckerMaker(proto, op_checker) {
67 68 69
    AddInput("Ids", "The index tensor of multiplex operator.");
    AddInput("X", "The candidate tensors of multiplex operator.")
        .AsDuplicable();
Y
Yibing Liu 已提交
70
    AddOutput("Out", "The output tensor of multiplex operator.");
K
kexinzhao 已提交
71 72
    AddComment(R"DOC(
Multiplex Operator.
Y
Yibing Liu 已提交
73

74
Multiplex multiple tensors according to the index provided by the index tensor.
Y
Yibing Liu 已提交
75

76 77
Ids: the index tensor.
X[0 : N - 1]: the candidate tensors for output (N >= 2).
Y
Yibing Liu 已提交
78
For each index i from 0 to batchSize - 1, the output is the i-th row of the
79
the (Ids[i])-th tensor.
Y
Yibing Liu 已提交
80

81
For i-th row of the output tensor:
Y
Yibing Liu 已提交
82

K
kexinzhao 已提交
83
$$y[i] = x_{k}[i]$$
Y
Yibing Liu 已提交
84

K
kexinzhao 已提交
85
where `y` is the output tensor, `x_{k}` is the k-th input tensor,
86
and `k = Ids[i]`.
K
kexinzhao 已提交
87

Y
Yibing Liu 已提交
88 89 90 91 92 93
)DOC");
  }
};

class MultiplexGradOp : public framework::OperatorWithKernel {
 public:
94
  using framework::OperatorWithKernel::OperatorWithKernel;
Y
Yibing Liu 已提交
95

96
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
97 98
    PADDLE_ENFORCE(!ctx->Inputs("X").empty(), "Input(X) should not be null.");
    PADDLE_ENFORCE(!ctx->Outputs(framework::GradVarName("X")).empty(),
Y
Yibing Liu 已提交
99
                   "Output(X@Grad) should not be null.");
Q
Qiao Longfei 已提交
100 101 102 103
    PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Out")),
                   "Input(Out@GRAD) should not be null.");
    std::vector<framework::DDim> d_ins;
    auto ins = ctx->GetInputsDim("X");
104 105
    // No need to compute gradient for Input(Ids)
    for (size_t i = 0; i < ins.size(); i++) {
Q
Qiao Longfei 已提交
106
      d_ins.push_back(ins[i]);
Y
Yibing Liu 已提交
107
    }
Q
Qiao Longfei 已提交
108
    ctx->SetOutputsDim(framework::GradVarName("X"), d_ins);
Y
Yibing Liu 已提交
109
  }
Y
Yu Yang 已提交
110

111
 protected:
Y
Yu Yang 已提交
112
  framework::OpKernelType GetKernelType(
Y
Yu Yang 已提交
113
      const framework::ExecutionContext& ctx) const override {
Y
Yu Yang 已提交
114 115 116
    return framework::OpKernelType(
        framework::ToDataType(ctx.MultiInput<Tensor>("X")[0]->type()),
        ctx.device_context());
Y
Yu Yang 已提交
117
  }
Y
Yibing Liu 已提交
118 119 120 121 122 123
};

}  // namespace operators
}  // namespace paddle
namespace ops = paddle::operators;

124 125 126
REGISTER_OPERATOR(multiplex, ops::MultiplexOp, ops::MultiplexOpMaker,
                  paddle::framework::DefaultGradOpDescMaker<false>);
REGISTER_OPERATOR(multiplex_grad, ops::MultiplexGradOp);
Y
Yibing Liu 已提交
127 128
REGISTER_OP_CPU_KERNEL(
    multiplex, ops::MultiplexCPUKernel<paddle::platform::CPUPlace, float>);
129 130
REGISTER_OP_CPU_KERNEL(
    multiplex_grad,
Y
Yibing Liu 已提交
131
    ops::MultiplexGradCPUKernel<paddle::platform::CPUPlace, float>);