elementwise_op.h 12.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

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

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

namespace paddle {
namespace operators {

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

  using Tensor = framework::Tensor;
C
chengduo 已提交
34 35

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

C
chengduo 已提交
43 44 45 46 47 48 49 50 51 52 53
    PADDLE_ENFORCE(
        ctx->GetInputsVarType("X").front() ==
            framework::proto::VarType::LOD_TENSOR,
        "The input var's type should be LoDTensor, but the received is %s",
        ctx->Inputs("X").front(), ctx->GetInputsVarType("X").front());
    PADDLE_ENFORCE(
        ctx->GetInputsVarType("Y").front() ==
            framework::proto::VarType::LOD_TENSOR,
        "The input var's type should be LoDTensor, but the received is %s",
        ctx->Inputs("Y").front(), ctx->GetInputsVarType("Y").front());

Q
Qiao Longfei 已提交
54 55
    auto x_dim = ctx->GetInputDim("X");
    auto y_dim = ctx->GetInputDim("Y");
G
gongweibao 已提交
56
    PADDLE_ENFORCE_GE(x_dim.size(), y_dim.size(),
57
                      "Rank of first input must >= rank of second input.");
58 59

    ctx->ShareDim("X", /*->*/ "Out");
Q
Qiao Longfei 已提交
60
    ctx->ShareLoD("X", /*->*/ "Out");
G
gongweibao 已提交
61
  }
62 63

  framework::OpKernelType GetExpectedKernelType(
C
chengduo 已提交
64 65
      const framework::ExecutionContext &ctx) const override {
    auto input_data_type = framework::GetDataTypeOfVar(ctx.InputVar("X"));
66 67 68 69 70 71 72 73 74 75

#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 已提交
76 77
};

C
chengduo 已提交
78 79 80 81 82 83
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"}};
84 85 86
  }
};

G
gongweibao 已提交
87 88
class ElementwiseOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
89
  void Make() final {
C
caoying03 已提交
90 91
    AddInput("X", "(Tensor), The first input tensor of elementwise op.");
    AddInput("Y", "(Tensor), The second input tensor of elementwise op.");
92
    AddOutput("Out", "The output of elementwise op.");
G
gongweibao 已提交
93
    AddAttr<int>("axis",
C
caoying03 已提交
94 95
                 "(int, default -1). The start dimension index "
                 "for broadcasting Y onto X.")
G
gongweibao 已提交
96 97
        .SetDefault(-1)
        .EqualGreaterThan(-1);
98 99
    AddAttr<bool>("use_mkldnn", "(bool, default false). Used by MKLDNN.")
        .SetDefault(false);
100 101 102 103 104 105 106 107 108 109 110 111 112 113
    AddAttr<std::string>(
      "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("");
    AddAttr<std::string>(
      "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 已提交
114
    AddComment(string::Sprintf(R"DOC(
T
Tao Luo 已提交
115
Elementwise %s Operator
K
kexinzhao 已提交
116 117 118

The equation is:

Y
Yu Yang 已提交
119
$$%s$$
K
kexinzhao 已提交
120

L
Luo Tao 已提交
121 122
- $X$: a tensor of any dimension. 
- $Y$: a tensor whose dimensions must be less than or equal to the dimensions of $X$.
K
kexinzhao 已提交
123 124

There are two cases for this operator:
125

L
Luo Tao 已提交
126 127
1. The shape of $Y$ is the same with $X$.
2. The shape of $Y$ is a continuous subsequence of $X$.
K
kexinzhao 已提交
128 129

For case 2:
130

L
Luo Tao 已提交
131 132 133 134 135
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 已提交
136

L
Luo Tao 已提交
137
For example:
138

139
  .. code-block:: python
G
gongweibao 已提交
140

141 142
    shape(X) = (2, 3, 4, 5), shape(Y) = (,)
    shape(X) = (2, 3, 4, 5), shape(Y) = (5,)
L
Luo Tao 已提交
143
    shape(X) = (2, 3, 4, 5), shape(Y) = (4, 5), with axis=-1(default) or axis=2
144 145
    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
146
    shape(X) = (2, 3, 4, 5), shape(Y) = (2, 1), with axis=0
147

L
Luo Tao 已提交
148 149
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 已提交
150

Y
Yu Yang 已提交
151 152
)DOC",
                               GetName(), GetEquation()));
G
gongweibao 已提交
153 154 155
  }

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

Y
Yu Yang 已提交
158
  virtual std::string GetEquation() const = 0;
G
gongweibao 已提交
159 160 161 162 163 164 165
};

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

