Add nn.AdaptiveAvgPool1d and nn.AdaptiveAvgPool3d

docs/source/experimental.rst

@@ -180,7 +180,12 @@ Experimental features
 .. autofunction:: oneflow.experimental.nn.AvgPool1d
 .. autofunction:: oneflow.experimental.nn.AvgPool2d
 .. autofunction:: oneflow.experimental.nn.AvgPool3d
+.. autofunction:: oneflow.experimental.nn.AdaptiveAvgPool1d
 .. autofunction:: oneflow.experimental.nn.AdaptiveAvgPool2d
+.. autofunction:: oneflow.experimental.nn.AdaptiveAvgPool3d
+.. autofunction:: oneflow.experimental.adaptive_avg_pool1d
+.. autofunction:: oneflow.experimental.adaptive_avg_pool2d
+.. autofunction:: oneflow.experimental.adaptive_avg_pool3d
 .. autofunction:: oneflow.experimental.nn.MaxPool1d
 .. autofunction:: oneflow.experimental.nn.MaxPool2d
 .. autofunction:: oneflow.experimental.nn.MaxPool3d

oneflow/core/autograd/gradient_funcs/adaptive_pool.cpp

@@ -18,6 +18,7 @@ limitations under the License.
 #include "oneflow/core/framework/op_expr.h"
 #include "oneflow/core/framework/op_expr_helper.h"
 #include "oneflow/core/framework/op_interpreter/op_interpreter_util.h"
