multiplex_op.cc 3.9 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
/* 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;
21
using LoDTensor = framework::LoDTensor;
Y
Yibing Liu 已提交
22 23 24

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

 protected:
  void InferShape(const framework::InferShapeContext &ctx) const override {
29 30 31 32
    PADDLE_ENFORCE(!ctx.MultiInputVar("X").empty(),
                   "Input(X) should not be null");
    PADDLE_ENFORCE_NOT_NULL(ctx.OutputVar("Out"),
                            "Output(Out) shouldn't be null.");
Y
Yibing Liu 已提交
33
    auto ins = ctx.MultiInput<Tensor>("X");
34
    auto *out = ctx.Output<LoDTensor>("Out");
Y
Yibing Liu 已提交
35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56
    auto num_ins = ins.size();
    PADDLE_ENFORCE(num_ins > 2,
                   "multiplex operator should have more than 2 inputs.");
    PADDLE_ENFORCE_EQ(ins[0]->dims().size(), 1,
                      "The first input must be a index vector.");
    auto in_dim = ins[1]->dims();

    for (size_t i = 2; i < num_ins; i++) {
      auto dim = ins[i]->dims();
      PADDLE_ENFORCE(
          in_dim == dim,
          "All the input tensors except the first one must have the same size");
    }
    out->Resize(in_dim);
  }
};

class MultiplexOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
  MultiplexOpMaker(framework::OpProto *proto,
                   framework::OpAttrChecker *op_checker)
      : OpProtoAndCheckerMaker(proto, op_checker) {
57
    AddInput("X", "The input tensors of multiplex operator.").AsDuplicable();
Y
Yibing Liu 已提交
58 59 60 61 62 63
    AddOutput("Out", "The output tensor of multiplex operator.");
    AddComment(R"DOC(Multiplex operator

Multiplex multiple tensors according to the index provided by the first
input tensor.

64 65
ins[0]: the index tensor.
ins[1:N]: the candidate output tensors.
Y
Yibing Liu 已提交
66 67 68
For each index i from 0 to batchSize - 1, the output is the i-th row of the
the (index[i] + 1)-th tensor.

69
For i-th row of the output tensor:
Y
Yibing Liu 已提交
70 71 72

y[i][j] = x_{k}[i][j], j = 0,1, ... , (x_{1}.width - 1)

73
where y is the output tensor. `x_{k}` is the k-th input tensor
Y
Yibing Liu 已提交
74 75 76 77 78 79 80 81
and `k = x{0}[i] + 1`.

)DOC");
  }
};

class MultiplexGradOp : public framework::OperatorWithKernel {
 public:
82
  using framework::OperatorWithKernel::OperatorWithKernel;
Y
Yibing Liu 已提交
83 84 85

 protected:
  void InferShape(const framework::InferShapeContext &ctx) const override {
86 87 88 89
    PADDLE_ENFORCE(!ctx.MultiInputVar("X").empty(),
                   "Input(X) should not be null");
    PADDLE_ENFORCE(!ctx.MultiOutputVar(framework::GradVarName("X")).empty(),
                   "Output(X@Grad) should not be null");
Y
Yibing Liu 已提交
90 91
    PADDLE_ENFORCE_NOT_NULL(ctx.InputVar(framework::GradVarName("Out")),
                            "Input(Out@GRAD) shouldn't be null.");
92
    auto d_ins = ctx.MultiOutput<LoDTensor>(framework::GradVarName("X"));
Y
Yibing Liu 已提交
93
    auto ins = ctx.MultiInput<Tensor>("X");
94
    // don't compute gradient for index (ins[0])
95 96 97 98
    for (size_t i = 1; i < ins.size(); i++) {
      if (d_ins[i]) {
        d_ins[i]->Resize(ins[i]->dims());
      }
Y
Yibing Liu 已提交
99 100 101 102 103 104 105 106 107 108
    }
  }
};

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

REGISTER_OP(multiplex, ops::MultiplexOp, ops::MultiplexOpMaker, multiplex_grad,
            ops::MultiplexGradOp);
Y
Yibing Liu 已提交
109 110
REGISTER_OP_CPU_KERNEL(
    multiplex, ops::MultiplexCPUKernel<paddle::platform::CPUPlace, float>);
111 112
REGISTER_OP_CPU_KERNEL(
    multiplex_grad,
Y
Yibing Liu 已提交
113
    ops::MultiplexGradCPUKernel<paddle::platform::CPUPlace, float>);