Dev maxpool series op 123d

oneflow/core/autograd/gradient_funcs/pooling.cpp

@@ -0,0 +1,120 @@
+/*
+Copyright 2020 The OneFlow 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 "oneflow/core/framework/attr_map.h"
+#include "oneflow/core/framework/op_expr_grad_function.h"
+#include "oneflow/core/framework/op_builder.h"
+#include "oneflow/core/framework/op_interpreter/op_interpreter_util.h"
+#include "oneflow/core/framework/op_expr.h"
+#include "oneflow/core/framework/op_expr_helper.h"
+#include "oneflow/core/functional/functional.h"
+
+namespace oneflow {
+namespace one {
+
+namespace {
+
+struct PoolingInterpState : public OpExprInterpState {
+  bool requires_grad;
+  size_t input_index;
+  size_t output_index;
+  size_t indice_index;
+
+  std::string data_format;
+  std::string padding;
+  std::vector<int32_t> padding_before;
+  std::vector<int32_t> padding_after;
+  std::vector<int32_t> kernel_size;
+  std::vector<int32_t> stride;
+  std::vector<int32_t> dilation;
+  bool return_indices;
+  bool ceil_mode;
+};
+
+class PoolingNdGrad : public OpExprGradFunction<PoolingInterpState> {
+ public:
+  virtual ~PoolingNdGrad() = default;
+  Maybe<void> Init(const OpExpr& op, const std::string& mode);
+  Maybe<void> Capture(PoolingInterpState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override;
+  Maybe<void> Apply(const PoolingInterpState* ctx, const TensorTuple& out_grads,
+                    TensorTuple* in_grads) const override;
+
+ private:
+  std::string mode_;
+  AttrMap base_attrs_;
+};
+
+Maybe<void> PoolingNdGrad::Init(const OpExpr& op, const std::string& mode) {
+  const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
+  CHECK_NOTNULL_OR_RETURN(fw_op_expr);
+  base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());
+  mode_ = mode;
+  return Maybe<void>::Ok();
+}
+
+Maybe<void> PoolingNdGrad::Capture(PoolingInterpState* ctx, const TensorTuple& inputs,
+                                   const TensorTuple& outputs, const AttrMap& attrs) const {
+  ctx->requires_grad = inputs.at(0)->requires_grad();
+  if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
+
+  ctx->input_index = ctx->SaveTensorForBackward(inputs.at(0));
+  ctx->output_index = ctx->SaveTensorForBackward(outputs.at(0));
+  ctx->indice_index = ctx->SaveTensorForBackward(outputs.at(1));
+
+  ComposedAttrMap composed_attrs(attrs, base_attrs_);
+  ctx->data_format = JUST(composed_attrs.GetAttr<std::string>("data_format"));
+  ctx->padding = JUST(composed_attrs.GetAttr<std::string>("padding"));
+  ctx->padding_before = JUST(composed_attrs.GetAttr<std::vector<int32_t>>("padding_before"));
+  ctx->padding_after = JUST(composed_attrs.GetAttr<std::vector<int32_t>>("padding_after"));
+  ctx->kernel_size = JUST(composed_attrs.GetAttr<std::vector<int32_t>>("kernel_size"));
+  ctx->stride = JUST(composed_attrs.GetAttr<std::vector<int32_t>>("stride"));
+  ctx->dilation = JUST(composed_attrs.GetAttr<std::vector<int32_t>>("dilation"));
+  ctx->return_indices = JUST(composed_attrs.GetAttr<bool>("return_indices"));
+  ctx->ceil_mode = JUST(composed_attrs.GetAttr<bool>("ceil_mode"));
+  return Maybe<void>::Ok();
+}
+
+Maybe<void> PoolingNdGrad::Apply(const PoolingInterpState* ctx, const TensorTuple& out_grads,
+                                 TensorTuple* in_grads) const {
+  if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
+  CHECK_LE_OR_RETURN(out_grads.size(), 2);
+
+  int32_t ndims = ctx->kernel_size.size();
+  const auto& input = ctx->SavedTensors().at(ctx->input_index);
+  const auto& output = ctx->SavedTensors().at(ctx->output_index);
+  const auto& indice = ctx->SavedTensors().at(ctx->indice_index);
+
+  in_grads->resize(1);
+  in_grads->at(0) = JUST(functional::PoolingNdGrad(
+      input, output, indice, out_grads.at(0), mode_, ndims, ctx->data_format, ctx->padding,
+      ctx->padding_before, ctx->padding_after, ctx->kernel_size, ctx->stride, ctx->dilation,
+      ctx->return_indices, ctx->ceil_mode));
+
+  return Maybe<void>::Ok();
+}
+
+}  // namespace
+
+class MaxpoolNdGrad final : public PoolingNdGrad {
+ public:
+  Maybe<void> Init(const OpExpr& op) override { return PoolingNdGrad::Init(op, "max"); }
+};
+
+REGISTER_OP_EXPR_GRAD_FUNCTION("maxpool_2d", MaxpoolNdGrad);
+REGISTER_OP_EXPR_GRAD_FUNCTION("maxpool_3d", MaxpoolNdGrad);
+
+}  // namespace one
+}  // namespace oneflow

oneflow/core/functional/functional_api.yaml

@@ -403,6 +403,29 @@
                        Int32List strides, Bool ceil_mode)"
   bind_python: False
 
+- name: "maxpool_2d"
+  signature:
+    "TensorTuple Maxpool2D(Tensor x, *, String data_format=\"channels_first\", String padding,
+                      Int32List padding_before, Int32List padding_after,
+                      Int32List kernel_size, Int32List stride, Int32List dilation, 
+                      Bool return_indices=True, Bool ceil_mode=False)"
+  bind_python: True
+
+- name: "maxpool_3d"
+  signature:
+    "TensorTuple Maxpool3D(Tensor x, *, String data_format=\"channels_first\", String padding,
+                      Int32List padding_before, Int32List padding_after,
+                      Int32List kernel_size, Int32List stride, Int32List dilation, 
+                      Bool return_indices=True, Bool ceil_mode=False)"
+  bind_python: True
+
+- name: "pooling_grad"
+  signature:
+    "Tensor PoolingNdGrad(Tensor x, Tensor y, Tensor indice, Tensor dy, *, String mode, Int32 ndims, String data_format,
+                       String padding, Int32List padding_before, Int32List padding_after, Int32List kernel_size,
+                       Int32List stride, Int32List dilation, Bool return_indices, Bool ceil_mode)"
+  bind_python: False
+
 - name: "prelu"
   signature: "Tensor PRelu(Tensor x, Tensor alpha)"
   bind_python: True

oneflow/core/functional/impl/nn_functor.cpp

@@ -227,6 +227,35 @@ class PoolNDFunctor {
   std::shared_ptr<OpExpr> op_;
 };
 
+class PoolingNDFunctor {
+ public:
+  PoolingNDFunctor() = default;
+  virtual ~PoolingNDFunctor() = default;
+  Maybe<TensorTuple> operator()(const std::shared_ptr<one::Tensor>& x,
+                                const std::string& data_format, const std::string& padding,
+                                const std::vector<int32_t>& padding_before,
+                                const std::vector<int32_t>& padding_after,
+                                const std::vector<int32_t>& kernel_size,
+                                const std::vector<int32_t>& stride,
+                                const std::vector<int32_t>& dilation, const bool& return_indices,
+                                const bool& ceil_mode) const {
+    MutableAttrMap attrs;
+    JUST(attrs.SetAttr<std::string>("padding", padding));
+    JUST(attrs.SetAttr<std::vector<int32_t>>("padding_before", padding_before));
+    JUST(attrs.SetAttr<std::vector<int32_t>>("padding_after", padding_after));
+    JUST(attrs.SetAttr<std::string>("data_format", data_format));
+    JUST(attrs.SetAttr<std::vector<int32_t>>("kernel_size", kernel_size));
+    JUST(attrs.SetAttr<std::vector<int32_t>>("stride", stride));
+    JUST(attrs.SetAttr<std::vector<int32_t>>("dilation", dilation));
+    JUST(attrs.SetAttr<bool>("return_indices", return_indices));
+    JUST(attrs.SetAttr<bool>("ceil_mode", ceil_mode));
+    return OpInterpUtil::Dispatch<TensorTuple>(*op_, {x}, attrs);
+  }
+
+ protected:
+  std::shared_ptr<OpExpr> op_;
+};
+
 class AvgPool2DFunctor : public PoolNDFunctor {
  public:
   AvgPool2DFunctor() {
@@ -241,6 +270,20 @@ class MaxPool2DFunctor : public PoolNDFunctor {
   }
 };
 
+class Maxpool2DFunctor : public PoolingNDFunctor {
+ public:
+  Maxpool2DFunctor() {
+    op_ = CHECK_JUST(one::OpBuilder("maxpool_2d").Input("x").Output("y").Output("indice").Build());
+  }
+};
+
+class Maxpool3DFunctor : public PoolingNDFunctor {
+ public:
+  Maxpool3DFunctor() {
+    op_ = CHECK_JUST(one::OpBuilder("maxpool_3d").Input("x").Output("y").Output("indice").Build());
+  }
+};
+
 class SparseSoftmaxCrossEntropyFunctor {
  public:
   SparseSoftmaxCrossEntropyFunctor() {
@@ -420,6 +463,8 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<impl::LayerNormFunctor>("LayerNorm");
   m.add_functor<impl::LayerNormAffineFunctor>("LayerNormAffine");
   m.add_functor<impl::AvgPool2DFunctor>("AvgPool2D");
+  m.add_functor<impl::Maxpool2DFunctor>("Maxpool2D");
+  m.add_functor<impl::Maxpool3DFunctor>("Maxpool3D");
   m.add_functor<impl::MaxPool2DFunctor>("MaxPool2D");
   m.add_functor<impl::SparseSoftmaxCrossEntropyFunctor>("SparseSoftmaxCrossEntropy");
   m.add_functor<impl::SmoothL1LossFunctor>("SmoothL1Loss");

oneflow/core/functional/impl/nn_grad_functor.cpp

@@ -109,6 +109,57 @@ class ConvDataGradFunctor {
   std::shared_ptr<OpExpr> op_;
 };
 
+class PoolingNdGradFunctor {
+ public:
+  PoolingNdGradFunctor() {
+    for (const auto& mode : {"max"}) {
+      for (int ndims = 2; ndims <= 3; ++ndims) {
+        const auto& op_type_name = GetOpTypeName(mode, ndims);
+        op_expr_map_[op_type_name] = CHECK_JUST(one::OpBuilder(op_type_name)
+                                                    .Input("x")
+                                                    .Input("y")
+                                                    .Input("indice")
+                                                    .Input("dy")
+                                                    .Output("dx")
+                                                    .Build());
+      }
+    }
+  }
+  static std::string GetOpTypeName(const std::string& mode, const int32_t& ndims) {
+    return mode + "pool_" + std::to_string(ndims) + "d_grad";
+  }
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x,
+                           const std::shared_ptr<one::Tensor>& y,
+                           const std::shared_ptr<one::Tensor>& indice,
+                           const std::shared_ptr<one::Tensor>& dy, const std::string& mode,
+                           const int32_t& ndims, const std::string& data_format,
+                           const std::string& padding, const std::vector<int32_t>& padding_before,
+                           const std::vector<int32_t>& padding_after,
+                           const std::vector<int32_t>& kernel_size,
+                           const std::vector<int32_t>& stride, const std::vector<int32_t>& dilation,
+                           const bool& return_indices, const bool& ceil_mode) const {
+    MutableAttrMap attrs;
+    JUST(attrs.SetAttr<std::string>("padding", padding));
+    JUST(attrs.SetAttr<std::vector<int32_t>>("padding_before", padding_before));
+    JUST(attrs.SetAttr<std::vector<int32_t>>("padding_after", padding_after));
+    JUST(attrs.SetAttr<std::string>("data_format", data_format));
+    JUST(attrs.SetAttr<std::vector<int32_t>>("kernel_size", kernel_size));
+    JUST(attrs.SetAttr<std::vector<int32_t>>("stride", stride));
+    JUST(attrs.SetAttr<std::vector<int32_t>>("dilation", dilation));
+    JUST(attrs.SetAttr<bool>("return_indices", return_indices));
+    JUST(attrs.SetAttr<bool>("ceil_mode", ceil_mode));
+    const auto& op_type_name = GetOpTypeName(mode, ndims);
+    const auto& it = op_expr_map_.find(op_type_name);
+    CHECK_OR_RETURN(it != op_expr_map_.end())
+        << "Encounter unsupported op " << op_type_name << " in PoolingNdGradFunctor.";
+    CHECK_NOTNULL_OR_RETURN(it->second);
+    return OpInterpUtil::Dispatch<Tensor>(*it->second, {x, y, indice, dy}, attrs);
+  }
+
+ protected:
+  std::unordered_map<std::string, std::shared_ptr<OpExpr>> op_expr_map_;
+};
+
 class PoolNdGradFunctor {
  public:
   PoolNdGradFunctor() {
@@ -225,6 +276,7 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<impl::ConvDataGradFunctor>("ConvDataGrad");
   m.add_functor<impl::PoolNdGradFunctor>("PoolNdGrad");
   m.add_functor<impl::SmoothL1LossGradFunctor>("SmoothL1LossGrad");
+  m.add_functor<impl::PoolingNdGradFunctor>("PoolingNdGrad");
   m.add_functor<impl::PadGradFunctor>("PadGrad");
 };