elementwise_op.h 13.0 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
G
gongweibao 已提交
2

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
G
gongweibao 已提交
6

7
    http://www.apache.org/licenses/LICENSE-2.0
G
gongweibao 已提交
8

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. */
G
gongweibao 已提交
14 15

#pragma once
C
chengduo 已提交
16

L
liuwei1031 已提交
17
#include <memory>
18
#include <string>
L
liuwei1031 已提交
19
#include <unordered_map>
20
#include "paddle/fluid/framework/data_layout.h"
Y
Yi Wang 已提交
21 22
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/operator.h"
C
chengduo 已提交
23

24 25 26
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif
G
gongweibao 已提交
27 28 29 30 31 32 33 34 35

namespace paddle {
namespace operators {

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

  using Tensor = framework::Tensor;
C
chengduo 已提交
36 37

  void InferShape(framework::InferShapeContext *ctx) const override {
Q
Qiao Longfei 已提交
38
    PADDLE_ENFORCE(ctx->HasInput("X"),
C
caoying03 已提交
39
                   "Input(X) of elementwise op should not be null.");
Q
Qiao Longfei 已提交
40
    PADDLE_ENFORCE(ctx->HasInput("Y"),
C
caoying03 已提交
41
                   "Input(Y) of elementwise op should not be null.");
Q
Qiao Longfei 已提交
42 43 44
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of elementwise op should not be null.");

C
chengduo 已提交
45 46 47
    PADDLE_ENFORCE(
        ctx->GetInputsVarType("Y").front() ==
            framework::proto::VarType::LOD_TENSOR,
C
chengduo 已提交
48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66
        "The input var's type should be LoDTensor, but the received is %s [%s]",
        ctx->GetInputsVarType("Y").front(), ctx->Inputs("Y").front());

    if (ctx->GetInputsVarType("X").front() ==
        framework::proto::VarType::LOD_TENSOR) {
      auto x_dim = ctx->GetInputDim("X");
      auto y_dim = ctx->GetInputDim("Y");
      PADDLE_ENFORCE_GE(x_dim.size(), y_dim.size(),
                        "Rank of first input must >= rank of second input.");
    } else if (ctx->GetInputsVarType("X").front() ==
               framework::proto::VarType::SELECTED_ROWS) {
      PADDLE_ENFORCE((ctx->GetInputDim("Y").size() == 1u) &&
                         (ctx->GetInputDim("Y")[0] == 1),
                     "For elementwise_op, if X is Sparse, "
                     "Y must be scalar.");
    } else {
      PADDLE_THROW("X's type[%s] is not supported by elementwise_op.",
                   ctx->GetInputsVarType("X").front());
    }
67 68

    ctx->ShareDim("X", /*->*/ "Out");
Q
Qiao Longfei 已提交
69
    ctx->ShareLoD("X", /*->*/ "Out");
G
gongweibao 已提交
70
  }
71 72

