【PaddlePaddle Hackathon 2】24、为 Paddle 新增 nn.ChannelShuffle 组网 API

paddle/fluid/operators/channel_shuffle_op.cc

@@ -1,21 +1,22 @@
-/* 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. */
+// Copyright (c) 2021 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 {
@@ -62,25 +63,6 @@ class ChannelShuffleOpMaker : public framework::OpProtoAndCheckerMaker {
 class ChannelShuffleGradOp : 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 dx_dims = do_dims;
-    ctx->SetOutputDim(framework::GradVarName("X"), dx_dims);
-  }
 };
 
 template <typename T>
@@ -110,4 +92,9 @@ REGISTER_OPERATOR(channel_shuffle, ops::ChannelShuffleOp,
                   ops::ChannelShuffleGradOpMaker<paddle::imperative::OpBase>,
                   ChannelShuffleInferShapeFunctor);
 
-REGISTER_OPERATOR(channel_shuffle_grad, ops::ChannelShuffleGradOp);
+DECLARE_INFER_SHAPE_FUNCTOR(channel_shuffle_grad,
+                            ChannelShuffleGradInferShapeFunctor,
+                            PD_INFER_META(phi::ChannelShuffleGradInferMeta));
+
+REGISTER_OPERATOR(channel_shuffle_grad, ops::ChannelShuffleGradOp,
+                  ChannelShuffleGradInferShapeFunctor);

paddle/phi/infermeta/backward.cc

@@ -67,6 +67,22 @@ void BilinearTensorProductGradInferMeta(const MetaTensor& x,
   }
 }
 
