space_to_depth_op.cc 5.1 KB
Newer Older
J
JiabinYang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.

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

J
JiabinYang 已提交
15
#include "paddle/fluid/operators/space_to_depth_op.h"
J
JiabinYang 已提交
16 17 18 19 20 21
#include <string>
#include <vector>

namespace paddle {
namespace operators {

J
JiabinYang 已提交
22
class SpaceToDepthOp : public framework::OperatorWithKernel {
J
JiabinYang 已提交
23 24 25 26 27
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

  void InferShape(framework::InferShapeContext* ctx) const override {
    PADDLE_ENFORCE(ctx->HasInput("X"),
J
JiabinYang 已提交
28
                   "Input(X) of SpaceToDepthOp should not be null.");
J
JiabinYang 已提交
29
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
J
JiabinYang 已提交
30
                   "Output(Out) of SpaceToDepthOp should not be null.");
J
JiabinYang 已提交
31 32 33 34 35

    auto x_dims = ctx->GetInputDim("X");
    PADDLE_ENFORCE_EQ(x_dims.size(), 4, "input should be a 4D tensor");
    auto stride = ctx->Attrs().Get<int64_t>("stride");

J
JiabinYang 已提交
36
    PADDLE_ENFORCE_GT(stride, 1, "The stride should be Greater than 1");
J
JiabinYang 已提交
37 38 39 40
    PADDLE_ENFORCE_GT(x_dims[1], 0, "input channel should be Greater than 0");
    PADDLE_ENFORCE_GT(x_dims[2], 0, "input Height should be Greater than 0");
    PADDLE_ENFORCE_GT(x_dims[3], 0, "input Width should be Greater than 0");

J
JiabinYang 已提交
41 42 43 44 45 46 47 48 49
    PADDLE_ENFORCE_EQ(x_dims[1] % (stride * stride), 0,
                      "input channel should be divisible of the square of "
                      "SpaceToDepthOp stride");
    PADDLE_ENFORCE_EQ(x_dims[2] % (stride), 0,
                      "input Height should be divisible of the square of "
                      "SpaceToDepthOp stride");
    PADDLE_ENFORCE_EQ(x_dims[3] % (stride), 0,
                      "input Width should be divisible of the square of "
                      "SpaceToDepthOp stride");
J
JiabinYang 已提交
50

J
JiabinYang 已提交
51 52
    VLOG(3) << "SpaceToDepthOp operator x.shape=" << x_dims
            << "Attribute stride" << stride << std::endl;
J
JiabinYang 已提交
53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71

    std::vector<int64_t> output_shape(4, 0);  // [B,C,H,W]
    output_shape[0] = x_dims[0];
    output_shape[1] = x_dims[1] * stride * stride;
    output_shape[2] = x_dims[2] / stride;
    output_shape[3] = x_dims[3] / stride;

    auto out_dims = framework::make_ddim(output_shape);

    ctx->SetOutputDim("Out", out_dims);

    if (x_dims[0] == out_dims[0]) {
      // Only pass LoD when the first dimension of output and Input(X)
      // are the same.
      ctx->ShareLoD("X", /*->*/ "Out");
    }
  }
};

J
JiabinYang 已提交
72
class SpaceToDepthOpMaker : public framework::OpProtoAndCheckerMaker {
J
JiabinYang 已提交
73 74 75
 public:
  void Make() override {
    AddInput("X",
J
JiabinYang 已提交
76 77
             "(Tensor). The input should be a 4D tensor B * C * W * H of "
             "SpaceToDepthOp "
J
JiabinYang 已提交
78 79 80
             "operator.");
    AddOutput("Out",
              "(Tensor), The output should be a 4D tensor B * C2 * W2 * H2 of "
J
JiabinYang 已提交
81 82 83 84 85 86
              "SpaceToDepthOp operator.");
    AddAttr<int64_t>(
        "stride",
        "(int64_t, default 2) stride used to do change Space To Depth.")
        .SetDefault(2)
        .GreaterThan(1);
J
JiabinYang 已提交
87 88 89 90 91 92 93 94 95
    AddComment(R"DOC(
        reorg operator used in Yolo v2.
        The equation is: C2 = C1/stride * stride, W2 = W1 ∗ stride + offset % stride, H2 = H1 ∗ stride + offset / stride, 

        Reshape Input(X) into the shape according to Attr(stride). The
        data in Input(X) are unchanged.

        Examples:

J
JiabinYang 已提交
96 97
            1. Given a 4-D tensor Input(X) with a shape [128, 2048, 26, 26], and the stride is 2, the reorg operator will transform Input(X)
            into a 4-D tensor with shape [128, 2048, 13, 13] and leaving Input(X)'s data unchanged.
J
JiabinYang 已提交
98 99 100 101 102

    )DOC");
  }
};

J
JiabinYang 已提交
103
class SpaceToDepthGradOp : public framework::OperatorWithKernel {
J
JiabinYang 已提交
104 105 106 107 108 109 110 111 112 113 114 115 116 117 118
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

  void InferShape(framework::InferShapeContext* ctx) const override {
    PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) shouldn't be null.");
    PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Out")),
                   "Input(Out@GRAD) shouldn't be null.");
    ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
  }
};
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

J
JiabinYang 已提交
119
REGISTER_OPERATOR(space_to_depth, ops::SpaceToDepthOp, ops::SpaceToDepthOpMaker,
J
JiabinYang 已提交
120
                  paddle::framework::DefaultGradOpDescMaker<true>);
J
JiabinYang 已提交
121
REGISTER_OPERATOR(space_to_depth_grad, ops::SpaceToDepthGradOp);
J
JiabinYang 已提交
122
REGISTER_OP_CPU_KERNEL(
J
JiabinYang 已提交
123 124 125 126
    space_to_depth,
    ops::SpaceToDepthKernel<paddle::platform::CPUDeviceContext, float>,
    ops::SpaceToDepthKernel<paddle::platform::CPUDeviceContext, double>,
    ops::SpaceToDepthKernel<paddle::platform::CPUDeviceContext, int64_t>);
J
JiabinYang 已提交
127
REGISTER_OP_CPU_KERNEL(
J
JiabinYang 已提交
128 129 130 131
    space_to_depth_grad,
    ops::SpaceToDepthGradKernel<paddle::platform::CPUDeviceContext, float>,
    ops::SpaceToDepthGradKernel<paddle::platform::CPUDeviceContext, double>,
    ops::SpaceToDepthGradKernel<paddle::platform::CPUDeviceContext, int64_t>);