multiplex_op.cc 7.1 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
Y
Yibing Liu 已提交
2

L
Luo Tao 已提交
3 4 5
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
Y
Yibing Liu 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
Y
Yibing Liu 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
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
Yibing Liu 已提交
14

Y
Yi Wang 已提交
15
#include "paddle/fluid/operators/multiplex_op.h"
S
sneaxiy 已提交
16
#include <memory>
S
sneaxiy 已提交
17
#include <vector>
Y
Yibing Liu 已提交
18 19 20 21 22 23 24 25

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;

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

28
  void InferShape(framework::InferShapeContext* ctx) const override {
29 30 31 32 33
    OP_INOUT_CHECK(ctx->HasInput("Ids"), "Input", "Ids", "Multiplex");
    PADDLE_ENFORCE_NE(
        ctx->Inputs("X").empty(), true,
        platform::errors::InvalidArgument("MultiInput(X) shouldn't be empty."));
    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "Multiplex");
Q
Qiao Longfei 已提交
34
    auto ids_dim = ctx->GetInputDim("Ids");
35 36 37 38 39 40 41 42
    PADDLE_ENFORCE_EQ(
        ids_dim.size(), 2,
        platform::errors::PreconditionNotMet(
            "The index tensor must be a vector with 2 dimensions"));
    PADDLE_ENFORCE_EQ(
        ids_dim[1], 1,
        platform::errors::PreconditionNotMet(
            "The index tensor must be a vector with batchSize x 1."));
43

Q
Qiao Longfei 已提交
44 45
    auto ins_dims = ctx->GetInputsDim("X");
    auto num_ins = ins_dims.size();
46 47 48 49
    PADDLE_ENFORCE_GT(num_ins, 1,
                      platform::errors::InvalidArgument(
                          "multiplex operator should have more than "
                          "one candidate input tensors."));
Y
Yibing Liu 已提交
50

Q
Qiao Longfei 已提交
51
    auto in_dim = ins_dims[0];
52 53 54 55
    PADDLE_ENFORCE_GE(
        in_dim.size(), 2,
        platform::errors::InvalidArgument(
            "The rank of candidate tensors must be not less than 2."));
56
    for (size_t i = 1; i < num_ins; i++) {
Q
Qiao Longfei 已提交
57
      auto dim = ins_dims[i];
58 59 60 61
      PADDLE_ENFORCE_EQ(
          in_dim, dim,
          platform::errors::PreconditionNotMet(
              "All the candidate tensors must have the same size."));
Y
Yibing Liu 已提交
62
    }
Q
Qiao Longfei 已提交
63
    ctx->SetOutputDim("Out", in_dim);
Y
Yibing Liu 已提交
64
  }
Y
Yu Yang 已提交
65

66
 protected:
67
  framework::OpKernelType GetExpectedKernelType(
Y
Yu Yang 已提交
68
      const framework::ExecutionContext& ctx) const override {
69 70 71
    return framework::OpKernelType(
        OperatorWithKernel::IndicateVarDataType(ctx, "X"),
        ctx.device_context());
Y
Yu Yang 已提交
72
  }
Y
Yibing Liu 已提交
73 74 75 76
};

class MultiplexOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
77
  void Make() override {
Y
yuyang18 已提交
78 79 80 81 82 83
    AddInput("Ids",
             "Tensor<int32>, index variable which is a 2-D tensor with shape "
             "[M, 1] where M is the batch size.");
    AddInput("X",
             "A list of variables to gather from. All variables have the same "
             "shape and the rank is at least 2.")
84
        .AsDuplicable();
Y
Yibing Liu 已提交
85
    AddOutput("Out", "The output tensor of multiplex operator.");
K
kexinzhao 已提交
86
    AddComment(R"DOC(
Y
yuyang18 已提交
87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107
Referring to the given index variable, this layer selects rows from the
input variables to construct a multiplex variable. Assuming that there are
:math:`m` input variables and :math:`I_i` represents the i-th input
variable and :math:`i` is in [0, :math:`m`). All input variables are
tensors with same shape [:math:`d_0`, :math:`d_1`, ..., :math:`d_R`].
Please note that rank of the input tensor should be at least 2. Each input
variable will be treated as a 2-D matrix with shape [:math:`M`, :math:`N`]
where :math:`M` for :math:`d_0` and :math:`N` for :math:`d_1` * :math:`d_2`
* ... * :math:`d_R`. Let :math:`I_i[j]` be the j-th row of the i-th input
variable. The given index variable should be a 2-D tensor with shape
[:math:`M`, 1]. Let `ID[i]` be the i-th index value of the index variable.
Then the output variable will be a tensor with shape [:math:`d_0`,
:math:`d_1`, ..., :math:`d_R`]. If we treat the output tensor as a 2-D
matrix with shape [:math:`M`, :math:`N`] and let :math:`O[i]` be the i-th
row of the matrix, then `O[i]` is equal to :math:`I_{ID[i]}[i]`.

* Ids: the index tensor.

* X[0 : N - 1]: the candidate tensors for output (N >= 2).

* For each index i from 0 to batchSize - 1, the output is the i-th row of the
108
the (Ids[i])-th tensor.
Y
Yibing Liu 已提交
109

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

Y
yuyang18 已提交
112 113 114
$$
y[i] = x_{k}[i]
$$
Y
Yibing Liu 已提交
115

Y
yuyang18 已提交
116 117
where $y$ is the output tensor, $x_{k}$ is the k-th input tensor,
and $k = Ids[i]$.
K
kexinzhao 已提交
118

Y
Yibing Liu 已提交
119 120 121 122 123 124
)DOC");
  }
};

