Add nn.AdaptiveAvgPool1d and nn.AdaptiveAvgPool3d

oneflow/core/autograd/gradient_funcs/adaptive_pool.cpp

@@ -26,9 +26,9 @@ struct AdaptivePoolInterpState : public OpExprInterpState {
   bool requires_grad;
 };
 
-class AdaptivePool : public OpExprGradFunction<AdaptivePoolInterpState> {
+class AdaptivePoolNd : public OpExprGradFunction<AdaptivePoolInterpState> {
  public:
-  Maybe<void> Init(const OpExpr& op) override;
+  Maybe<void> Init(const OpExpr& op, 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,
@@ -39,16 +39,16 @@ class AdaptivePool : public OpExprGradFunction<AdaptivePoolInterpState> {
   std::shared_ptr<OpExpr> grad_op_;
 };
 
-Maybe<void> AdaptivePool::Init(const OpExpr& op) {
+Maybe<void> AdaptivePoolNd::Init(const OpExpr& op, 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)));
+  grad_op_ = JUST(op_expr_helper::AdaptivePoolNdGradOp(ndims, GradientOpName(op_name)));
   return Maybe<void>::Ok();
 }
 
-Maybe<void> AdaptivePool::Capture(AdaptivePoolInterpState* ctx, const TensorTuple& inputs,
+Maybe<void> AdaptivePoolNd::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(); }
@@ -57,7 +57,7 @@ Maybe<void> AdaptivePool::Capture(AdaptivePoolInterpState* ctx, const TensorTupl
   return Maybe<void>::Ok();
 }
 
-Maybe<void> AdaptivePool::Apply(const AdaptivePoolInterpState* ctx, const TensorTuple& out_grads,
+Maybe<void> AdaptivePoolNd::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);
@@ -68,7 +68,24 @@ Maybe<void> AdaptivePool::Apply(const AdaptivePoolInterpState* ctx, const Tensor
   return Maybe<void>::Ok();
 }
 
-REGISTER_OP_EXPR_GRAD_FUNCTION("adaptive_avg_pool2d", AdaptivePool);
+class AdaptivePool1d final : public AdaptivePoolNd {
+ public:
+  Maybe<void> Init(const OpExpr& op) override { return AdaptivePoolNd::Init(op, 1); }
+};
+
+class AdaptivePool2d final : public AdaptivePoolNd {
+ public:
+  Maybe<void> Init(const OpExpr& op) override { return AdaptivePoolNd::Init(op, 2); }
+};
+
+class AdaptivePool3d final : public AdaptivePoolNd {
+ public:
+  Maybe<void> Init(const OpExpr& op) override { return AdaptivePoolNd::Init(op, 3); }
+};
+
+REGISTER_OP_EXPR_GRAD_FUNCTION("adaptive_avg_pool1d", AdaptivePool1d);
+REGISTER_OP_EXPR_GRAD_FUNCTION("adaptive_avg_pool2d", AdaptivePool2d);
+REGISTER_OP_EXPR_GRAD_FUNCTION("adaptive_avg_pool3d", AdaptivePool3d);
 
 }  // namespace one
 }  // namespace oneflow

oneflow/core/framework/op_expr_helper.cpp

@@ -851,11 +851,11 @@ 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> AdaptivePoolNdGradOp(const int& ndims) {
+  return AdaptivePoolNdGradOp(ndims, UniqueOpName("adaptive_avg_pool_nd_grad"));
 }
-Maybe<one::UserOpExpr> AdaptivePoolGradOp(const std::string& name) {
-  return one::OpBuilder("adaptive_avg_pool2d_grad", name)
+Maybe<one::UserOpExpr> AdaptivePoolNdGradOp(const int& ndims, const std::string& name) {
+  return one::OpBuilder("adaptive_avg_pool" + std::to_string(ndims) + "d_grad", name)
       .Input("x")
       .Input("dy")
       .Output("dx")

oneflow/core/framework/op_expr_helper.h

@@ -276,8 +276,8 @@ 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> AdaptivePoolNdGradOp();
+Maybe<one::UserOpExpr> AdaptivePoolNdGradOp(const int& ndims, 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/python/nn/modules/adaptive_pool.py

@@ -18,6 +18,37 @@
 from oneflow.python.oneflow_export import oneflow_export, experimental_api
 
 
+@oneflow_export("nn.AdaptiveAvgPool1d")
+@experimental_api
+class AdaptiveAvgPool1d(Module):
+    def __init__(self, output_size) -> None:
+        super().__init__()
+        self.output_size = output_size
+
+        self._op = (
+            flow.builtin_op("adaptive_avg_pool1d")
+                .Input("x")
+                .Attr("output_size", [])
+                .Output("y")
+                .Build()
+        )
+
+    def forward(self, x):
+        new_output_size = None
+        assert len(x.shape) == 3
+
+        if isinstance(self.output_size, int):
+            new_output_size = tuple([self.output_size])
+        elif isinstance(self.output_size, tuple):
+            pass
+        else:
+            raise NotImplementedError("output_size param wrong, please check!")
+
+        assert (new_output_size is not None and new_output_size[0] <= x.shape[2]), f"output_size param wrong, please check!"
+
+        return self._op(x, output_size=new_output_size)[0]
+
+
 @oneflow_export("nn.AdaptiveAvgPool2d")
 @experimental_api
 class AdaptiveAvgPool2d(Module):
@@ -100,6 +131,53 @@ def forward(self, x):
         return self._op(x, output_size=new_output_size)[0]
 
 
+@oneflow_export("nn.AdaptiveAvgPool3d")
+@experimental_api
+class AdaptiveAvgPool3d(Module):
+    def __init__(self, output_size) -> None:
+        super().__init__()
+        self.output_size = output_size
+
+        self._op = (
+            flow.builtin_op("adaptive_avg_pool3d")
+                .Input("x")
+                .Attr("output_size", [])
+                .Output("y")
+                .Build()
+        )
+
+    def forward(self, x):
+        new_output_size = []
+        assert len(x.shape) == 5
+
+        if isinstance(self.output_size, int):
+            for _ in range(len(x.shape) - 2):
+                new_output_size.append(self.output_size)
+        elif isinstance(self.output_size, tuple):
+            new_output_size = list(self.output_size)
+            if self.output_size[0] is None:
+                new_output_size[0] = x.shape[2]
+            if self.output_size[1] is None:
+                new_output_size[1] = x.shape[3]
+            if self.output_size[2] is None:
+                new_output_size[2] = x.shape[4]
+        else:
+            raise NotImplementedError("output_size param wrong, please check!")
+
+        new_output_size = tuple(new_output_size)
+        assert (
+                new_output_size[0] <= x.shape[2]
+        ), f"output_size param wrong, please check!"
+        assert (
+                new_output_size[1] <= x.shape[3]
+        ), f"output_size param wrong, please check!"
+        assert (
+                new_output_size[2] <= x.shape[4]
+        ), f"output_size param wrong, please check!"
+
+        return self._op(x, output_size=new_output_size)[0]
+
+
 if __name__ == "__main__":
     import doctest