concat_op.cc 8.8 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2 3 4 5 6 7 8 9 10 11 12 13 14

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

http://www.apache.org/licenses/LICENSE-2.0

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
Yi Wang 已提交
15
#include "paddle/fluid/operators/concat_op.h"
1
123malin 已提交
16

P
phlrain 已提交
17
#include <memory>
S
Siddharth Goyal 已提交
18
#include <string>
19 20
#include <vector>

P
phlrain 已提交
21 22 23 24
#ifdef PADDLE_WITH_MKLDNN
#include <paddle/fluid/platform/mkldnn_helper.h>
#endif

25 26
namespace paddle {
namespace operators {
27
using Tensor = framework::Tensor;
28 29 30 31 32

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

33
  void InferShape(framework::InferShapeContext *ctx) const override {
34 35
    OP_INOUT_CHECK(ctx->HasInputs("X"), "Input", "X", "Concat");
    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "Concat");
36

37
    auto inputs_dims = ctx->GetInputsDim("X");
38

39
    const size_t inputs_num = inputs_dims.size();
40 41 42 43 44
    PADDLE_ENFORCE_GT(
        inputs_num, static_cast<size_t>(0),
        platform::errors::InvalidArgument(
            "The number of input tensors in concat op should > 0. But "
            "received inputs' length is 0."));
45
    if (inputs_num == 1) {
46 47
      VLOG(3) << "Warning: concat op have only one input, may waste memory";
    }
48

49 50 51 52 53 54 55 56 57 58 59 60 61
    if (ctx->HasInput("AxisTensor")) {
      auto out_dims =
          framework::make_ddim(std::vector<int>(inputs_dims[0].size(), -1));
      ctx->SetOutputDim("Out", out_dims);
      ctx->ShareLoD("X", /*->*/ "Out");
    } else {
      size_t axis =
          ComputeAxis(static_cast<int64_t>(ctx->Attrs().Get<int>("axis")),
                      static_cast<int64_t>(inputs_dims[0].size()));
      framework::DDim out_dims =
          ComputeAndCheckShape(ctx->IsRuntime(), inputs_dims, axis);
      if (out_dims[axis] < 0) {
        out_dims[axis] = -1;
62
      }
63 64
      ctx->SetOutputDim("Out", out_dims);
      ctx->ShareLoD("X", /*->*/ "Out");
65 66
    }
  }
P
phlrain 已提交
67 68 69 70

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
71
    auto inputs = ctx.MultiInput<Tensor>("X");
72 73
    auto input_data_type = framework::proto::VarType::Type(0);
    bool flag = 0;
74 75 76
    for (auto *input : inputs) {
      if (input->IsInitialized() && input->numel() > 0) {
        input_data_type = input->type();
77 78 79 80 81
        flag = 1;
        break;
      }
    }
    if (flag == 0) {
1
123malin 已提交
82 83
      PADDLE_THROW(platform::errors::InvalidArgument(
          "All Inputs of Concat OP are Empty!"));
84
    }
P
phlrain 已提交
85
#ifdef PADDLE_WITH_MKLDNN
86
    if (this->CanMKLDNNBeUsed(ctx)) {
P
phlrain 已提交
87 88 89 90 91 92 93
      return framework::OpKernelType(input_data_type, ctx.GetPlace(),
                                     framework::DataLayout::kMKLDNN,
                                     framework::LibraryType::kMKLDNN);
    }
#endif
    return framework::OpKernelType(input_data_type, ctx.GetPlace());
  }
94 95 96 97 98 99 100 101 102 103

  framework::OpKernelType GetKernelTypeForVar(
      const std::string &var_name, const Tensor &tensor,
      const framework::OpKernelType &expected_kernel_type) const override {
    if (var_name == "AxisTensor") {
      return expected_kernel_type;
    }
    return framework::OpKernelType(expected_kernel_type.data_type_,
                                   tensor.place(), tensor.layout());
  }
104 105 106 107
};

class ConcatOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
108
  void Make() override {
109 110
    AddInput("X", "Input tensors of concat operator.").AsDuplicable();
    AddOutput("Out", "Output tensor of concat operator.");
P
phlrain 已提交
111 112 113 114
    AddAttr<bool>(
        "use_mkldnn",
        "(bool, default false) Indicates if MKL-DNN kernel will be used")
        .SetDefault(false);
115
    AddAttr<int>("axis",
116 117 118 119
                 "The axis along which the input tensors will be concatenated."
                 "The axis could also be negative numbers. Negative axis is "
                 "interpreted as counting from the end of the rank."
                 "i.e., axis + rank(X) th dimension.")
120
        .SetDefault(0);
121 122 123 124 125 126
    AddInput("AxisTensor",
             "(Tensor) The axis along which the input tensors will be "
             "concatenated.  "
             "It has higher priority than Attr(axis). "
             "The shape of AxisTensor must be [1].")
        .AsDispensable();
127 128 129 130
    AddAttr<bool>(
        "use_quantizer",
        "(bool, default false) "
        "This parameter is no longer used. Use 'mkldnn_data_type' instead.")
131
        .SetDefault(false);
132 133 134 135 136
    AddAttr<std::string>(
        "mkldnn_data_type",
        "(string, default \"float32\"). Data type of mkldnn kernel")
        .SetDefault("float32")
        .InEnum({"float32", "int8", "bfloat16"});
137 138 139 140 141 142 143 144 145 146 147 148 149
    AddComment(R"DOC(
Concat Operator.

Concatenate the input tensors along dimension axis.
Examples:
  Input[0] = [[1,2],[3,4]]
  Input[1] = [[5,6]]
  axis = 0
  Output = [[1,2],
            [3,4],
            [5,6]]

)DOC");
150 151 152
  }
};

