elementwise_op.h 11.1 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
16
#include <string>
17
#include "paddle/fluid/framework/data_layout.h"
Y
Yi Wang 已提交
18 19
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/operator.h"
20 21 22
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif
G
gongweibao 已提交
23 24 25 26 27 28 29 30 31

namespace paddle {
namespace operators {

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

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

    auto x_dim = ctx->GetInputDim("X");
    auto y_dim = ctx->GetInputDim("Y");
G
gongweibao 已提交
42
    PADDLE_ENFORCE_GE(x_dim.size(), y_dim.size(),
43
                      "Rank of first input must >= rank of second input.");
Q
Qiao Longfei 已提交
44 45
    ctx->SetOutputDim("Out", x_dim);
    ctx->ShareLoD("X", /*->*/ "Out");
G
gongweibao 已提交
46
  }
47 48 49 50 51 52 53 54 55 56 57 58 59 60 61

  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override {
    auto input_data_type =
        framework::ToDataType(ctx.Input<Tensor>("X")->type());

#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 已提交
62 63
};

64 65 66 67 68 69 70 71 72 73 74 75
class ElementwiseOpInferVarType : public framework::VarTypeInference {
 public:
  void operator()(const framework::OpDesc& op_desc,
                  framework::BlockDesc* block) const override {
    auto x_name = op_desc.Input("X")[0];
    auto out_name = op_desc.Output("Out")[0];
    auto& x = block->FindRecursiveOrCreateVar(x_name);
    auto& out = block->FindRecursiveOrCreateVar(out_name);
    out.SetType(x.GetType());
  }
};

G
gongweibao 已提交
76 77
class ElementwiseOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
78
  void Make() final {
C
caoying03 已提交
79 80
    AddInput("X", "(Tensor), The first input tensor of elementwise op.");
    AddInput("Y", "(Tensor), The second input tensor of elementwise op.");
81 82 83
    // AddOutput("SavedShape", "(Tensor), save X, Y shape for grad to save
    // memory.").AsIntermediate();
    AddOutput("Out", "The output of elementwise op.");
G
gongweibao 已提交
84
    AddAttr<int>("axis",
C
caoying03 已提交
85 86
                 "(int, default -1). The start dimension index "
                 "for broadcasting Y onto X.")
G
gongweibao 已提交
87 88
        .SetDefault(-1)
        .EqualGreaterThan(-1);
89 90
    AddAttr<bool>("use_mkldnn", "(bool, default false). Used by MKLDNN.")
        .SetDefault(false);
Y
Yu Yang 已提交
91
    AddComment(string::Sprintf(R"DOC(
L
Luo Tao 已提交
92
Limited Elementwise %s Operator
K
kexinzhao 已提交
93 94 95

The equation is:

Y
Yu Yang 已提交
96
$$%s$$
K
kexinzhao 已提交
97

L
Luo Tao 已提交
98 99
- $X$: a tensor of any dimension. 
- $Y$: a tensor whose dimensions must be less than or equal to the dimensions of $X$.
K
kexinzhao 已提交
100 101

There are two cases for this operator:
102

L
Luo Tao 已提交
103 104
1. The shape of $Y$ is the same with $X$.
2. The shape of $Y$ is a continuous subsequence of $X$.
K
kexinzhao 已提交
105 106

For case 2:
107

L
Luo Tao 已提交
108 109 110 111 112
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 已提交
113

L
Luo Tao 已提交
114
For example:
115

116
  .. code-block:: python
G
gongweibao 已提交
117

118 119
    shape(X) = (2, 3, 4, 5), shape(Y) = (,)
    shape(X) = (2, 3, 4, 5), shape(Y) = (5,)
L
Luo Tao 已提交
120
    shape(X) = (2, 3, 4, 5), shape(Y) = (4, 5), with axis=-1(default) or axis=2
121 122
    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
123
    shape(X) = (2, 3, 4, 5), shape(Y) = (2, 1), with axis=0
124

L
Luo Tao 已提交
125 126
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 已提交
127

Y
Yu Yang 已提交
128 129
)DOC",
                               GetName(), GetEquation()));
130
    SetReuse();
G
gongweibao 已提交
131 132 133
  }

 protected:
Y
Yu Yang 已提交
134 135
  virtual std::string GetName() const = 0;
  virtual std::string GetEquation() const = 0;
136
  virtual void SetReuse() {}
G
gongweibao 已提交
137 138 139 140 141 142 143
};

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

144
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
145 146 147 148 149 150 151 152
    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 已提交
153 154

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

Q
Qiao Longfei 已提交
157 158 159 160
    auto x_grad_name = framework::GradVarName("X");
    auto y_grad_name = framework::GradVarName("Y");
    if (ctx->HasOutput(x_grad_name)) {
      ctx->SetOutputDim(x_grad_name, x_dims);
G
gongweibao 已提交
161
    }
Q
Qiao Longfei 已提交
162 163
    if (ctx->HasOutput(y_grad_name)) {
      ctx->SetOutputDim(y_grad_name, y_dims);
G
gongweibao 已提交
164 165
    }
  }
166 167 168

  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override {
169 170
    auto input_data_type = framework::ToDataType(
        ctx.Input<Tensor>(framework::GradVarName("Out"))->type());
171 172 173 174 175 176 177 178 179 180

#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 已提交
181
};
182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207

// 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;

  void InferShape(framework::InferShapeContext* ctx) const override {
    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)) {
      auto out_dims = ctx->GetInputDim(framework::GradVarName("Out"));
      ctx->SetOutputDim(x_grad_name, out_dims);
    }
    auto y_grad_name = framework::GradVarName("Y");
    if (ctx->HasOutput(y_grad_name)) {
      PADDLE_ENFORCE(ctx->HasInput("Y"), "Input(Y) should not be null");
      auto y_dims = ctx->GetInputDim("Y");
      ctx->SetOutputDim(y_grad_name, y_dims);
    }
  }
};

208 209 210 211 212 213 214 215 216 217 218 219 220 221
template <typename T>
class ElemwiseGradKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    auto* dx =
        context.Output<framework::LoDTensor>(framework::GradVarName("X"));
    if (dx != nullptr) {
      auto& dout =
          *context.Input<framework::LoDTensor>(framework::GradVarName("Out"));
      dx->set_lod(dout.lod());
    }
  }
};

G
gongweibao 已提交
222 223
}  // namespace operators
}  // namespace paddle
Y
Yu Yang 已提交
224

225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247
/*
*/

#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 {      \
      auto* op = new paddle::framework::OpDesc();                            \
      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 已提交
248 249 250 251 252 253 254 255 256
#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__,                     \
257
                    ::paddle::operators::ElementwiseOpInferVarType,     \
Y
Yu Yang 已提交
258 259
                    ::paddle::framework::DefaultGradOpDescMaker<true>); \
  REGISTER_OPERATOR(op_type##_grad, ::paddle::operators::ElementwiseOpGrad)
260 261 262 263 264 265 266 267 268 269 270 271 272 273 274

#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; }       \
    virtual void SetReuse() { Reuse(__VA_ARGS__); }                    \
  };                                                                   \
  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)