【PaddlePaddle Hackathon 2】29 新增 PixelUnshuffle 组网 API

paddle/fluid/operators/pixel_unshuffle_op.cc

@@ -0,0 +1,142 @@
+/*Copyright (c) 2019 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. */
+
+#include <memory>
+#include "paddle/fluid/framework/infershape_utils.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/op_version_registry.h"
+#include "paddle/phi/core/infermeta_utils.h"
+#include "paddle/phi/infermeta/unary.h"
+
+namespace paddle {
+namespace operators {
+
+class PixelUnshuffleOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+};
+
+class PixelUnshuffleOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("X",
+             "(Tensor, default Tensor<float>), "
+             "the input feature data of PixelUnshuffleOp, the layout is [N, C, "
+             "H, W] or [N, H, W, C].");
+    AddOutput("Out",
+              "(Tensor, default Tensor<float>), the output of "
+              "PixelUnshuffleOp. The layout is [N, C*factor^2, H/factor, "
+              "W/factor] or [N, H/factor, W/factor, C*factor^2].");
+    AddAttr<int>("downscale_factor",
+                 "the factor to descrease spatial resolution by.")
+        .SetDefault(1)
+        .AddCustomChecker([](const int& downscale_factor) {
+          PADDLE_ENFORCE_GE(downscale_factor, 1,
+                            platform::errors::InvalidArgument(
+                                "downscale_factor should be larger than 0."));
+        });
+    AddAttr<std::string>(
+        "data_format",
+        "An optional string from: \"NHWC\", \"NCHW\". "
+        "Defaults to \"NHWC\", Specify the data format of the input data.")
+        .SetDefault("NCHW");
+
+    AddComment(R"DOC(
+		Pixel Unshuffle operator
+		This operator rearranges elements in a tensor of shape :math:`(*, C , H \times r, W \times r)`
+    		to a tensor of shape :math:`( C \times r^2, H, W)`, where r is downscale factor.
+
+		This is taken as the inverse operator of Pixel Shuffle .
+
+		Please refer to the paper:
+		 `Real-Time Single Image and Video Super-Resolution Using an Efficient 
+		 Sub-Pixel Convolutional Neural Network <https://arxiv.org/abs/1609.05158v2>`_
+    		by Shi et. al (2016) for more details. 
+
+        )DOC");
+  }
+};
+
+template <typename T>
+class PixelUnshuffleGradMaker : public framework::SingleGradOpMaker<T> {
+ public:
+  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
+
+  void Apply(GradOpPtr<T> op) const override {
+    op->SetType("pixel_unshuffle_grad");
+    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
+    op->SetAttrMap(this->Attrs());
+    op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
+  }
+};
+
+class PixelUnshuffleGradOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+  void InferShape(framework::InferShapeContext* ctx) const override {
+    PADDLE_ENFORCE_EQ(
+        ctx->HasInput(framework::GradVarName("Out")), true,
+        platform::errors::NotFound("Input(Out@Grad) should not be null"));
+    PADDLE_ENFORCE_EQ(
+        ctx->HasOutput(framework::GradVarName("X")), true,
+        platform::errors::NotFound("Output(X@Grad) should not be null"));
+
+    auto do_dims = ctx->GetInputDim(framework::GradVarName("Out"));
+    PADDLE_ENFORCE_EQ(do_dims.size(), 4,
+                      platform::errors::InvalidArgument(
+                          "Input should be a 4-D tensor of format [N, C, H, W] "
+                          "or [N, H, W, C], but got %u.",
+                          do_dims.size()));
+
+    auto downscale_factor = ctx->Attrs().Get<int>("downscale_factor");
+
+    const std::string data_format =
+        ctx->Attrs().Get<std::string>("data_format");
+    const bool channel_last = (data_format == "NHWC");
+
+    auto dx_dims = do_dims;
+    dx_dims[0] = do_dims[0];
+
+    if (!channel_last) {
+      dx_dims[1] = do_dims[1] / (downscale_factor * downscale_factor);
+      dx_dims[2] = do_dims[2] * downscale_factor;
+      dx_dims[3] = do_dims[3] * downscale_factor;
+    } else {
+      dx_dims[1] = do_dims[1] * downscale_factor;
+      dx_dims[2] = do_dims[2] * downscale_factor;
+      dx_dims[3] = do_dims[3] / (downscale_factor * downscale_factor);
+    }
+    ctx->SetOutputDim(framework::GradVarName("X"), dx_dims);
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+DECLARE_INFER_SHAPE_FUNCTOR(pixel_unshuffle, PixelUnshuffleInferShapeFunctor,
+                            PD_INFER_META(phi::PixelUnshuffleInferMeta));
+
+REGISTER_OPERATOR(pixel_unshuffle, ops::PixelUnshuffleOp,
+                  ops::PixelUnshuffleOpMaker,
+                  ops::PixelUnshuffleGradMaker<paddle::framework::OpDesc>,
+                  ops::PixelUnshuffleGradMaker<paddle::imperative::OpBase>,
+                  PixelUnshuffleInferShapeFunctor);
+
+REGISTER_OPERATOR(pixel_unshuffle_grad, ops::PixelUnshuffleGradOp);
+
+REGISTER_OP_VERSION(pixel_unshuffle)
+    .AddCheckpoint(
+        R"ROC(
+               Compatible upgrade of pixel_unshuffle, add a new attribute [data_format])ROC",
+        paddle::framework::compatible::OpVersionDesc().NewAttr(
+            "data_format", "Specify the data format of the input data", true));

