concat_op.cc 6.0 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"
S
sneaxiy 已提交
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 26 27 28 29 30 31
namespace paddle {
namespace operators {
using framework::Tensor;

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,
34
                      "Inputs(X) of ConcatOp should be empty.");
Q
Qiao Longfei 已提交
35 36
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of ConcatOp should not be null.");
37

Q
Qiao Longfei 已提交
38 39
    auto ins = ctx->GetInputsDim("X");
    size_t axis = static_cast<size_t>(ctx->Attrs().Get<int>("axis"));
40
    const size_t n = ins.size();
41

42 43 44 45
    PADDLE_ENFORCE_GT(n, 0, "Input tensors count should > 0.");
    if (n == 1) {
      VLOG(3) << "Warning: concat op have only one input, may waste memory";
    }
46

Q
Qiao Longfei 已提交
47
    auto out_dims = ins[0];
48 49 50 51
    size_t in_zero_dims_size = out_dims.size();
    for (size_t i = 1; i < n; i++) {
      for (size_t j = 0; j < in_zero_dims_size; j++) {
        if (j == axis) {
Q
Qiao Longfei 已提交
52
          out_dims[axis] += ins[i][j];
Q
Qiao Longfei 已提交
53
        } else {
P
phlrain 已提交
54 55 56 57 58 59
          if (ctx->IsRuntime()) {
            // check all shape in run time
            PADDLE_ENFORCE_EQ(out_dims[j], ins[i][j],
                              "Input tensors should have the same "
                              "elements except the specify axis.");
          } else {
P
phlrain 已提交
60
            // not check -1 with other in compile time
P
phlrain 已提交
61
            if (out_dims[j] > 0 && ins[i][j] > 0) {
P
phlrain 已提交
62 63 64 65 66
              PADDLE_ENFORCE_EQ(out_dims[j], ins[i][j],
                                "Input tensors should have the same "
                                "elements except the specify axis.");
            }
          }
67 68 69
        }
      }
    }
Q
Qiao Longfei 已提交
70 71 72
    if (out_dims[axis] < 0) {
      out_dims[axis] = -1;
    }
Q
Qiao Longfei 已提交
73
    ctx->SetOutputDim("Out", out_dims);
Q
Qiao Longfei 已提交
74
    ctx->ShareLoD("X", /*->*/ "Out");
75
  }
P
phlrain 已提交
76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
    auto input_data_type =
        framework::GetDataTypeOfVar(ctx.MultiInputVar("X")[0]);

#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
};

class ConcatOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
96
  void Make() override {
97 98
    AddInput("X", "Input tensors of concat operator.").AsDuplicable();
    AddOutput("Out", "Output tensor of concat operator.");
P
phlrain 已提交
99 100 101 102
    AddAttr<bool>(
        "use_mkldnn",
        "(bool, default false) Indicates if MKL-DNN kernel will be used")
        .SetDefault(false);
103 104
    AddAttr<int>("axis",
                 "The axis along which the input tensors will be concatenated.")
105
        .SetDefault(0);
106 107 108 109 110 111 112 113 114 115 116 117 118
    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");
119 120 121
  }
};

122 123
class ConcatOpGrad : public framework::OperatorWithKernel {
 public:
S
sneaxiy 已提交
124
  using framework::OperatorWithKernel::OperatorWithKernel;
125

126
  void InferShape(framework::InferShapeContext *ctx) const override {
C
chengduo 已提交
127 128 129 130 131 132 133 134 135 136 137 138 139 140
    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);
      }
    }
141
  }
S
sneaxiy 已提交
142 143 144 145 146 147 148 149

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
    return framework::OpKernelType(
        ctx.Input<Tensor>(framework::GradVarName("Out"))->type(),
        ctx.GetPlace());
  }
150 151
};

S
sneaxiy 已提交
152 153 154
DECLARE_NO_NEED_BUFFER_VARS_INFERENCE(ConcatOpGradNoNeedBufferVarInference,
                                      "X");

155 156 157 158
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
159 160
REGISTER_OPERATOR(concat, ops::ConcatOp, ops::ConcatOpMaker,
                  paddle::framework::DefaultGradOpDescMaker<
161
                      false> /* set false to disable empty grad */);
S
sneaxiy 已提交
162 163
REGISTER_OPERATOR(concat_grad, ops::ConcatOpGrad,
                  ops::ConcatOpGradNoNeedBufferVarInference);
C
chengduoZH 已提交
164
REGISTER_OP_CPU_KERNEL(
165 166 167 168
    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 已提交
169 170
REGISTER_OP_CPU_KERNEL(
    concat_grad,
171 172 173 174
    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>);