+void ChannelShuffleGradInferMeta(const MetaTensor& out_grad,
+                                 int groups,
+                                 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()));
+  auto dx_dims = do_dims;
+  x_grad->set_dims(dx_dims);
+  x_grad->set_dtype(out_grad.dtype());
+}
+
 void ConvTransposeGradInferMeta(const MetaTensor& x,
                                 const MetaTensor& filter,
                                 const MetaTensor& dout,

paddle/phi/infermeta/backward.h

@@ -37,6 +37,11 @@ void BilinearTensorProductGradInferMeta(const MetaTensor& x,
                                         MetaTensor* dweight,
                                         MetaTensor* dbias);
 
+void ChannelShuffleGradInferMeta(const MetaTensor& out_grad,
+                                 int groups,
+                                 const std::string& data_format,
+                                 MetaTensor* x_grad);
+
 void ConvTransposeGradInferMeta(const MetaTensor& x,
                                 const MetaTensor& filter,
                                 const MetaTensor& dout,

paddle/phi/kernels/channel_shuffle_grad_kernel.h

@@ -20,7 +20,7 @@
 namespace phi {
 
 template <typename T, typename Context>
-void ChannelShuffleGradKernel(const Context& ctx,
+void ChannelShuffleGradKernel(const Context& dev_ctx,
                               const DenseTensor& out_grad,
                               int groups,
                               const std::string& data_format,

paddle/phi/kernels/channel_shuffle_kernel.h

@@ -20,7 +20,7 @@
 namespace phi {
 
 template <typename T, typename Context>
-void ChannelShuffleKernel(const Context& ctx,
+void ChannelShuffleKernel(const Context& dev_ctx,
                           const DenseTensor& x,
                           int groups,
                           const std::string& data_format,

paddle/phi/kernels/cpu/channel_shuffle_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/channel_shuffle_grad_kernel.h"
+#include "paddle/phi/kernels/impl/channel_shuffle_grad_kernel_impl.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(channel_shuffle_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::ChannelShuffleGradKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/cpu/channel_shuffle_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/channel_shuffle_kernel.h"
+#include "paddle/phi/kernels/impl/channel_shuffle_kernel_impl.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(channel_shuffle,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::ChannelShuffleKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/gpu/channel_shuffle_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/channel_shuffle_grad_kernel.h"
+#include "paddle/phi/kernels/impl/channel_shuffle_grad_kernel_impl.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(channel_shuffle_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::ChannelShuffleGradKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/gpu/channel_shuffle_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/channel_shuffle_kernel.h"
+#include "paddle/phi/kernels/impl/channel_shuffle_kernel_impl.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(channel_shuffle,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::ChannelShuffleKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/channel_shuffle_grad_kernel.cc → paddle/phi/kernels/impl/channel_shuffle_grad_kernel_impl.h

@@ -12,25 +12,25 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/phi/kernels/channel_shuffle_grad_kernel.h"
+#pragma once
+
 #include <string>
 #include <vector>
-#include "paddle/phi/backends/all_context.h"
+
 #include "paddle/phi/core/dense_tensor.h"
-#include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
 namespace phi {
 
 template <typename T, typename Context>
-void ChannelShuffleGradKernel(const Context& ctx,
+void ChannelShuffleGradKernel(const Context& dev_ctx,
                               const DenseTensor& out_grad,
                               int groups,
                               const std::string& data_format,
                               DenseTensor* x_grad) {
   auto* dout = &out_grad;
   auto* dx = x_grad;
-  ctx.template Alloc<T>(dx);
+  dev_ctx.template Alloc<T>(dx);
   bool channel_last = (data_format == "NHWC");
   auto do_dims = dout->dims();
   auto dx_dims = dx->dims();
@@ -51,24 +51,8 @@ void ChannelShuffleGradKernel(const Context& ctx,
     o.Resize({dx_dims[0], dx_dims[1], dx_dims[2], groups, dx_dims[3] / groups});
   }
   phi::funcs::Transpose<Context, T, 5> trans;
-  trans(ctx, t, &o, axis);
+  trans(dev_ctx, t, &o, axis);
   dx->Resize(dx_dims);
 }
 
 }  // namespace phi
-
-PD_REGISTER_KERNEL(channel_shuffle_grad,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::ChannelShuffleGradKernel,
-                   float,
-                   double) {}
-
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-PD_REGISTER_KERNEL(channel_shuffle_grad,
-                   GPU,
-                   ALL_LAYOUT,
-                   phi::ChannelShuffleGradKernel,
-                   float,
-                   double) {}
-#endif

paddle/phi/kernels/channel_shuffle_kernel.cc → paddle/phi/kernels/impl/channel_shuffle_kernel_impl.h

@@ -12,24 +12,24 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/phi/kernels/channel_shuffle_kernel.h"
+#pragma once
+
 #include <string>
 #include <vector>
-#include "paddle/phi/backends/all_context.h"
+
 #include "paddle/phi/core/dense_tensor.h"
-#include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
 namespace phi {
 
 template <typename T, typename Context>
-void ChannelShuffleKernel(const Context& ctx,
+void ChannelShuffleKernel(const Context& dev_ctx,
                           const DenseTensor& x,
                           int groups,
                           const std::string& data_format,
                           DenseTensor* out) {
   auto* in = &x;
-  ctx.template Alloc<T>(out);
+  dev_ctx.template Alloc<T>(out);
   bool channel_last = (data_format == "NHWC");
   auto in_dims = in->dims();
   auto o_dims = out->dims();
@@ -50,24 +50,8 @@ void ChannelShuffleKernel(const Context& ctx,
     o.Resize({in_dims[0], in_dims[1], in_dims[2], in_dims[3] / groups, groups});
   }
   phi::funcs::Transpose<Context, T, 5> trans;
-  trans(ctx, t, &o, axis);
+  trans(dev_ctx, t, &o, axis);
   out->Resize(o_dims);
 }
 
 }  // namespace phi
-
-PD_REGISTER_KERNEL(channel_shuffle,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::ChannelShuffleKernel,
-                   float,
-                   double) {}
-
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-PD_REGISTER_KERNEL(channel_shuffle,
-                   GPU,
-                   ALL_LAYOUT,
-                   phi::ChannelShuffleKernel,
-                   float,
-                   double) {}
-#endif

paddle/phi/ops/compat/channel_shuffle_sig.cc

@@ -16,12 +16,6 @@
 
 namespace phi {
 
-KernelSignature ChannelShuffleOpArgumentMapping(
-    const ArgumentMappingContext& ctx) {
-  return KernelSignature(
-      "channel_shuffle", {"X"}, {"groups", "data_format"}, {"Out"});
-}
-
 KernelSignature ChannelShuffleGradOpArgumentMapping(
     const ArgumentMappingContext& ctx) {
   return KernelSignature("channel_shuffle_grad",
@@ -32,7 +26,5 @@ KernelSignature ChannelShuffleGradOpArgumentMapping(
 
 }  // namespace phi
 
-PD_REGISTER_ARG_MAPPING_FN(channel_shuffle,
-                           phi::ChannelShuffleOpArgumentMapping);
 PD_REGISTER_ARG_MAPPING_FN(channel_shuffle_grad,
                            phi::ChannelShuffleGradOpArgumentMapping);