concat_op.cc 3.7 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"
16

S
Siddharth Goyal 已提交
17
#include <string>
18 19 20 21 22 23 24 25 26 27
#include <vector>

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

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

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

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

38 39 40 41
    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";
    }
42

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

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

89 90 91 92 93 94 95 96
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) {}

97
  void InferShape(framework::InferShapeContext *ctx) const override {
98 99 100 101
    ctx->SetOutputsDim(framework::GradVarName("X"), ctx->GetInputsDim("X"));
  }
};

102 103 104 105
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
106 107 108 109
REGISTER_OPERATOR(concat, ops::ConcatOp, ops::ConcatOpMaker,
                  paddle::framework::DefaultGradOpDescMaker<
                      false> /* set false to disable empty grad */)
REGISTER_OPERATOR(concat_grad, ops::ConcatOpGrad)
C
chengduoZH 已提交
110 111 112 113 114
REGISTER_OP_CPU_KERNEL(
    concat, ops::ConcatKernel<paddle::platform::CPUDeviceContext, float>)
REGISTER_OP_CPU_KERNEL(
    concat_grad,
    ops::ConcatGradKernel<paddle::platform::CPUDeviceContext, float>)