153 154
class ConcatOpGrad : public framework::OperatorWithKernel {
 public:
P
phlrain 已提交
155
  using framework::OperatorWithKernel::OperatorWithKernel;
156

157
  void InferShape(framework::InferShapeContext *ctx) const override {
C
chengduo 已提交
158 159 160
    auto in_x = "X";
    auto out_x_g_n = framework::GradVarName(in_x);
    ctx->SetOutputsDim(out_x_g_n, ctx->GetInputsDim(in_x));
H
hong 已提交
161 162

    ctx->ShareAllLoD(in_x, out_x_g_n);
163
  }
P
phlrain 已提交
164 165 166 167

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
168 169 170
    return framework::OpKernelType(OperatorWithKernel::IndicateVarDataType(
                                       ctx, framework::GradVarName("Out")),
                                   ctx.GetPlace());
P
phlrain 已提交
171
  }
172 173 174 175 176 177 178 179 180 181

  framework::OpKernelType GetKernelTypeForVar(
      const std::string &var_name, const Tensor &tensor,
      const framework::OpKernelType &expected_kernel_type) const override {
    if (var_name == "AxisTensor") {
      return expected_kernel_type;
    }
    return framework::OpKernelType(expected_kernel_type.data_type_,
                                   tensor.place(), tensor.layout());
  }
P
phlrain 已提交
182 183
};

184
DECLARE_NO_NEED_BUFFER_VARS_INFERER(ConcatOpGradNoNeedBufferVarInferer, "X");
P
phlrain 已提交
185

H
hong 已提交
186 187
template <typename T>
class ConcatGradOpMaker : public framework::SingleGradOpMaker<T> {
P
phlrain 已提交
188
 public:
H
hong 已提交
189
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
P
phlrain 已提交
190 191

 protected:
192
  void Apply(GradOpPtr<T> op) const override {
P
phlrain 已提交
193
    op->SetType("concat_grad");
H
hong 已提交
194
    op->SetInput("X", this->Input("X"));
H
hong 已提交
195 196 197
    if (this->HasInput("AxisTensor")) {
      op->SetInput("AxisTensor", this->Input("AxisTensor"));
    }
H
hong 已提交
198 199 200
    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
    op->SetOutput(framework::GradVarName("X"), this->InputGrad("X", false));
    op->SetAttrMap(this->Attrs());
P
phlrain 已提交
201
  }
202 203
};

C
ceci3 已提交
204 205 206 207 208 209 210 211 212 213 214 215 216 217
template <typename T>
class ConcatDoubleGradOpMaker : public framework::SingleGradOpMaker<T> {
 public:
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;

 protected:
  void Apply(GradOpPtr<T> grad_op) const override {
    grad_op->SetType("concat");
    grad_op->SetInput("X", this->OutputGrad(framework::GradVarName("X")));
    grad_op->SetOutput("Out", this->InputGrad(framework::GradVarName("Out")));
    grad_op->SetAttrMap(this->Attrs());
  }
};

218 219 220 221
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
222
REGISTER_OPERATOR(concat, ops::ConcatOp, ops::ConcatOpMaker,
H
hong 已提交
223 224
                  ops::ConcatGradOpMaker<paddle::framework::OpDesc>,
                  ops::ConcatGradOpMaker<paddle::imperative::OpBase>);
P
phlrain 已提交
225
REGISTER_OPERATOR(concat_grad, ops::ConcatOpGrad,
C
ceci3 已提交
226 227
                  ops::ConcatDoubleGradOpMaker<paddle::framework::OpDesc>,
                  ops::ConcatDoubleGradOpMaker<paddle::imperative::OpBase>,
228
                  ops::ConcatOpGradNoNeedBufferVarInferer);
C
chengduoZH 已提交
229
REGISTER_OP_CPU_KERNEL(
230 231
    concat, ops::ConcatKernel<paddle::platform::CPUDeviceContext, double>,
    ops::ConcatKernel<paddle::platform::CPUDeviceContext, float>,
232
    ops::ConcatKernel<paddle::platform::CPUDeviceContext, bool>,
233
    ops::ConcatKernel<paddle::platform::CPUDeviceContext, int64_t>,
234 235
    ops::ConcatKernel<paddle::platform::CPUDeviceContext,
                      paddle::platform::float16>,
236
    ops::ConcatKernel<paddle::platform::CPUDeviceContext, int>);
C
chengduoZH 已提交
237 238
REGISTER_OP_CPU_KERNEL(
    concat_grad,
239 240
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext, double>,
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext, float>,
241
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext, bool>,
242
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext, int64_t>,
243 244
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext,
                          paddle::platform::float16>,
245
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext, int>);