+#include "oneflow/core/functional/functional.h"
 
 namespace oneflow {
 namespace one {
@@ -26,49 +27,67 @@ struct AdaptivePoolInterpState : public OpExprInterpState {
   bool requires_grad;
 };
 
-class AdaptivePool : public OpExprGradFunction<AdaptivePoolInterpState> {
+class AdaptivePoolNdGrad : public OpExprGradFunction<AdaptivePoolInterpState> {
  public:
-  Maybe<void> Init(const OpExpr& op) override;
+  Maybe<void> Init(const OpExpr& op, std::string mode, const int& ndims);
   Maybe<void> Capture(AdaptivePoolInterpState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override;
   Maybe<void> Apply(const AdaptivePoolInterpState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override;
 
  private:
   AttrMap base_attrs_;
-  std::shared_ptr<OpExpr> grad_op_;
+  std::string mode_;
+  int32_t ndims_;
 };
 
-Maybe<void> AdaptivePool::Init(const OpExpr& op) {
+Maybe<void> AdaptivePoolNdGrad::Init(const OpExpr& op, std::string mode, const int& ndims) {
   const UserOpExpr* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
   CHECK_NOTNULL_OR_RETURN(fw_op_expr);
   base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());
-  const std::string& op_name = fw_op_expr->op_name();
-  grad_op_ = JUST(op_expr_helper::AdaptivePoolGradOp(GradientOpName(op_name)));
+  mode_ = mode;
+  ndims_ = ndims;
   return Maybe<void>::Ok();
 }
 
-Maybe<void> AdaptivePool::Capture(AdaptivePoolInterpState* ctx, const TensorTuple& inputs,
-                                  const TensorTuple& outputs, const AttrMap& attrs) const {
+Maybe<void> AdaptivePoolNdGrad::Capture(AdaptivePoolInterpState* 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->SaveTensorForBackward(inputs.at(0));
   return Maybe<void>::Ok();
 }
 
-Maybe<void> AdaptivePool::Apply(const AdaptivePoolInterpState* ctx, const TensorTuple& out_grads,
-                                TensorTuple* in_grads) const {
+Maybe<void> AdaptivePoolNdGrad::Apply(const AdaptivePoolInterpState* ctx,
+                                      const TensorTuple& out_grads, TensorTuple* in_grads) const {
   if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
   CHECK_EQ_OR_RETURN(out_grads.size(), 1);
 
   const std::shared_ptr<oneflow::one::Tensor>& x = ctx->SavedTensors().at(0);
   in_grads->resize(1);
-  in_grads->at(0) = JUST(OpInterpUtil::Dispatch<Tensor>(*grad_op_, {x, out_grads.at(0)}));
+  in_grads->at(0) = JUST(functional::AdaptivePoolNdGrad(x, out_grads.at(0), mode_, ndims_));
   return Maybe<void>::Ok();
 }
 
-REGISTER_OP_EXPR_GRAD_FUNCTION("adaptive_avg_pool2d", AdaptivePool);
+class AdaptiveAvgPool1dGrad final : public AdaptivePoolNdGrad {
+ public:
+  Maybe<void> Init(const OpExpr& op) override { return AdaptivePoolNdGrad::Init(op, "avg", 1); }
+};
+
+class AdaptiveAvgPool2dGrad final : public AdaptivePoolNdGrad {
+ public:
+  Maybe<void> Init(const OpExpr& op) override { return AdaptivePoolNdGrad::Init(op, "avg", 2); }
+};
+
+class AdaptiveAvgPool3dGrad final : public AdaptivePoolNdGrad {
+ public:
+  Maybe<void> Init(const OpExpr& op) override { return AdaptivePoolNdGrad::Init(op, "avg", 3); }
+};
+
+REGISTER_OP_EXPR_GRAD_FUNCTION("adaptive_avg_pool1d", AdaptiveAvgPool1dGrad);
+REGISTER_OP_EXPR_GRAD_FUNCTION("adaptive_avg_pool2d", AdaptiveAvgPool2dGrad);
+REGISTER_OP_EXPR_GRAD_FUNCTION("adaptive_avg_pool3d", AdaptiveAvgPool3dGrad);
 
 }  // namespace one
 }  // namespace oneflow

oneflow/core/framework/op_expr_helper.cpp

@@ -852,17 +852,6 @@ Maybe<one::UserOpExpr> PoolNdGradOp(const std::string& mode, const std::string&
       .Build();
 }
 
-Maybe<one::UserOpExpr> AdaptivePoolGradOp() {
-  return AdaptivePoolGradOp(UniqueOpName("adaptive_pool_grad"));
-}
-Maybe<one::UserOpExpr> AdaptivePoolGradOp(const std::string& name) {
-  return one::OpBuilder("adaptive_avg_pool2d_grad", name)
-      .Input("x")
-      .Input("dy")
-      .Output("dx")
-      .Build();
-}
-
 Maybe<one::UserOpExpr> UnsortedSegmentSumLikeOp(const int64_t& axis) {
   return UnsortedSegmentSumLikeOp(axis, UniqueOpName("unsorted_segment_sum_like"));
 }

oneflow/core/framework/op_expr_helper.h

@@ -276,9 +276,6 @@ Maybe<one::UserOpExpr> PoolNdGradOp(const std::string& mode, const std::string&
                                     const std::vector<int32_t>& strides, const bool& ceil_mode,
                                     const std::string& name);
 
-Maybe<one::UserOpExpr> AdaptivePoolGradOp();
-Maybe<one::UserOpExpr> AdaptivePoolGradOp(const std::string& name);
-
 Maybe<one::UserOpExpr> UnsortedSegmentSumLikeOp(const int64_t& axis);
 Maybe<one::UserOpExpr> UnsortedSegmentSumLikeOp(const int64_t& axis, const std::string& name);
 

oneflow/core/functional/functional_api.yaml

@@ -404,6 +404,25 @@
                       String data_format=\"channels_first\", Bool ceil_mode=False)"
   bind_python: True
 
+- name: "adaptive_avg_pool1d"
+  signature:
+    "Tensor AdaptiveAvgPool1D(Tensor x, *, Int64List output_size)"
+  bind_python: True
+
+- name: "adaptive_avg_pool2d"
+  signature:
+    "Tensor AdaptiveAvgPool2D(Tensor x, *, Int64List output_size)"
+  bind_python: True
+
+- name: "adaptive_avg_pool3d"
+  signature:
+    "Tensor AdaptiveAvgPool3D(Tensor x, *, Int64List output_size)"
+  bind_python: True
+
+- name: "adaptive_pool_grad"
+  signature:
+    "Tensor AdaptivePoolNdGrad(Tensor x, Tensor dy, *, String mode, Int32 ndims)"
+
 - name: "pool_grad"
   signature:
     "Tensor PoolNdGrad(Tensor x, Tensor y, Tensor dy, *, String mode, Int32 ndims, String data_format,

oneflow/core/functional/impl/nn_functor.cpp

@@ -287,6 +287,42 @@ class Maxpool3DFunctor : public PoolingNDFunctor {
   }
 };
 
+class AdaptivePoolNDFunctor {
+ public:
+  AdaptivePoolNDFunctor() = default;
+  virtual ~AdaptivePoolNDFunctor() = default;
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x,
+                           const std::vector<int64_t>& output_size) const {
+    MutableAttrMap attrs;
+    JUST(attrs.SetAttr<std::vector<int64_t>>("output_size", output_size));
+    return OpInterpUtil::Dispatch<Tensor>(*op_, {x}, attrs);
+  }
+
+ protected:
+  std::shared_ptr<OpExpr> op_;
+};
+
+class AdaptiveAvgPool1DFunctor : public AdaptivePoolNDFunctor {
+ public:
+  AdaptiveAvgPool1DFunctor() {
+    op_ = CHECK_JUST(one::OpBuilder("adaptive_avg_pool1d").Input("x").Output("y").Build());
+  }
+};
+
+class AdaptiveAvgPool2DFunctor : public AdaptivePoolNDFunctor {
+ public:
+  AdaptiveAvgPool2DFunctor() {
+    op_ = CHECK_JUST(one::OpBuilder("adaptive_avg_pool2d").Input("x").Output("y").Build());
+  }
+};
+
+class AdaptiveAvgPool3DFunctor : public AdaptivePoolNDFunctor {
+ public:
+  AdaptiveAvgPool3DFunctor() {
+    op_ = CHECK_JUST(one::OpBuilder("adaptive_avg_pool3d").Input("x").Output("y").Build());
+  }
+};
+
 class SparseSoftmaxCrossEntropyFunctor {
  public:
   SparseSoftmaxCrossEntropyFunctor() {
@@ -479,6 +515,9 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<impl::Maxpool2DFunctor>("Maxpool2D");
   m.add_functor<impl::Maxpool3DFunctor>("Maxpool3D");
   m.add_functor<impl::MaxPool2DFunctor>("MaxPool2D");
+  m.add_functor<impl::AdaptiveAvgPool1DFunctor>("AdaptiveAvgPool1D");
+  m.add_functor<impl::AdaptiveAvgPool2DFunctor>("AdaptiveAvgPool2D");
+  m.add_functor<impl::AdaptiveAvgPool3DFunctor>("AdaptiveAvgPool3D");
   m.add_functor<impl::SparseSoftmaxCrossEntropyFunctor>("SparseSoftmaxCrossEntropy");
   m.add_functor<impl::SmoothL1LossFunctor>("SmoothL1Loss");
   m.add_functor<impl::NormalizationFunctor>("Normalization");

oneflow/core/functional/impl/nn_grad_functor.cpp

@@ -199,6 +199,35 @@ class PoolNdGradFunctor {
   std::unordered_map<std::string, std::shared_ptr<OpExpr>> op_expr_map_;
 };
 
+class AdaptivePoolNdGradFunctor {
+ public:
+  AdaptivePoolNdGradFunctor() {
+    for (const auto& mode : {"avg"}) {
+      for (int ndims = 1; 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("dy").Output("dx").Build());
+      }
+    }
+  }
+  static std::string GetOpTypeName(const std::string& mode, const int32_t& ndims) {
+    return "adaptive_" + mode + "_pool" + std::to_string(ndims) + "d_grad";
+  }
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x,
+                           const std::shared_ptr<one::Tensor>& dy, const std::string& mode,
+                           const int32_t& ndims) const {
+    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 AdaptivePoolNdGradFunctor.";
+    CHECK_NOTNULL_OR_RETURN(it->second);
+    return OpInterpUtil::Dispatch<Tensor>(*it->second, {x, dy});
+  }
+
+ protected:
+  std::unordered_map<std::string, std::shared_ptr<OpExpr>> op_expr_map_;
+};
+
 class SmoothL1LossGradFunctor {
  public:
   SmoothL1LossGradFunctor() {
@@ -281,6 +310,7 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<impl::ConvFilterGradFunctor>("ConvFilterGrad");
   m.add_functor<impl::ConvDataGradFunctor>("ConvDataGrad");
   m.add_functor<impl::PoolNdGradFunctor>("PoolNdGrad");
+  m.add_functor<impl::AdaptivePoolNdGradFunctor>("AdaptivePoolNdGrad");
   m.add_functor<impl::SmoothL1LossGradFunctor>("SmoothL1LossGrad");
   m.add_functor<impl::PoolingNdGradFunctor>("PoolingNdGrad");
   m.add_functor<impl::PadGradFunctor>("PadGrad");