transpose_op.cc 4.1 KB
Newer Older
X
xzl 已提交
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/transpose_op.h"

namespace paddle {
namespace operators {

using framework::Tensor;

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

26
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
27 28 29 30
    PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) should not be null");
    PADDLE_ENFORCE(ctx->HasOutput("Out"), "Output(Out) should not be null");
    auto x_dims = ctx->GetInputDim("X");
    std::vector<int> axis = ctx->Attrs().Get<std::vector<int>>("axis");
X
xzl 已提交
31
    size_t x_rank = x_dims.size();
X
xzl 已提交
32
    size_t axis_size = axis.size();
X
xzl 已提交
33

X
xzl 已提交
34
    PADDLE_ENFORCE_EQ(x_rank, axis_size,
35
                      "The input tensor's rank(%d) "
36
                      "should be equal to the axis's size(%d)",
X
xzl 已提交
37
                      x_rank, axis_size);
38 39 40 41 42 43 44 45

    std::vector<int> count(axis_size, 0);
    for (size_t i = 0; i < axis_size; i++) {
      PADDLE_ENFORCE(
          axis[i] < static_cast<int>(axis_size) && ++count[axis[i]] == 1,
          "Each element of Attribute axis should be a unique value "
          "range from 0 to (dims - 1), "
          "where the dims is the axis's size");
X
xzl 已提交
46
    }
X
xzl 已提交
47

X
xzl 已提交
48
    framework::DDim out_dims(x_dims);
49
    for (size_t i = 0; i < axis_size; i++) {
X
xzl 已提交
50
      out_dims[i] = x_dims[axis[i]];
X
xzl 已提交
51
    }
Q
Qiao Longfei 已提交
52
    ctx->SetOutputDim("Out", out_dims);
X
xzl 已提交
53 54 55 56 57
  }
};

class TransposeOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Q
Qiao Longfei 已提交
58 59
  TransposeOpMaker(framework::OpProto* proto,
                   framework::OpAttrChecker* op_checker)
X
xzl 已提交
60
      : OpProtoAndCheckerMaker(proto, op_checker) {
61
    AddInput(
X
xzl 已提交
62
        "X",
63
        "(Tensor)The input tensor, tensors with rank at most 6 are supported");
X
xzl 已提交
64
    AddOutput("Out", "(Tensor)The output tensor");
X
xzl 已提交
65 66
    AddAttr<std::vector<int>>(
        "axis",
67
        "(vector<int>)A list of values, and the size of the list should be "
68
        "the same with the input tensor rank, the tensor will "
X
xzl 已提交
69
        "permute the axes according the the values given");
X
xzl 已提交
70
    AddComment(R"DOC(
71 72 73 74
Transpose Operator.

The input tensor will be permuted according to the axis values given.
The op functions similar to how numpy.transpose works in python.
75 76 77 78 79 80 81
For example:
 >> input = numpy.arange(6).reshape((2,3))
 >> input
 array([[0, 1, 2],
        [3, 4, 5]])
 >> axis = [1, 0]
 >> output = input.transpose(axis)
Q
Qiao Longfei 已提交
82
 >> output
83 84 85 86
 array([[0, 3],
        [1, 4],
		[2, 5]])
So, given a input tensor of shape(N, C, H, W) and the axis is {0, 2, 3, 1},
X
xzl 已提交
87
the output tensor shape will be (N, H, W, C)
88

X
xzl 已提交
89 90 91 92 93 94 95 96
)DOC");
  }
};

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

97
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
98 99 100 101 102 103 104 105
    PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) should not be null");
    PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Out")),
                   "Input(Out@GRAD) should not be null");
    auto x_dims = ctx->GetInputDim("X");
    ctx->SetOutputDim(framework::GradVarName("X"), x_dims);
    if (ctx->HasOutput(framework::GradVarName("X"))) {
      ctx->SetOutputDim(framework::GradVarName("X"), x_dims);
    }
X
xzl 已提交
106 107 108 109 110 111 112 113 114 115 116 117 118 119
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
REGISTER_OP(transpose, ops::TransposeOp, ops::TransposeOpMaker, transpose_grad,
            ops::TransposeOpGrad);
REGISTER_OP_CPU_KERNEL(transpose,
                       ops::TransposeKernel<paddle::platform::CPUPlace, float>);
REGISTER_OP_CPU_KERNEL(
    transpose_grad,
    ops::TransposeGradKernel<paddle::platform::CPUPlace, float>);