  framework::OpKernelType GetExpectedKernelType(
C
chengduo 已提交
73 74
      const framework::ExecutionContext &ctx) const override {
    auto input_data_type = framework::GetDataTypeOfVar(ctx.InputVar("X"));
75 76 77 78 79 80 81 82 83 84

#ifdef PADDLE_WITH_MKLDNN
    if (platform::CanMKLDNNBeUsed(ctx)) {
      return framework::OpKernelType(input_data_type, ctx.GetPlace(),
                                     framework::DataLayout::kMKLDNN,
                                     framework::LibraryType::kMKLDNN);
    }
#endif
    return framework::OpKernelType(input_data_type, ctx.GetPlace());
  }
G
gongweibao 已提交
85 86
};

C
chengduo 已提交
87 88 89 90 91 92
class ElementwiseOpInferVarType
    : public framework::PassInDtypeAndVarTypeToOutput {
 protected:
  std::unordered_map<std::string, std::string> GetInputOutputWithSameType()
      const override {
    return std::unordered_map<std::string, std::string>{{"X", /*->*/ "Out"}};
93 94 95
  }
};

G
gongweibao 已提交
96 97
class ElementwiseOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
98
  void Make() final {
C
caoying03 已提交
99 100
    AddInput("X", "(Tensor), The first input tensor of elementwise op.");
    AddInput("Y", "(Tensor), The second input tensor of elementwise op.");
101
    AddOutput("Out", "The output of elementwise op.");
G
gongweibao 已提交
102
    AddAttr<int>("axis",
C
caoying03 已提交
103 104
                 "(int, default -1). The start dimension index "
                 "for broadcasting Y onto X.")
G
gongweibao 已提交
105 106
        .SetDefault(-1)
        .EqualGreaterThan(-1);
107 108
    AddAttr<bool>("use_mkldnn", "(bool, default false). Used by MKLDNN.")
        .SetDefault(false);
109
    AddAttr<std::string>(
110 111 112 113 114 115
        "x_data_format",
        "(string, default NCHW) Only used in mkldnn"
        "An optional string from: \"NHWC\", \"NCHW\", \"NCHW16C\", \"NCHW8C\". "
        "Defaults to \"\". Specify the data format of the output data, "
        "the input will be transformed automatically. ")
        .SetDefault("");
116
    AddAttr<std::string>(
117 118 119 120 121 122
        "y_data_format",
        "(string, default \"\") Only used in mkldnn"
        "An optional string from: \"NHWC\", \"NCHW\", \"NCHW16C\", \"NCHW8C\". "
        "Defaults to \"\". Specify the data format of the output data, "
        "the input will be transformed automatically. ")
        .SetDefault("");
Y
Yu Yang 已提交
123
    AddComment(string::Sprintf(R"DOC(
T
Tao Luo 已提交
124
Elementwise %s Operator
K
kexinzhao 已提交
125 126 127

The equation is:

Y
Yu Yang 已提交
128
$$%s$$
K
kexinzhao 已提交
129

L
Luo Tao 已提交
130 131
- $X$: a tensor of any dimension. 
- $Y$: a tensor whose dimensions must be less than or equal to the dimensions of $X$.
K
kexinzhao 已提交
132 133

There are two cases for this operator:
134

L
Luo Tao 已提交
135 136
1. The shape of $Y$ is the same with $X$.
2. The shape of $Y$ is a continuous subsequence of $X$.
K
kexinzhao 已提交
137 138

For case 2:
139

L
Luo Tao 已提交
140 141 142 143 144
1. Broadcast $Y$ to match the shape of $X$, where $axis$ is the start dimension index 
   for broadcasting $Y$ onto $X$. 
2. If $axis$ is -1 (default), $axis = rank(X) - rank(Y)$.
3. The trailing dimensions of size 1 for $Y$ will be ignored for the consideration of 
   subsequence, such as shape(Y) = (2, 1) => (2).
K
kexinzhao 已提交
145

L
Luo Tao 已提交
146
For example:
147

148
  .. code-block:: python
G
gongweibao 已提交
149

150 151
    shape(X) = (2, 3, 4, 5), shape(Y) = (,)
    shape(X) = (2, 3, 4, 5), shape(Y) = (5,)
L
Luo Tao 已提交
152
    shape(X) = (2, 3, 4, 5), shape(Y) = (4, 5), with axis=-1(default) or axis=2
153 154
    shape(X) = (2, 3, 4, 5), shape(Y) = (3, 4), with axis=1
    shape(X) = (2, 3, 4, 5), shape(Y) = (2), with axis=0
155
    shape(X) = (2, 3, 4, 5), shape(Y) = (2, 1), with axis=0
156

L
Luo Tao 已提交
157 158
The inputs $X$ and $Y$ can carry the different LoD information. 
But the output only shares the LoD information with the input $X$.
K
kexinzhao 已提交
159

Y
Yu Yang 已提交
160 161
)DOC",
                               GetName(), GetEquation()));
G
gongweibao 已提交
162 163 164
  }

 protected:
Y
Yu Yang 已提交
165
  virtual std::string GetName() const = 0;
C
chengduo 已提交
166

Y
Yu Yang 已提交
167
  virtual std::string GetEquation() const = 0;
G
gongweibao 已提交
168 169 170 171 172 173 174
};

class ElementwiseOpGrad : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;
  using Tensor = framework::Tensor;

