concat_op.cc 3.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

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. */

#include "paddle/operators/concat_op.h"
#include <vector>

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

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

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

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

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

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

class ConcatOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
61
  ConcatOpMaker(OpProto *proto, OpAttrChecker *op_checker)
62
      : OpProtoAndCheckerMaker(proto, op_checker) {
63 64 65 66
    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.")
67
        .SetDefault(0);
68 69 70 71 72 73 74 75 76 77 78 79 80
    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");
81 82 83
  }
};

84 85 86 87 88 89 90 91
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) {}

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

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

namespace ops = paddle::operators;
101 102
REGISTER_OP_EX(concat, ops::ConcatOp, ops::ConcatOpMaker, concat_grad,
               ops::ConcatOpGrad, false)
103 104
REGISTER_OP_CPU_KERNEL(concat,
                       ops::ConcatKernel<paddle::platform::CPUPlace, float>)
105 106
REGISTER_OP_CPU_KERNEL(concat_grad,
                       ops::ConcatGradKernel<paddle::platform::CPUPlace, float>)