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

paddle/fluid/operators/channel_shuffle_op.cc

@@ -0,0 +1,113 @@
+/* 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/unary.h"
+
+namespace paddle {
+namespace operators {
+
+class ChannelShuffleOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+};
+
+class ChannelShuffleOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("X",
+             "(Tensor, default Tensor<float>), "
+             "the input feature data of ChannelShuffleOp, the layout is "
+             "[N, C, H, W] or [N, H, W, C].");
+    AddOutput("Out",
+              "(Tensor, default Tensor<float>), the output of "
+              "ChannelShuffleOp. The layout is also [N, C, "
+              "H, W] or [N, H, W, C].");
+    AddAttr<int>("groups", "number of groups to divide channels in.");
+    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(
+    Channel Shuffle operator
+    This operator divides channels in a tensor of shape :math:`(*, C, H, W)`
+        into :math:`g` groups and rearranges them as :math:`(*, C/g, g, H, W)`
+        while keeping the original tensor shape.
+
+    Please refer to the paper:
+        `ShuffleNet: An Extremely Efficient Convolutional Neural Network for 
+        Mobile Devices <https://arxiv.org/abs/1707.01083>`_
+        by Zhang et. al (2017) for more details. 
+
+        )DOC");
+  }
+};
+
+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>
+class ChannelShuffleGradOpMaker : public framework::SingleGradOpMaker<T> {
+ public:
+  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
+
+ protected:
+  void Apply(GradOpPtr<T> op) const override {
+    op->SetType("channel_shuffle_grad");
+    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
+    op->SetAttrMap(this->Attrs());
+    op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+DECLARE_INFER_SHAPE_FUNCTOR(channel_shuffle, ChannelShuffleInferShapeFunctor,
+                            PD_INFER_META(phi::ChannelShuffleInferMeta));
+
+REGISTER_OPERATOR(channel_shuffle, ops::ChannelShuffleOp,
+                  ops::ChannelShuffleOpMaker,
+                  ops::ChannelShuffleGradOpMaker<paddle::framework::OpDesc>,
+                  ops::ChannelShuffleGradOpMaker<paddle::imperative::OpBase>,
+                  ChannelShuffleInferShapeFunctor);
+
+REGISTER_OPERATOR(channel_shuffle_grad, ops::ChannelShuffleGradOp);

paddle/phi/infermeta/unary.cc

@@ -2863,6 +2863,52 @@ void WhereIndexInferMeta(const MetaTensor& condition, MetaTensor* out) {
   out->set_dtype(DataType::INT64);
 }
 
+void ChannelShuffleInferMeta(const MetaTensor& x,
+                             int groups,
+                             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(
+      groups,
+      1,
+      phi::errors::InvalidArgument("groups 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[1] % groups,
+                      0,
+                      phi::errors::InvalidArgument(
+                          "The number of groups to divide channels in [%u] "
+                          "should divide the number of channel [%u]",
+                          groups,
+                          input_dims[1]));
+  } else {
+    PADDLE_ENFORCE_EQ(input_dims[3] % groups,
+                      0,
+                      phi::errors::InvalidArgument(
+                          "The number of groups to divide channels in [%u] "
+                          "should divide the number of channel [%u]",
+                          groups,
+                          input_dims[3]));
+  }
+  auto output_dims = input_dims;
+  out->set_dtype(x.dtype());
+  out->set_dims(output_dims);
+}
+
 }  // namespace phi
 
 PD_REGISTER_INFER_META_FN(copy_to, phi::CopyToInferMeta);

paddle/phi/infermeta/unary.h

@@ -426,4 +426,9 @@ void OneHotInferMeta(const MetaTensor& x, const Scalar& depth, MetaTensor* out);
 
 void WhereIndexInferMeta(const MetaTensor& condition, MetaTensor* out);
 
+void ChannelShuffleInferMeta(const MetaTensor& x,
+                             int groups,
+                             const std::string& data_format,
+                             MetaTensor* out);
+
 }  // namespace phi

paddle/phi/kernels/channel_shuffle_grad_kernel.cc

@@ -0,0 +1,74 @@
+// 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 <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,
+                              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);
+  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], do_dims[1] / groups, groups, do_dims[2], do_dims[3]});
+  } else {
+    t.Resize({do_dims[0], do_dims[1], do_dims[2], do_dims[3] / groups, groups});
+  }
+  auto axis = !channel_last ? std::vector<int>{0, 2, 1, 3, 4}
+                            : std::vector<int>{0, 1, 2, 4, 3};
+
+  DenseTensor o(*dx);
+  if (!channel_last) {
+    o.Resize({dx_dims[0], groups, dx_dims[1] / groups, dx_dims[2], dx_dims[3]});
+  } else {
+    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);
+  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_grad_kernel.h

@@ -0,0 +1,29 @@
+// 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 "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void ChannelShuffleGradKernel(const Context& ctx,
+                              const DenseTensor& out_grad,
+                              int groups,
+                              const std::string& data_format,
+                              DenseTensor* x_grad);
+
+}  // namespace phi

paddle/phi/kernels/channel_shuffle_kernel.cc

@@ -0,0 +1,73 @@
+// 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 <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,
+                          const DenseTensor& x,
+                          int groups,
+                          const std::string& data_format,
+                          DenseTensor* out) {
+  auto* in = &x;
+  ctx.template Alloc<T>(out);
+  bool channel_last = (data_format == "NHWC");
+  auto in_dims = in->dims();
+  auto o_dims = out->dims();
+
+  DenseTensor t(*in);
+  if (!channel_last) {
+    t.Resize({in_dims[0], groups, in_dims[1] / groups, in_dims[2], in_dims[3]});
+  } else {
+    t.Resize({in_dims[0], in_dims[1], in_dims[2], groups, in_dims[3] / groups});
+  }
+  auto axis = !channel_last ? std::vector<int>{0, 2, 1, 3, 4}
+                            : std::vector<int>{0, 1, 2, 4, 3};
+
+  DenseTensor o(*out);
+  if (!channel_last) {
+    o.Resize({in_dims[0], in_dims[1] / groups, groups, in_dims[2], in_dims[3]});
+  } else {
+    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);
+  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/kernels/channel_shuffle_kernel.h

@@ -0,0 +1,29 @@
+// 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 "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void ChannelShuffleKernel(const Context& ctx,
+                          const DenseTensor& x,
+                          int groups,
+                          const std::string& data_format,
+                          DenseTensor* out);
+
+}  // namespace phi