Repository: pytorch/pytorch
Pull request: allow output_size to contain None in adaptive pooling methods
This PR fixes #7884.
Since #3127 we support output_size to contain None (use corresponding number in input_size), but this got lost in TH refactor. Adding it back and add some test cases in test_nn.py.

Problem statement:
[pytorch] A bug for torch.nn.AdaptiveMaxPool2d
## Issue description
I run the example code of pytorch, but there was a bug.
```
>>> # target output size of 10x7
>>> m = nn.AdaptiveMaxPool2d((None, 7))
>>> input = torch.randn(1, 64, 10, 9)
>>> output = m(input)
```
When i used "nn.AdaptiveMaxPool2d((10, 7))", it was OK.

## Code example

```
In [17]: input = torch.randn(1, 64, 10, 9)

In [18]: m = nn.AdaptiveMaxPool2d((None, 7))

In [19]: output = m(input)   
```

## System Info
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-19-0b450ddad2f3> in <module>()
----> 1 output = m(input)

/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.pyc in __call__(self, *input, **kwargs)
    489             result = self._slow_forward(*input, **kwargs)
    490         else:
--> 491             result = self.forward(*input, **kwargs)
    492         for hook in self._forward_hooks.values():
    493             hook_result = hook(self, input, result)

/usr/local/lib/python2.7/dist-packages/torch/nn/modules/pooling.pyc in forward(self, input)
    848 
    849     def forward(self, input):
--> 850         return F.adaptive_max_pool2d(input, self.output_size, self.return_indices)
    851 
    852 

/usr/local/lib/python2.7/dist-packages/torch/nn/functional.pyc in adaptive_max_pool2d(input, output_size, return_indices)
    489         return_indices: whether to return pooling indices. Default: ``False``
    490     """
--> 491     ret = torch._C._nn.adaptive_max_pool2d(input, output_size)
    492     return ret if return_indices else ret[0]
    493 

TypeError: an integer is required
```

## Code example

```
In [20]: input = torch.randn(1, 64, 10, 9)

In [21]: m = nn.AdaptiveMaxPool2d((10, 7))

In [22]: output = m(input)

In [23]: output.size()
```

## System Info

`Out[23]: torch.Size([1, 64, 10, 7])`

Hints from issue discussion:
I can repro this on master branch. Will work on a fix.

Commit message: fix lint

Patch to review:
diff --git a/test/test_nn.py b/test/test_nn.py
index bea6820abab1..cda8694ca122 100644
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -7206,6 +7206,12 @@ def multimarginloss_weights_no_reduce_test():
         input_fn=lambda: _rand_tensor_non_equal(1, 3, 5, 6),
         desc='tuple',
     ),
+    dict(
+        module_name='AdaptiveMaxPool2d',
+        constructor_args=((3, None),),
+        input_fn=lambda: _rand_tensor_non_equal(1, 3, 5, 6),
+        desc='tuple_none',
+    ),
     dict(
         module_name='AdaptiveMaxPool3d',
         constructor_args=(3,),
@@ -7218,6 +7224,12 @@ def multimarginloss_weights_no_reduce_test():
         input_fn=lambda: _rand_tensor_non_equal(2, 3, 5, 6, 7),
         desc='tuple',
     ),
+    dict(
+        module_name='AdaptiveMaxPool3d',
+        constructor_args=((3, None, 5),),
+        input_fn=lambda: _rand_tensor_non_equal(2, 3, 5, 6, 7),
+        desc='tuple_none',
+    ),
     dict(
         module_name='AdaptiveMaxPool3d',
         constructor_args=(3,),
@@ -7247,6 +7259,12 @@ def multimarginloss_weights_no_reduce_test():
         input_fn=lambda: torch.rand(1, 3, 5, 6),
         desc='tuple',
     ),
+    dict(
+        module_name='AdaptiveAvgPool2d',
+        constructor_args=((3, None),),
+        input_fn=lambda: torch.rand(1, 3, 5, 6),
+        desc='tuple_none',
+    ),
     dict(
         module_name='AdaptiveAvgPool3d',
         constructor_args=(3,),
@@ -7259,6 +7277,12 @@ def multimarginloss_weights_no_reduce_test():
         input_fn=lambda: torch.rand(2, 3, 5, 3, 7),
         desc='tuple',
     ),
+    dict(
+        module_name='AdaptiveAvgPool3d',
+        constructor_args=((None, 4, 5),),
+        input_fn=lambda: torch.rand(2, 3, 5, 3, 7),
+        desc='tuple_none',
+    ),
     dict(
         module_name='SELU',
         input_size=(3, 2, 5),
diff --git a/torch/nn/functional.py b/torch/nn/functional.py
index 4f09a5eeec69..724a3394061d 100644
--- a/torch/nn/functional.py
+++ b/torch/nn/functional.py
@@ -12,7 +12,7 @@
 from ._functions.padding import ConstantPadNd
 from ._functions import vision
 from ._functions.thnn.fold import Col2Im, Im2Col
-from .modules.utils import _single, _pair, _triple
+from .modules.utils import _single, _pair, _triple, set_output_size
 from . import grad
 
 
@@ -490,6 +490,7 @@ def adaptive_max_pool2d(input, output_size, return_indices=False):
             double-integer tuple)
         return_indices: whether to return pooling indices. Default: ``False``
     """
