Align Torch.nn.functional poolXd (#6184)

* fix backward index for gamma beta

* move layernormParamGrad to functional

* fix bug in functional

* still error in autotest

* fix(LayerNorm): fix LayerNorm cpu backward bug

* rename test layernorm

* use functional grad and use elementwise affine

* fix comment and skip autotest temporary

* check shape

* merge

* align torch functional pool

Co-authored-by: Yinggang Wang <wyg19970408@gmail.com>
Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

oneflow/core/functional/functional_api.yaml

@@ -621,23 +621,26 @@
 
 - name: "max_pool1d"
   signature:
-    'TensorTuple (Tensor x, String data_format="channels_first", Int32List[1] padding,
-    Int32List[1] kernel_size, Int32List[1] stride, Int32List[1] dilation,
-    Bool return_indices=True, Bool ceil_mode=False) => Maxpool1D'
+    'TensorTuple (Tensor x, Int32List[1] kernel_size, Int32List[1] stride=1, 
+                  Int32List[1] padding=0, Int32List[1] dilation=1, 
+                  Bool return_indices=True, Bool ceil_mode=False, 
+                  String data_format="channels_first") => Maxpool1D'
   bind_python: True
 
 - name: "max_pool2d"
   signature:
-    'TensorTuple (Tensor x, String data_format="channels_first", Int32List[2] padding,
-    Int32List[2] kernel_size, Int32List[2] stride, Int32List[2] dilation,
-    Bool return_indices=True, Bool ceil_mode=False) => Maxpool2D'
+    'TensorTuple (Tensor x, Int32List[2] kernel_size, Int32List[2] stride=1, 
+                  Int32List[2] padding=0, Int32List[2] dilation=1, 
+                  Bool return_indices=True, Bool ceil_mode=False, 
+                  String data_format="channels_first") => Maxpool2D'
   bind_python: True
 
 - name: "max_pool3d"
   signature:
-    'TensorTuple (Tensor x, String data_format="channels_first", Int32List[3] padding,
-    Int32List[3] kernel_size, Int32List[3] stride, Int32List[3] dilation,
-    Bool return_indices=True, Bool ceil_mode=False) => Maxpool3D'
+    'TensorTuple (Tensor x, Int32List[3] kernel_size, Int32List[3] stride=1, 
+                  Int32List[3] padding=0, Int32List[3] dilation=1, 
+                  Bool return_indices=True, Bool ceil_mode=False, 
+                  String data_format="channels_first") => Maxpool3D'
   bind_python: True
 
 - name: "pooling_grad"
@@ -976,23 +979,23 @@
 
 - name: "avg_pool1d"
   signature:
-    'Tensor (Tensor x, String data_format="channels_first", Int32List[1] padding,
-    Int32List[1] kernel_size, Int32List[1] stride, Bool ceil_mode=False, Bool count_include_pad=True,
-    Int64 divisor_override=0) => Avgpool1D'
+    'Tensor (Tensor x, Int32List[1] kernel_size, Int32List[1] stride=1, 
+             Int32List[1] padding=0, Bool ceil_mode=False, Bool count_include_pad=True,
+             Int64 divisor_override=0, String data_format="channels_first") => Avgpool1D'
   bind_python: True
 
 - name: "avg_pool2d"
   signature:
-    'Tensor (Tensor x, String data_format="channels_first", Int32List[2] padding,
-    Int32List[2] kernel_size, Int32List[2] stride, Bool ceil_mode=False, Bool count_include_pad=True,
-    Int64 divisor_override=0) => Avgpool2D'
+    'Tensor (Tensor x, Int32List[2] kernel_size, Int32List[2] stride=1, 
+             Int32List[2] padding=0, Bool ceil_mode=False, Bool count_include_pad=True,
+             Int64 divisor_override=0, String data_format="channels_first") => Avgpool2D'
   bind_python: True
 
 - name: "avg_pool3d"
   signature:
-    'Tensor (Tensor x, String data_format="channels_first", Int32List[3] padding,
-    Int32List[3] kernel_size, Int32List[3] stride, Bool ceil_mode=False, Bool count_include_pad=True,
-    Int64 divisor_override=0) => Avgpool3D'
+    'Tensor (Tensor x, Int32List[3] kernel_size, Int32List[3] stride=1, 
+             Int32List[3] padding=0, Bool ceil_mode=False, Bool count_include_pad=True,
+             Int64 divisor_override=0, String data_format="channels_first") => Avgpool3D'
   bind_python: True
 
 - name: "avg_pooling_grad"

oneflow/core/functional/impl/nn_functor.cpp

