space_to_depth_op.cc 5.2 KB
Newer Older
J
JiabinYang 已提交
1
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
J
JiabinYang 已提交
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. */

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

    auto x_dims = ctx->GetInputDim("X");
    PADDLE_ENFORCE_EQ(x_dims.size(), 4, "input should be a 4D tensor");
J
JiabinYang 已提交
34
    auto blocksize = ctx->Attrs().Get<int64_t>("blocksize");
J
JiabinYang 已提交
35

J
JiabinYang 已提交
36
    PADDLE_ENFORCE_GT(blocksize, 1, "The blocksize 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
    PADDLE_ENFORCE_EQ(x_dims[1] % (blocksize * blocksize), 0,
J
JiabinYang 已提交
42
                      "input channel should be divisible of the square of "
J
JiabinYang 已提交
43 44
                      "SpaceToDepthOp blocksize");
    PADDLE_ENFORCE_EQ(x_dims[2] % (blocksize), 0,
J
JiabinYang 已提交
45
                      "input Height should be divisible of the square of "
J
JiabinYang 已提交
46 47
                      "SpaceToDepthOp blocksize");
    PADDLE_ENFORCE_EQ(x_dims[3] % (blocksize), 0,
J
JiabinYang 已提交
48
                      "input Width should be divisible of the square of "
J
JiabinYang 已提交
49
                      "SpaceToDepthOp blocksize");
J
JiabinYang 已提交
50

J
JiabinYang 已提交
51
    VLOG(3) << "SpaceToDepthOp operator x.shape=" << x_dims
J
JiabinYang 已提交
52
            << "Attribute blocksize" << blocksize << std::endl;
J
JiabinYang 已提交
53 54 55

    std::vector<int64_t> output_shape(4, 0);  // [B,C,H,W]
    output_shape[0] = x_dims[0];
J
JiabinYang 已提交
56 57 58
    output_shape[1] = x_dims[1] * blocksize * blocksize;
    output_shape[2] = x_dims[2] / blocksize;
    output_shape[3] = x_dims[3] / blocksize;
J
JiabinYang 已提交
59 60 61 62 63 64 65 66 67 68 69 70 71

    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
              "SpaceToDepthOp operator.");
    AddAttr<int64_t>(
J
JiabinYang 已提交
83 84
        "blocksize",
        "(int64_t, default 2) blocksize used to do change Space To Depth.")
J
JiabinYang 已提交
85 86
        .SetDefault(2)
        .GreaterThan(1);
J
JiabinYang 已提交
87 88
    AddComment(R"DOC(
        reorg operator used in Yolo v2.
J
JiabinYang 已提交
89
        The equation is: C2 = C1/blocksize * blocksize, W2 = W1 ∗ blocksize + offset % blocksize, H2 = H1 ∗ blocksize + offset / blocksize, 
J
JiabinYang 已提交
90

J
JiabinYang 已提交
91
        Reshape Input(X) into the shape according to Attr(blocksize). The
J
JiabinYang 已提交
92 93 94 95
        data in Input(X) are unchanged.

        Examples:

J
JiabinYang 已提交
96
            1. Given a 4-D tensor Input(X) with a shape [128, 2048, 26, 26], and the blocksize is 2, the reorg operator will transform Input(X)
J
JiabinYang 已提交
97
            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>);