elementwise_op.h 11.1 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
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
6

7
    http://www.apache.org/licenses/LICENSE-2.0
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. */
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
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 {
33
    PADDLE_ENFORCE(ctx->HasInput("X"),
34
                   "Input(X) of elementwise op should not be null.");
35
    PADDLE_ENFORCE(ctx->HasInput("Y"),
36
                   "Input(Y) of elementwise op should not be null.");
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");
42
    PADDLE_ENFORCE_GE(x_dim.size(), y_dim.size(),
43
                      "Rank of first input must >= rank of second input.");
44 45
    ctx->SetOutputDim("Out", x_dim);
    ctx->ShareLoD("X", /*->*/ "Out");
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());
  }
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());
  }
};

76 77
class ElementwiseOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
78
  void Make() final {
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.");
84
    AddAttr<int>("axis",
85 86
                 "(int, default -1). The start dimension index "
                 "for broadcasting Y onto X.")
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
93 94 95

The equation is:

Y
Yu Yang 已提交
96
$$%s$$
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$.
100 101

There are two cases for this operator:
102

103 104
1. The shape of $Y$ is the same with $X$.
2. The shape of $Y$ is a continuous subsequence of $X$.
105 106

For case 2:
107

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).
113

114
For example:
115

116
  .. code-block:: python
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

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$.
127

Y
Yu Yang 已提交
128 129
)DOC",
                               GetName(), GetEquation()));
130
    SetReuse();
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() {}
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 {
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"));
153 154

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

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);
161
    }
162 163
    if (ctx->HasOutput(y_grad_name)) {
      ctx->SetOutputDim(y_grad_name, y_dims);
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());
  }
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());
    }
  }
};

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)
反馈
建议
客服 返回
顶部