+    output_size = set_output_size(input.size()[2:], output_size)
     ret = torch._C._nn.adaptive_max_pool2d(input, output_size)
     return ret if return_indices else ret[0]
 
@@ -505,6 +506,7 @@ def adaptive_max_pool3d(input, output_size, return_indices=False):
             triple-integer tuple)
         return_indices: whether to return pooling indices. Default: ``False``
     """
+    output_size = set_output_size(input.size()[2:], output_size)
     ret = torch._C._nn.adaptive_max_pool3d(input, output_size)
     return ret if return_indices else ret[0]
 
@@ -521,35 +523,38 @@ def adaptive_max_pool3d(input, output_size, return_indices=False):
     output_size: the target output size (single integer)
 """)
 
-adaptive_avg_pool2d = _add_docstr(torch._C._nn.adaptive_avg_pool2d, r"""
-adaptive_avg_pool2d(input, output_size) -> Tensor
 
-Applies a 2D adaptive average pooling over an input signal composed of
-several input planes.
+def adaptive_avg_pool2d(input, output_size):
+    r"""
+    Applies a 2D adaptive average pooling over an input signal composed of
+    several input planes.
 
-See :class:`~torch.nn.AdaptiveAvgPool2d` for details and output shape.
+    See :class:`~torch.nn.AdaptiveAvgPool2d` for details and output shape.
 
-Args:
-    output_size: the target output size (single integer or
-        double-integer tuple)
-""")
+    Args:
+        output_size: the target output size (single integer or
+            double-integer tuple)
+    """
+    output_size = set_output_size(input.size()[2:], output_size)
+    return torch._C._nn.adaptive_avg_pool2d(input, output_size)
 
-adaptive_avg_pool3d = _add_docstr(torch._C._nn.adaptive_avg_pool3d, r"""
-adaptive_avg_pool3d(input, output_size) -> Tensor
 
-Applies a 3D adaptive average pooling over an input signal composed of
-several input planes.
+def adaptive_avg_pool3d(input, output_size):
+    r"""
+    Applies a 3D adaptive average pooling over an input signal composed of
+    several input planes.
 
-See :class:`~torch.nn.AdaptiveAvgPool3d` for details and output shape.
+    See :class:`~torch.nn.AdaptiveAvgPool3d` for details and output shape.
 
-Args:
-    output_size: the target output size (single integer or
-        triple-integer tuple)
-""")
+    Args:
+        output_size: the target output size (single integer or
+            triple-integer tuple)
+    """
+    output_size = set_output_size(input.size()[2:], output_size)
+    return torch._C._nn.adaptive_avg_pool3d(input, output_size)
 
 
 # Activation functions
-
 def dropout(input, p=0.5, training=False, inplace=False):
     return _functions.dropout.Dropout.apply(input, p, training, inplace)
 
diff --git a/torch/nn/modules/utils.py b/torch/nn/modules/utils.py
index c3dc5126e654..4faaee8d11f2 100644
--- a/torch/nn/modules/utils.py
+++ b/torch/nn/modules/utils.py
@@ -13,3 +13,18 @@ def parse(x):
 _pair = _ntuple(2)
 _triple = _ntuple(3)
 _quadruple = _ntuple(4)
+
+
+def set_output_size(input_size, output_size):
+    dim = len(input_size)
+    if isinstance(output_size, int):
+        return _ntuple(dim)(output_size)
+    if len(output_size) == 1:
+        return _ntuple(dim)(output_size[0])
+    assert(len(output_size) == dim)
+
+    shape = ()
+    for i, s in enumerate(output_size):
+        shape += ((s or input_size[i]), )
+
+    return shape


Code review comments:
diff_hunk: @@ -13,3 +13,18 @@ def parse(x):
 _pair = _ntuple(2)
 _triple = _ntuple(3)
 _quadruple = _ntuple(4)
+
+
+def set_output_size(input_size, output_size):
+    dim = len(input_size)
+    if isinstance(output_size, int):
+        return _ntuple(dim)(output_size)
+    if len(output_size) == 1:
+        return _ntuple(dim)(output_size[0])
+    assert(len(output_size) == dim)
+
+    shape = ()
+    for i, s in enumerate(output_size):
+        shape += ((s or input_size[i]), )
path: torch/nn/modules/utils.py
text: @user1:
If `s = 0` this will be `input_size[i]`

diff_hunk: @@ -13,3 +13,18 @@ def parse(x):
 _pair = _ntuple(2)
 _triple = _ntuple(3)
 _quadruple = _ntuple(4)
+
+
+def set_output_size(input_size, output_size):
path: torch/nn/modules/utils.py
text: @user1:
Helper functions that aren't part of public API should be prefixed with `_`.

diff_hunk: @@ -521,35 +523,38 @@ def adaptive_max_pool3d(input, output_size, return_indices=False):
     output_size: the target output size (single integer)
 """)
 
-adaptive_avg_pool2d = _add_docstr(torch._C._nn.adaptive_avg_pool2d, r"""
-adaptive_avg_pool2d(input, output_size) -> Tensor
 
-Applies a 2D adaptive average pooling over an input signal composed of
-several input planes.
+def adaptive_avg_pool2d(input, output_size):
+    r"""
+    Applies a 2D adaptive average pooling over an input signal composed of
+    several input planes.
 
-See :class:`~torch.nn.AdaptiveAvgPool2d` for details and output shape.
+    See :class:`~torch.nn.AdaptiveAvgPool2d` for details and output shape.
 
-Args:
-    output_size: the target output size (single integer or
-        double-integer tuple)
-""")
+    Args:
+        output_size: the target output size (single integer or
+            double-integer tuple)
+    """
+    output_size = set_output_size(input.size()[2:], output_size)
+    return torch._C._nn.adaptive_avg_pool2d(input, output_size)
 
-adaptive_avg_pool3d = _add_docstr(torch._C._nn.adaptive_avg_pool3d, r"""
-adaptive_avg_pool3d(input, output_size) -> Tensor
 
-Applies a 3D adaptive average pooling over an input signal composed of
-several input planes.
+def adaptive_avg_pool3d(input, output_size):
+    r"""
+    Applies a 3D adaptive average pooling over an input signal composed of
+    several input planes.
 
-See :class:`~torch.nn.AdaptiveAvgPool3d` for details and output shape.
+    See :class:`~torch.nn.AdaptiveAvgPool3d` for details and output shape.
 
-Args:
-    output_size: the target output size (single integer or
-        triple-integer tuple)
-""")
+    Args:
+        output_size: the target output size (single integer or
+            triple-integer tuple)
+    """
+    output_size = set_output_size(input.size()[2:], output_size)
path: torch/nn/functional.py
text: @user1:
This will be problematic if input has less than `len(output_size) - 2` dimensions. (loop in `set_output_size` will throw).

How about writing such a helper?
```py
# uses `defaults[-len(input):]` as default values for `input`
def _list_with_default(input, defaults):
  if len(defaults) < len(input):
     raise
  return [v if v is not None else d for v, d in zip(input, default[-len(input):])]
```
You can tweak the name/imple/error message to make it nicer. What do you think?

With something like this you should also add 1 or 2 test that asserts some inputs triggers error.