C
chengduo 已提交
175
  void InferShape(framework::InferShapeContext *ctx) const override {
Q
Qiao Longfei 已提交
176 177 178 179 180 181 182
    PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) should not be null");
    PADDLE_ENFORCE(ctx->HasInput("Y"), "Input(Y) should not be null");
    PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Out")),
                   "Input(Out@GRAD) should not be null");

    auto x_dims = ctx->GetInputDim("X");
    auto y_dims = ctx->GetInputDim("Y");
G
gongweibao 已提交
183 184

    PADDLE_ENFORCE_GE(x_dims.size(), y_dims.size(),
185
                      "Rank of first input must >= rank of second input.");
G
gongweibao 已提交
186

Q
Qiao Longfei 已提交
187 188 189
    auto x_grad_name = framework::GradVarName("X");
    auto y_grad_name = framework::GradVarName("Y");
    if (ctx->HasOutput(x_grad_name)) {
190 191
      ctx->ShareDim("X", /*->*/ x_grad_name);
      ctx->ShareLoD("X", /*->*/ x_grad_name);
G
gongweibao 已提交
192
    }
Q
Qiao Longfei 已提交
193
    if (ctx->HasOutput(y_grad_name)) {
194 195
      ctx->ShareDim("Y", /*->*/ y_grad_name);
      ctx->ShareLoD("Y", /*->*/ y_grad_name);
G
gongweibao 已提交
196 197
    }
  }
198 199

  framework::OpKernelType GetExpectedKernelType(
C
chengduo 已提交
200
      const framework::ExecutionContext &ctx) const override {
Y
Yu Yang 已提交
201 202
    auto input_data_type =
        ctx.Input<Tensor>(framework::GradVarName("Out"))->type();
203 204 205 206 207 208 209 210 211 212

#ifdef PADDLE_WITH_MKLDNN
    if (platform::CanMKLDNNBeUsed(ctx)) {
      return framework::OpKernelType(input_data_type, ctx.GetPlace(),
                                     framework::DataLayout::kMKLDNN,
                                     framework::LibraryType::kMKLDNN);
    }
#endif
    return framework::OpKernelType(input_data_type, ctx.GetPlace());
  }
G
gongweibao 已提交
213
};
214 215 216 217 218 219 220 221

// For Add, Sub op, the X, Out is not needed.
class ElementwiseOpExplicitGrad : public ElementwiseOpGrad {
 public:
  using operators::ElementwiseOpGrad::ElementwiseOpGrad;
  using operators::ElementwiseOpGrad::GetExpectedKernelType;
  using Tensor = framework::Tensor;

C
chengduo 已提交
222
  void InferShape(framework::InferShapeContext *ctx) const override {
223 224 225 226 227
    PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Out")),
                   "Input(Out@GRAD) should not be null");

    auto x_grad_name = framework::GradVarName("X");
    if (ctx->HasOutput(x_grad_name)) {
228 229
      ctx->ShareDim(framework::GradVarName("Out"), /*->*/ x_grad_name);
      ctx->ShareLoD(framework::GradVarName("Out"), /*->*/ x_grad_name);
230 231 232 233
    }
    auto y_grad_name = framework::GradVarName("Y");
    if (ctx->HasOutput(y_grad_name)) {
      PADDLE_ENFORCE(ctx->HasInput("Y"), "Input(Y) should not be null");
234 235 236

      ctx->ShareDim("Y", /*->*/ y_grad_name);
      ctx->ShareLoD("Y", /*->*/ y_grad_name);
237 238 239 240
    }
  }
};

241 242 243
template <typename T>
class ElemwiseGradKernel : public framework::OpKernel<T> {
 public:
C
chengduo 已提交
244 245
  void Compute(const framework::ExecutionContext &context) const override {
    auto *dx =
246 247
        context.Output<framework::LoDTensor>(framework::GradVarName("X"));
    if (dx != nullptr) {
C
chengduo 已提交
248
      auto &dout =
249 250 251 252 253 254
          *context.Input<framework::LoDTensor>(framework::GradVarName("Out"));
      dx->set_lod(dout.lod());
    }
  }
};

