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

paddle/fluid/operators/pixel_unshuffle_op.cc

@@ -0,0 +1,103 @@
+// 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/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/backward.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 decrease spatial resolution by.")
+        .SetDefault(1);
+    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, W)`
+    		to a tensor of shape :math:`(*, C\times r^2, H / r, W / r)`.
+
+		This operation is the reversion of PixelShuffle operation.
+
+		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 PixelUnshuffleGradOpMaker : public framework::SingleGradOpMaker<T> {
+ public:
+  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
+
+ protected:
+  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;
+};
+
+}  // 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::PixelUnshuffleGradOpMaker<paddle::framework::OpDesc>,
+                  ops::PixelUnshuffleGradOpMaker<paddle::imperative::OpBase>,
+                  PixelUnshuffleInferShapeFunctor);
+
+DECLARE_INFER_SHAPE_FUNCTOR(pixel_unshuffle_grad,
+                            PixelUnshuffleGradInferShapeFunctor,
+                            PD_INFER_META(phi::PixelUnshuffleGradInferMeta));
+
+REGISTER_OPERATOR(pixel_unshuffle_grad, ops::PixelUnshuffleGradOp,
+                  PixelUnshuffleGradInferShapeFunctor);

paddle/phi/infermeta/backward.cc

@@ -427,6 +427,36 @@ void NllLossGradInferMeta(const MetaTensor& x,
   }
 }
 
+void PixelUnshuffleGradInferMeta(const MetaTensor& out_grad,
+                                 int downscale_factor,
+                                 const std::string& data_format,
+                                 MetaTensor* x_grad) {
+  auto do_dims = out_grad.dims();
+  PADDLE_ENFORCE_EQ(do_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.",
+                        do_dims.size()));
+
+  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);
+  }
+  x_grad->set_dims(dx_dims);
+  x_grad->set_dtype(out_grad.dtype());
+}
+
 void PoolGradInferMeta(const MetaTensor& x,
                        const MetaTensor& out,
                        const MetaTensor& dout,

paddle/phi/infermeta/backward.h

@@ -173,6 +173,11 @@ void NllLossGradInferMeta(const MetaTensor& input,
                           MetaTensor* intput_grad,
                           MetaConfig config = MetaConfig());
 
+void PixelUnshuffleGradInferMeta(const MetaTensor& out_grad,
+                                 int downscale_factor,
+                                 const std::string& data_format,
+                                 MetaTensor* x_grad);
+
 void PsroiPoolGradInferMeta(const MetaTensor& x,
                             const MetaTensor& rois,
                             paddle::optional<const MetaTensor&> rois_num,

paddle/phi/infermeta/unary.cc

@@ -1423,6 +1423,66 @@ void PixelShuffleGradInferMeta(const MetaTensor& out_grad,
   x_grad->set_dtype(out_grad.dtype());
 }
 
+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()));
+  PADDLE_ENFORCE_GE(downscale_factor,
+                    1,
+                    phi::errors::InvalidArgument(
+                        "downscale_factor should be larger than 0."));
+  PADDLE_ENFORCE_EQ(data_format == "NCHW" || data_format == "NHWC",
+                    true,
+                    phi::errors::InvalidArgument(
+                        "data_format must be one of "
+                        "NCHW and NHWC. But recevied data_format: %s",
+                        data_format));
+
+  const bool channel_last = (data_format == "NHWC");
+
+  if (!channel_last) {
+    PADDLE_ENFORCE_EQ(
+        (input_dims[2] % downscale_factor) == 0 &&
+            (input_dims[3] % downscale_factor) == 0,
+        true,
+        phi::errors::InvalidArgument("Downscale factor[%u] should divide both "
+                                     "height[%u] and width[%u]",
+                                     downscale_factor,
+                                     input_dims[2],
+                                     input_dims[3]));
+  } else {
+    PADDLE_ENFORCE_EQ(
+        (input_dims[1] % downscale_factor) == 0 &&
+            (input_dims[2] % downscale_factor) == 0,
+        true,
+        phi::errors::InvalidArgument("Downscale factor[%u] should divide both "
+                                     "height[%u] and width[%u]",
+                                     downscale_factor,
+                                     input_dims[1],
+                                     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

@@ -209,6 +209,11 @@ void PixelShuffleGradInferMeta(const MetaTensor& out_grad,
                                const std::string& data_format,
                                MetaTensor* x_grad);
 
+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 <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& dev_ctx,
+                              const DenseTensor& out_grad,
+                              int downscale_factor,
+                              const std::string& data_format,
+                              DenseTensor* x_grad) {
+  auto* dout = &out_grad;
+  auto* dx = x_grad;
+  dev_ctx.template Alloc<T>(dx);
+  int factor = downscale_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(dev_ctx, t, &o, axis);
+  dx->Resize(dx_dims);
+}
+
+}  // namespace phi