class MultiplexGradOp : public framework::OperatorWithKernel {
 public:
125
  using framework::OperatorWithKernel::OperatorWithKernel;
Y
Yibing Liu 已提交
126

127
  void InferShape(framework::InferShapeContext* ctx) const override {
H
hong 已提交
128
    auto dxs = ctx->Outputs(framework::GradVarName("X"));
129 130 131 132 133
    PADDLE_ENFORCE_NE(dxs.empty(), true,
                      platform::errors::InvalidArgument(
                          "Output(X@Grad) should not be null."));
    OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Out")), "Input",
                   framework::GradVarName("Out"), "MultiplexGrad");
S
sneaxiy 已提交
134 135 136
    auto dout_dim = ctx->GetInputDim(framework::GradVarName("Out"));
    ctx->SetOutputsDim(framework::GradVarName("X"),
                       std::vector<framework::DDim>(dxs.size(), dout_dim));
Y
Yibing Liu 已提交
137
  }
Y
Yu Yang 已提交
138

139
 protected:
140
  framework::OpKernelType GetExpectedKernelType(
Y
Yu Yang 已提交
141
      const framework::ExecutionContext& ctx) const override {
142 143 144
    return framework::OpKernelType(OperatorWithKernel::IndicateVarDataType(
                                       ctx, framework::GradVarName("Out")),
                                   ctx.device_context());
S
sneaxiy 已提交
145 146 147
  }
};

H
hong 已提交
148 149
template <typename T>
class MultiplexGradMaker : public framework::SingleGradOpMaker<T> {
S
sneaxiy 已提交
150
 public:
H
hong 已提交
151
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
S
sneaxiy 已提交
152 153

 protected:
154
  void Apply(GradOpPtr<T> op) const override {
S
sneaxiy 已提交
155
    op->SetType("multiplex_grad");
H
hong 已提交
156 157 158 159
    op->SetInput("Ids", this->Input("Ids"));
    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
    op->SetOutput(framework::GradVarName("X"), this->InputGrad("X", false));
    op->SetAttrMap(this->Attrs());
Y
Yu Yang 已提交
160
  }
Y
Yibing Liu 已提交
161 162 163 164
};

}  // namespace operators
}  // namespace paddle
S
sneaxiy 已提交
165

Y
Yibing Liu 已提交
166 167
namespace ops = paddle::operators;

168
REGISTER_OPERATOR(multiplex, ops::MultiplexOp, ops::MultiplexOpMaker,
H
hong 已提交
169 170
                  ops::MultiplexGradMaker<paddle::framework::OpDesc>,
                  ops::MultiplexGradMaker<paddle::imperative::OpBase>);
171
REGISTER_OPERATOR(multiplex_grad, ops::MultiplexGradOp);
Y
Yibing Liu 已提交
172
REGISTER_OP_CPU_KERNEL(
Q
QI JUN 已提交
173
    multiplex,
174 175 176 177
    ops::MultiplexCPUKernel<paddle::platform::CPUDeviceContext, float>,
    ops::MultiplexCPUKernel<paddle::platform::CPUDeviceContext, double>,
    ops::MultiplexCPUKernel<paddle::platform::CPUDeviceContext, int>,
    ops::MultiplexCPUKernel<paddle::platform::CPUDeviceContext, int64_t>);
178 179
REGISTER_OP_CPU_KERNEL(
    multiplex_grad,
180 181 182 183
    ops::MultiplexGradCPUKernel<paddle::platform::CPUDeviceContext, float>,
    ops::MultiplexGradCPUKernel<paddle::platform::CPUDeviceContext, double>,
    ops::MultiplexGradCPUKernel<paddle::platform::CPUDeviceContext, int>,
    ops::MultiplexGradCPUKernel<paddle::platform::CPUDeviceContext, int64_t>);