L
liuwei1031 已提交
255
class ElementwiseOpInplace : public framework::InplaceOpInference {
D
dzhwinter 已提交
256
 public:
L
liuwei1031 已提交
257 258
  std::unordered_map<std::string, std::string> operator()(
      const framework::OpDesc &op_desc) const override {
D
dzhwinter 已提交
259 260 261 262 263 264
    return std::unordered_map<std::string, std::string>{
        {"X", "Out"},
    };
  }
};

L
liuwei1031 已提交
265
class ElementwiseGradOpInplace : public framework::InplaceOpInference {
D
dzhwinter 已提交
266
 public:
L
liuwei1031 已提交
267 268 269 270 271
  std::unordered_map<std::string, std::string> operator()(
      const framework::OpDesc &op_desc) const override {
    return std::unordered_map<std::string, std::string>{
        {framework::GradVarName("Out"), framework::GradVarName("X")},
    };
D
dzhwinter 已提交
272 273 274
  }
};

G
gongweibao 已提交
275 276
}  // namespace operators
}  // namespace paddle
Y
Yu Yang 已提交
277

278 279 280 281 282 283 284 285 286 287 288
/*
*/

#define REGISTER_ELEMWISE_GRAD_MAKER(kernel_type, op_name)                   \
  class kernel_type##GradMaker                                               \
      : public paddle::framework::SingleGradOpDescMaker {                    \
   public:                                                                   \
    using ::paddle::framework::SingleGradOpDescMaker::SingleGradOpDescMaker; \
                                                                             \
   protected:                                                                \
    std::unique_ptr<paddle::framework::OpDesc> Apply() const override {      \
C
chengduo 已提交
289
      auto *op = new paddle::framework::OpDesc();                            \
290 291 292 293 294 295 296 297 298 299 300
      op->SetType(#kernel_type "_grad");                                     \
      op->SetInput("Y", Input("Y"));                                         \
      op->SetInput(::paddle::framework::GradVarName("Out"),                  \
                   OutputGrad("Out"));                                       \
      op->SetAttrMap(Attrs());                                               \
      op->SetOutput(::paddle::framework::GradVarName("X"), InputGrad("X"));  \
      op->SetOutput(::paddle::framework::GradVarName("Y"), InputGrad("Y"));  \
      return std::unique_ptr<::paddle::framework::OpDesc>(op);               \
    }                                                                        \
  }

Y
Yu Yang 已提交
301 302 303 304 305 306 307 308 309
#define REGISTER_ELEMWISE_OP(op_type, op_name, equation)                \
  class __ElemwiseOp##op_type##Maker__                                  \
      : public ::paddle::operators::ElementwiseOpMaker {                \
   protected:                                                           \
    virtual std::string GetName() const { return op_name; }             \
    virtual std::string GetEquation() const { return equation; }        \
  };                                                                    \
  REGISTER_OPERATOR(op_type, ::paddle::operators::ElementwiseOp,        \
                    __ElemwiseOp##op_type##Maker__,                     \
310
                    ::paddle::operators::ElementwiseOpInferVarType,     \
Y
Yu Yang 已提交
311 312
                    ::paddle::framework::DefaultGradOpDescMaker<true>); \
  REGISTER_OPERATOR(op_type##_grad, ::paddle::operators::ElementwiseOpGrad)
313 314 315 316 317 318 319 320 321 322 323

#define REGISTER_ELEMWISE_EXPLICIT_OP(op_type, op_name, equation, ...) \
  class __ElemwiseOp##op_type##Maker__                                 \
      : public ::paddle::operators::ElementwiseOpMaker {               \
   protected:                                                          \
    virtual std::string GetName() const { return op_name; }            \
    virtual std::string GetEquation() const { return equation; }       \
  };                                                                   \
  REGISTER_OPERATOR(op_type, ::paddle::operators::ElementwiseOp,       \
                    __ElemwiseOp##op_type##Maker__,                    \
                    ::paddle::operators::ElementwiseOpInferVarType,    \
D
dzhwinter 已提交
324 325
                    op_type##GradMaker,                                \
                    ::paddle::operators::ElementwiseOpInplace);        \
326
  REGISTER_OPERATOR(op_type##_grad,                                    \
D
dzhwinter 已提交
327 328
                    ::paddle::operators::ElementwiseOpExplicitGrad,    \
                    ::paddle::operators::ElementwiseGradOpInplace)