paddle/phi/infermeta/unary.cc

@@ -1204,6 +1204,71 @@ void PixelShuffleInferMeta(const MetaTensor& x,
   out->set_dims(output_dims);
 }
 
+void PixelUnshuffleInferMeta(const MetaTensor& x,
+                             int downscale_factor,
+                             const std::string& data_format,
+                             MetaTensor* out) {
+  auto input_dims = x.dims();
+  PADDLE_ENFORCE_EQ(input_dims.size(),
+                    4,
+                    phi::errors::InvalidArgument(
+                        "Input should be a 4-D tensor of format [N, C, H, W] "
+                        "or [N, H, W, C], but got %u.",
+                        input_dims.size()));
+
+  const bool channel_last = (data_format == "NHWC");
+
+  if (!channel_last) {
+    PADDLE_ENFORCE_EQ(
+        input_dims[2] % downscale_factor,
+        0,
+        phi::errors::InvalidArgument(
+            "The square of downscale_factor[%u] should divide the "
+            "height[%u]",
+            downscale_factor,
+            input_dims[2]));
+    PADDLE_ENFORCE_EQ(
+        input_dims[3] % downscale_factor,
+        0,
+        phi::errors::InvalidArgument(
+            "The square of downscale_factor[%u] should divide the "
+            "height[%u]",
+            downscale_factor,
+            input_dims[3]));
+  } else {
+    PADDLE_ENFORCE_EQ(
+        input_dims[1] % downscale_factor,
+        0,
+        phi::errors::InvalidArgument(
+            "The square of downscale_factor[%u] should divide the "
+            "width[%u]",
+            downscale_factor,
+            input_dims[1]));
+    PADDLE_ENFORCE_EQ(
+        input_dims[2] % downscale_factor,
+        0,
+        phi::errors::InvalidArgument(
+            "The square of downscale_factor[%u] should divide the "
+            "width[%u]",
+            downscale_factor,
+            input_dims[2]));
+  }
+
+  auto output_dims = input_dims;
+  output_dims[0] = input_dims[0];
+  if (!channel_last) {
+    output_dims[1] = input_dims[1] * (downscale_factor * downscale_factor);
+    output_dims[2] = input_dims[2] / downscale_factor;
+    output_dims[3] = input_dims[3] / downscale_factor;
+  } else {
+    output_dims[1] = input_dims[1] / downscale_factor;
+    output_dims[2] = input_dims[2] / downscale_factor;
+    output_dims[3] = input_dims[3] * (downscale_factor * downscale_factor);
+  }
+  out->set_dtype(x.dtype());
+  out->set_dims(output_dims);
+}
+
 void PNormInferMeta(const MetaTensor& x,
                     float porder,
                     int axis,

paddle/phi/infermeta/unary.h

@@ -191,6 +191,11 @@ void PixelShuffleInferMeta(const MetaTensor& x,
                            const std::string& data_format,
                            MetaTensor* out);
 
+void PixelUnshuffleInferMeta(const MetaTensor& x,
+                             int downscale_factor,
+                             const std::string& data_format,
+                             MetaTensor* out);
+
 void PNormInferMeta(const MetaTensor& x,
                     float porder,
                     int axis,

paddle/phi/kernels/cpu/pixel_unshuffle_grad_kernel.cc

@@ -0,0 +1,26 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/kernels/pixel_unshuffle_grad_kernel.h"
+#include "paddle/phi/kernels/impl/pixel_unshuffle_grad_kernel_impl.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(pixel_unshuffle_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::PixelUnshuffleGradKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/cpu/pixel_unshuffle_kernel.cc

@@ -0,0 +1,26 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/kernels/pixel_unshuffle_kernel.h"
+#include "paddle/phi/kernels/impl/pixel_unshuffle_kernel_impl.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(pixel_unshuffle,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::PixelUnshuffleKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/gpu/pixel_unshuffle_grad_kernel.cu

@@ -0,0 +1,26 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/kernels/impl/pixel_unshuffle_grad_kernel_impl.h"
+#include "paddle/phi/kernels/pixel_unshuffle_grad_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(pixel_unshuffle_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::PixelUnshuffleGradKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/gpu/pixel_unshuffle_kernel.cu

@@ -0,0 +1,26 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/kernels/impl/pixel_unshuffle_kernel_impl.h"
+#include "paddle/phi/kernels/pixel_unshuffle_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(pixel_unshuffle,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::PixelUnshuffleKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/impl/pixel_unshuffle_grad_kernel_impl.h

@@ -0,0 +1,58 @@
+// Copyright (c) 2022 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.
+
+#pragma once
+#include <algorithm>
+#include <string>
+#include <vector>
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void PixelUnshuffleGradKernel(const Context& ctx,
+                              const DenseTensor& out_grad,
+                              int upscale_factor,
+                              const std::string& data_format,
+                              DenseTensor* x_grad) {
+  auto* dout = &out_grad;
+  auto* dx = x_grad;
+  ctx.template Alloc<T>(dx);
+  int factor = upscale_factor;
+  bool channel_last = (data_format == "NHWC");
+  auto do_dims = dout->dims();
+  auto dx_dims = dx->dims();
+
+  DenseTensor t(*dout);
+  if (!channel_last) {
+    t.Resize({do_dims[0], dx_dims[1], factor, factor, do_dims[2], do_dims[3]});
+  } else {
+    t.Resize({do_dims[0], do_dims[1], do_dims[2], dx_dims[3], factor, factor});
+  }
+
+  std::vector<int> axis = {0, 1, 4, 2, 5, 3};
+
+  DenseTensor o(*dx);
+  if (!channel_last) {
+    o.Resize({do_dims[0], dx_dims[1], do_dims[2], factor, do_dims[3], factor});
+  } else {
+    o.Resize({do_dims[0], do_dims[1], factor, do_dims[2], factor, dx_dims[3]});
+  }
+  phi::funcs::Transpose<Context, T, 6> trans;
+  trans(ctx, t, &o, axis);
+  dx->Resize(dx_dims);
+}
+
+}  // namespace phi