C
chengduo 已提交
166
  void InferShape(framework::InferShapeContext *ctx) const override {
Q
Qiao Longfei 已提交
167 168 169 170 171 172 173 174
    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");
    auto out_dims = ctx->GetInputDim(framework::GradVarName("Out"));
G
gongweibao 已提交
175 176

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

Q
Qiao Longfei 已提交
179 180 181
    auto x_grad_name = framework::GradVarName("X");
    auto y_grad_name = framework::GradVarName("Y");
    if (ctx->HasOutput(x_grad_name)) {
182 183
      ctx->ShareDim("X", /*->*/ x_grad_name);
      ctx->ShareLoD("X", /*->*/ x_grad_name);
G
gongweibao 已提交
184
    }
Q
Qiao Longfei 已提交
185
    if (ctx->HasOutput(y_grad_name)) {
186 187
      ctx->ShareDim("Y", /*->*/ y_grad_name);
      ctx->ShareLoD("Y", /*->*/ y_grad_name);
G
gongweibao 已提交
188 189
    }
  }
190 191

  framework::OpKernelType GetExpectedKernelType(
C
chengduo 已提交
192
      const framework::ExecutionContext &ctx) const override {
193 194
    auto input_data_type = framework::ToDataType(
        ctx.Input<Tensor>(framework::GradVarName("Out"))->type());
195 196 197 198 199 200 201 202 203 204

#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 已提交
205
};
206 207 208 209 210 211 212 213

// 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 已提交
214
  void InferShape(framework::InferShapeContext *ctx) const override {
215 216 217 218 219
    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)) {
220 221
      ctx->ShareDim(framework::GradVarName("Out"), /*->*/ x_grad_name);
      ctx->ShareLoD(framework::GradVarName("Out"), /*->*/ x_grad_name);
222 223 224 225
    }
    auto y_grad_name = framework::GradVarName("Y");
    if (ctx->HasOutput(y_grad_name)) {
      PADDLE_ENFORCE(ctx->HasInput("Y"), "Input(Y) should not be null");
226 227 228

      ctx->ShareDim("Y", /*->*/ y_grad_name);
      ctx->ShareLoD("Y", /*->*/ y_grad_name);
229 230 231 232
    }
  }
};

233 234 235
template <typename T>
class ElemwiseGradKernel : public framework::OpKernel<T> {
 public:
C
chengduo 已提交
236 237
  void Compute(const framework::ExecutionContext &context) const override {
    auto *dx =
238 239
        context.Output<framework::LoDTensor>(framework::GradVarName("X"));
    if (dx != nullptr) {
C
chengduo 已提交
240
      auto &dout =
241 242 243 244 245 246
          *context.Input<framework::LoDTensor>(framework::GradVarName("Out"));
      dx->set_lod(dout.lod());
    }
  }
};

G
gongweibao 已提交
247 248
}  // namespace operators
}  // namespace paddle
Y
Yu Yang 已提交
249

250 251 252 253 254 255 256 257 258 259 260
/*
*/

#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 已提交
261
      auto *op = new paddle::framework::OpDesc();                            \
262 263 264 265 266 267 268 269 270 271 272
      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 已提交
273 274 275 276 277 278 279 280 281
#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__,                     \
282
                    ::paddle::operators::ElementwiseOpInferVarType,     \
Y
Yu Yang 已提交
283 284
                    ::paddle::framework::DefaultGradOpDescMaker<true>); \
  REGISTER_OPERATOR(op_type##_grad, ::paddle::operators::ElementwiseOpGrad)
285 286 287 288 289 290 291 292 293 294 295 296 297 298

#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,    \
                    op_type##GradMaker);                               \
  REGISTER_OPERATOR(op_type##_grad,                                    \
                    ::paddle::operators::ElementwiseOpExplicitGrad)