@@ -265,11 +265,11 @@ class PoolingNDFunctor {
   PoolingNDFunctor() = default;
   virtual ~PoolingNDFunctor() = default;
   Maybe<TensorTuple> operator()(const std::shared_ptr<one::Tensor>& x,
-                                const std::string& data_format, const std::vector<int32_t>& padding,
                                 const std::vector<int32_t>& kernel_size,
                                 const std::vector<int32_t>& stride,
+                                const std::vector<int32_t>& padding,
                                 const std::vector<int32_t>& dilation, const bool& return_indices,
-                                const bool& ceil_mode) const {
+                                const bool& ceil_mode, const std::string& data_format) const {
     MutableAttrMap attrs;
     JUST(attrs.SetAttr<std::string>("data_format", data_format));
     JUST(attrs.SetAttr<std::vector<int32_t>>("padding", padding));
@@ -682,11 +682,11 @@ class AvgPoolingNDFunctor {
  public:
   AvgPoolingNDFunctor() = default;
   virtual ~AvgPoolingNDFunctor() = default;
-  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const std::string& data_format,
-                           const std::vector<int32_t>& padding,
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x,
                            const std::vector<int32_t>& kernel_size,
-                           const std::vector<int32_t>& stride, const bool& ceil_mode,
-                           const bool& count_include_pad, const int64_t& divisor_override) const {
+                           const std::vector<int32_t>& stride, const std::vector<int32_t>& padding,
+                           const bool& ceil_mode, const bool& count_include_pad,
+                           const int64_t& divisor_override, const std::string& data_format) const {
     MutableAttrMap attrs;
     JUST(attrs.SetAttr<std::string>("data_format", data_format));
     JUST(attrs.SetAttr<std::vector<int32_t>>("padding", padding));

python/oneflow/nn/functional/__init__.py

@@ -23,9 +23,9 @@
 from oneflow._C import avg_pool1d
 from oneflow._C import avg_pool2d
 from oneflow._C import avg_pool3d
-from oneflow._C import max_pool1d
-from oneflow._C import max_pool2d
-from oneflow._C import max_pool3d
+from .functional_maxpool import max_pool1d
+from .functional_maxpool import max_pool2d
+from .functional_maxpool import max_pool3d
 from oneflow._C import adaptive_avg_pool1d
 from oneflow._C import adaptive_avg_pool2d
 from oneflow._C import adaptive_avg_pool3d

python/oneflow/nn/functional/functional_maxpool.py

@@ -0,0 +1,95 @@
+"""
+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.
+"""
+import oneflow
+
+# oneflow._C.max_poolXd returns a TensorTuple, to align torch,
+# here we return different result according to the param `return_indices`.
+def max_pool1d(
+    x,
+    kernel_size,
+    stride=1,
+    padding=0,
+    dilation=1,
+    return_indices=False,
+    ceil_mode=False,
+    data_format="channels_first",
+):
+    _max_pool_out = oneflow._C.max_pool1d(
+        x,
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        return_indices,
+        ceil_mode,
+        data_format,
+    )
+    if return_indices:
+        return _max_pool_out
+    else:
+        return _max_pool_out[0]
+
+
+def max_pool2d(
+    x,
+    kernel_size,
+    stride=1,
+    padding=0,
+    dilation=1,
+    return_indices=False,
+    ceil_mode=False,
+    data_format="channels_first",
+):
+    _max_pool_out = oneflow._C.max_pool2d(
+        x,
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        return_indices,
+        ceil_mode,
+        data_format,
+    )
+    if return_indices:
+        return _max_pool_out
+    else:
+        return _max_pool_out[0]
+
+
+def max_pool3d(
+    x,
+    kernel_size,
+    stride=1,
+    padding=0,
+    dilation=1,
+    return_indices=False,
+    ceil_mode=False,
+    data_format="channels_first",
+):
+    _max_pool_out = oneflow._C.max_pool3d(
+        x,
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        return_indices,
+        ceil_mode,
+        data_format,
+    )
+    if return_indices:
+        return _max_pool_out
+    else:
+        return _max_pool_out[0]

python/oneflow/nn/modules/pooling.py

@@ -103,13 +103,13 @@ def __init__(
     def forward(self, x):
         y, indice = flow._C.max_pool1d(
             x,
-            data_format=self.channel_pos,
-            padding=self.padding,
             kernel_size=self.kernel_size,
             stride=self.stride,
+            padding=self.padding,
             dilation=self.dilation,
             return_indices=True,
             ceil_mode=self.ceil_mode,
+            data_format=self.channel_pos,
         )
         if self.return_indices:
             return y, indice
@@ -210,13 +210,13 @@ def __init__(
     def forward(self, x):
         y, indice = flow._C.max_pool2d(
             x,
-            data_format=self.channel_pos,
-            padding=self.padding,
             kernel_size=self.kernel_size,
             stride=self.stride,
+            padding=self.padding,
             dilation=self.dilation,
             return_indices=True,
             ceil_mode=self.ceil_mode,
+            data_format=self.channel_pos,
         )
         if self.return_indices:
             return y, indice
@@ -324,13 +324,13 @@ def __init__(
     def forward(self, x):
         y, indice = flow._C.max_pool3d(
             x,
-            data_format=self.channel_pos,
-            padding=self.padding,
             kernel_size=self.kernel_size,
             stride=self.stride,
+            padding=self.padding,
             dilation=self.dilation,
             return_indices=True,
             ceil_mode=self.ceil_mode,
+            data_format=self.channel_pos,
         )
 
         if self.return_indices: