concat_op.cc 8.1 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"
P
phlrain 已提交
16
#include <memory>
S
Siddharth Goyal 已提交
17
#include <string>
18 19
#include <vector>

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

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

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

32
  void InferShape(framework::InferShapeContext *ctx) const override {
33
    PADDLE_ENFORCE_GE(ctx->Inputs("X").size(), 1UL,
H
hutuxian 已提交
34
                      "Inputs(X) of ConcatOp should not be empty.");
35 36
    PADDLE_ENFORCE_EQ(ctx->HasOutput("Out"), true,
                      "Output(Out) of ConcatOp should not be null.");
37

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

40 41
    const size_t inputs_num = inputs_dims.size();
    PADDLE_ENFORCE_GT(inputs_num, 0,
42 43
                      "ShapeError: Input tensors count should > 0. But "
                      "recevied inputs' length is 0.");
44
    if (inputs_num == 1) {
45 46
      VLOG(3) << "Warning: concat op have only one input, may waste memory";
    }
47

48 49 50 51 52 53 54 55 56 57 58 59 60
    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;
61
      }
62 63
      ctx->SetOutputDim("Out", out_dims);
      ctx->ShareLoD("X", /*->*/ "Out");
64 65
    }
  }
P
phlrain 已提交
66 67 68 69

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
70
    auto inputs = ctx.MultiInput<Tensor>("X");
71 72
    auto input_data_type = framework::proto::VarType::Type(0);
    bool flag = 0;
73 74 75
    for (auto *input : inputs) {
      if (input->IsInitialized() && input->numel() > 0) {
        input_data_type = input->type();
76 77 78 79 80 81 82
        flag = 1;
        break;
      }
    }
    if (flag == 0) {
      PADDLE_THROW("All Inputs of Concat OP are Empty!");
    }
P
phlrain 已提交
83 84 85 86 87 88 89 90 91
#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());
  }
92 93 94 95 96 97 98 99 100 101

  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());
  }
102 103 104 105
};

class ConcatOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
106
  void Make() override {
107 108
    AddInput("X", "Input tensors of concat operator.").AsDuplicable();
    AddOutput("Out", "Output tensor of concat operator.");
P
phlrain 已提交
109 110 111 112
    AddAttr<bool>(
        "use_mkldnn",
        "(bool, default false) Indicates if MKL-DNN kernel will be used")
        .SetDefault(false);
113
    AddAttr<int>("axis",
114 115 116 117
                 "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.")
118
        .SetDefault(0);
119 120 121 122 123 124
    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();
125 126 127 128 129 130
    AddAttr<bool>("use_quantizer",
                  "(bool, default false) "
                  "Set to true for operators that should be quantized and use "
                  "int8 kernel. "
                  "Only used on CPU.")
        .SetDefault(false);
131 132 133 134 135 136 137 138 139 140 141 142 143
    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");
144 145 146
  }
};

147 148
class ConcatOpGrad : public framework::OperatorWithKernel {
 public:
P
phlrain 已提交
149
  using framework::OperatorWithKernel::OperatorWithKernel;
150

151
  void InferShape(framework::InferShapeContext *ctx) const override {
C
chengduo 已提交
152 153 154 155 156 157 158 159 160 161 162 163 164 165
    auto in_x = "X";
    auto out_x_g_n = framework::GradVarName(in_x);
    ctx->SetOutputsDim(out_x_g_n, ctx->GetInputsDim(in_x));
    auto &in_names = ctx->Inputs(in_x);
    auto &out_names = ctx->Outputs(out_x_g_n);
    PADDLE_ENFORCE_EQ(
        in_names.size(), out_names.size(),
        "The number of arguments in %s[%d] and %s[%d] is not equal.", in_x,
        in_names.size(), out_x_g_n, out_names.size());
    for (size_t i = 0; i < in_names.size(); ++i) {
      if (out_names[i] != framework::kEmptyVarName) {
        ctx->ShareLoD(in_x, out_x_g_n, i, i);
      }
    }
166
  }
P
phlrain 已提交
167 168 169 170

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

  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 已提交
185 186 187 188 189
};

DECLARE_NO_NEED_BUFFER_VARS_INFERENCE(ConcatOpGradNoNeedBufferVarInference,
                                      "X");

H
hong 已提交
190 191
template <typename T>
class ConcatGradOpMaker : public framework::SingleGradOpMaker<T> {
P
phlrain 已提交
192
 public:
H
hong 已提交
193
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
P
phlrain 已提交
194 195

 protected:
H
hong 已提交
196 197
  std::unique_ptr<T> Apply() const override {
    std::unique_ptr<T> op(new T());
P
phlrain 已提交
198
    op->SetType("concat_grad");
H
hong 已提交
199 200 201 202 203
    op->SetInput("X", this->Input("X"));
    op->SetInput("AxisTensor", this->Input("AxisTensor"));
    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
    op->SetOutput(framework::GradVarName("X"), this->InputGrad("X", false));
    op->SetAttrMap(this->Attrs());
P
phlrain 已提交
204 205
    return op;
  }
206 207
};

208 209 210 211
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
212
REGISTER_OPERATOR(concat, ops::ConcatOp, ops::ConcatOpMaker,
H
hong 已提交
213 214
                  ops::ConcatGradOpMaker<paddle::framework::OpDesc>,
                  ops::ConcatGradOpMaker<paddle::imperative::OpBase>);
P
phlrain 已提交
215 216
REGISTER_OPERATOR(concat_grad, ops::ConcatOpGrad,
                  ops::ConcatOpGradNoNeedBufferVarInference);
C
chengduoZH 已提交
217
REGISTER_OP_CPU_KERNEL(
218 219 220 221
    concat, ops::ConcatKernel<paddle::platform::CPUDeviceContext, double>,
    ops::ConcatKernel<paddle::platform::CPUDeviceContext, float>,
    ops::ConcatKernel<paddle::platform::CPUDeviceContext, int64_t>,
    ops::ConcatKernel<paddle::platform::CPUDeviceContext, int>);
C
chengduoZH 已提交
222 223
REGISTER_OP_CPU_KERNEL(
    concat_grad,
224 225 226 227
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext, double>,
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext, float>,
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext, int64_t>,
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext, int>);