concat_op.cc 3.5 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
Siddharth Goyal 已提交
16
#include <string>
17 18 19 20 21 22 23 24 25 26
#include <vector>

namespace paddle {
namespace operators {
using framework::Tensor;

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

27
  void InferShape(framework::InferShapeContext *ctx) const override {
28
    PADDLE_ENFORCE_GE(ctx->Inputs("X").size(), 1UL,
29
                      "Inputs(X) of ConcatOp should be empty.");
Q
Qiao Longfei 已提交
30 31
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of ConcatOp should not be null.");
32

Q
Qiao Longfei 已提交
33 34
    auto ins = ctx->GetInputsDim("X");
    size_t axis = static_cast<size_t>(ctx->Attrs().Get<int>("axis"));
35
    const size_t n = ins.size();
36 37 38

    PADDLE_ENFORCE_GT(n, 1, "Input tensors count should > 1.");

Q
Qiao Longfei 已提交
39
    auto out_dims = ins[0];
40 41 42 43
    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 已提交
44
          out_dims[axis] += ins[i][j];
Q
Qiao Longfei 已提交
45 46 47 48
        } else {
          PADDLE_ENFORCE_EQ(out_dims[j], ins[i][j],
                            "Input tensors should have the same "
                            "elements except the specify axis.");
49 50 51
        }
      }
    }
Q
Qiao Longfei 已提交
52 53 54
    if (out_dims[axis] < 0) {
      out_dims[axis] = -1;
    }
Q
Qiao Longfei 已提交
55
    ctx->SetOutputDim("Out", out_dims);
Q
Qiao Longfei 已提交
56
    ctx->ShareLoD("X", /*->*/ "Out");
57 58 59 60 61
  }
};

class ConcatOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
62
  ConcatOpMaker(OpProto *proto, OpAttrChecker *op_checker)
63
      : OpProtoAndCheckerMaker(proto, op_checker) {
64 65 66 67
    AddInput("X", "Input tensors of concat operator.").AsDuplicable();
    AddOutput("Out", "Output tensor of concat operator.");
    AddAttr<int>("axis",
                 "The axis along which the input tensors will be concatenated.")
68
        .SetDefault(0);
69 70 71 72 73 74 75 76 77 78 79 80 81
    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");
82 83 84
  }
};

85 86 87 88 89 90 91 92
class ConcatOpGrad : public framework::OperatorWithKernel {
 public:
  ConcatOpGrad(const std::string &type,
               const framework::VariableNameMap &inputs,
               const framework::VariableNameMap &outputs,
               const framework::AttributeMap &attrs)
      : OperatorWithKernel(type, inputs, outputs, attrs) {}

93
  void InferShape(framework::InferShapeContext *ctx) const override {
94 95 96 97
    ctx->SetOutputsDim(framework::GradVarName("X"), ctx->GetInputsDim("X"));
  }
};

98 99 100 101
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
102 103
REGISTER_OP_EX(concat, ops::ConcatOp, ops::ConcatOpMaker, concat_grad,
               ops::ConcatOpGrad, false)
C
chengduoZH 已提交
104 105 106 107 108
REGISTER_OP_CPU_KERNEL(
    concat, ops::ConcatKernel<paddle::platform::CPUDeviceContext, float>)
REGISTER_OP_CPU_KERNEL(
    concat_grad,
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext, float>)