Add Unflatten Module (#41564)

alvgaona · facebook-github-bot · commit c89c294ef93b · 2020-07-21T07:43:02.000-07:00
Summary: This PR implements a feature extension discussed in #41516. I followed this other PR #22245 to add this other module. While I was at it, I also added `extra_repr()` method in `Flatten` which was missing. I see there are no unit tests for these modules. Should I add those too? If so, what is the best place I should place these? Pull Request resolved: #41564 Reviewed By: gchanan Differential Revision: D22636766 Pulled By: albanD fbshipit-source-id: f9efdefd3ffe7d9af9482087625344af8f990943
diff --git a/docs/source/nn.rst b/docs/source/nn.rst
@@ -355,6 +355,7 @@ Utility functions in other modules
     nn.utils.rnn.pack_sequence
 
     nn.Flatten
+    nn.Unflatten
 
 Quantized Functions
 --------------------
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -8326,6 +8326,51 @@ def test_functional_grad_conv(self):
             torch.nn.grad._grad_input_padding(torch.rand(1, 2, 3), [1, 2, 5], (1,), (0,), (3,))
         self.assertEqual(len(w), 1)
 
+    def test_unflatten(self):
+        tensor_input = torch.randn(2, 50)
+
+        # Unflatten Tensor
+
+        unflatten = nn.Unflatten(dim=1, unflattened_size=(2, 5, 5))
+        tensor_output = unflatten(tensor_input)
+        self.assertEqual(tensor_output.size(), torch.Size([2, 2, 5, 5]))
+
+        # Unflatten NamedTensor
+
+        unflatten = nn.Unflatten(dim='features', unflattened_size=(('C', 2), ('H', 5), ('W', 5)))
+        named_tensor_input = tensor_input.refine_names('N', 'features')
+        named_tensor_output = unflatten(named_tensor_input)
+        self.assertEqual(tensor_output.size(), torch.Size([2, 2, 5, 5]))
+
+    def test_unflatten_invalid_arg(self):
+        # Wrong type for unflattened_size (tuple of floats)
+
+        with self.assertRaisesRegex(
+                TypeError,
+                r"unflattened_size must be tuple of ints, but found element of type float at pos 2"):
+            nn.Unflatten(dim=1, unflattened_size=(2, 5, 5.0))
+
+        # Wrong type for unflattened_size (tuple of lists)
+
+        with self.assertRaisesRegex(
+                TypeError,
+                r"unflattened_size must be tuple of tuples, but found element of type list at pos 0"):
+            nn.Unflatten(dim='features', unflattened_size=(['C', 2], ['W', 5], ['H', 5]))
+
+        # Wrong type for unflattened_size (list of ints)
+
+        with self.assertRaisesRegex(
+                TypeError,
+                r"unflattened_size must be a tuple of ints, but found type list"):
+            nn.Unflatten(dim=1, unflattened_size=[2, 5, 5])
+
+        # Wrong type for unflattened_size (list of lists)
+
+        with self.assertRaisesRegex(
+                TypeError,
+                r"unflattened_size must be a tuple of tuples, but found type list"):
+            nn.Unflatten(dim='features', unflattened_size=[['C', 2], ['W', 5], ['H', 5]])
+
 
 class TestNNInit(TestCase):
     def setUp(self):
diff --git a/torch/nn/modules/__init__.py b/torch/nn/modules/__init__.py
@@ -30,7 +30,7 @@
 from .adaptive import AdaptiveLogSoftmaxWithLoss
 from .transformer import TransformerEncoder, TransformerDecoder, \
     TransformerEncoderLayer, TransformerDecoderLayer, Transformer
-from .flatten import Flatten
+from .flatten import Flatten, Unflatten
 
 __all__ = [
     'Module', 'Identity', 'Linear', 'Conv1d', 'Conv2d', 'Conv3d', 'ConvTranspose1d',
@@ -54,5 +54,5 @@
     'ConstantPad3d', 'Bilinear', 'CosineSimilarity', 'Unfold', 'Fold',
     'AdaptiveLogSoftmaxWithLoss', 'TransformerEncoder', 'TransformerDecoder',
     'TransformerEncoderLayer', 'TransformerDecoderLayer', 'Transformer',
-    'Flatten', 'Hardsigmoid', 'Hardswish', 'SiLU',
+    'Flatten', 'Unflatten', 'Hardsigmoid', 'Hardswish', 'SiLU',
 ]
diff --git a/torch/nn/modules/flatten.py b/torch/nn/modules/flatten.py
@@ -1,6 +1,9 @@
 from .module import Module
 
+from typing import Union
 from torch import Tensor
+from torch import Size
+
 
 class Flatten(Module):
     r"""
@@ -31,3 +34,106 @@ def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None:
 
     def forward(self, input: Tensor) -> Tensor:
         return input.flatten(self.start_dim, self.end_dim)
+
+    def extra_repr(self) -> str:
+        return 'start_dim={}, end_dim={}'.format(
+            self.start_dim, self.end_dim
+        )
+
+
+class Unflatten(Module):
+    r"""
+    Unflattens a tensor into another tensor of a desired shape. For use with :class:`~nn.Sequential`.
+
+    * :attr:`dim` specifies the dimension of the input tensor to be flattened, and it can 
+    be either `str` or `int` when `NamedTensor` or `Tensor` is used, respectively.
+
+    * :attr:`unflattened_size` is the size of the unflattened dimension of the tensor and it can be a
+    `namedshape` (`tuple` of tuples) if :attr:`dim` is `str` or a `tuple` of ints as well as `torch.Size` if
+    :attr:`dim` is an `int`.
+
+    Shape:
+        - Input: :math:`(N, *dims)`
+        - Output: :math:`(N, C_out, H_out, W_out)`
+
+    Args:
+        dim (Union[int, str]): Dimension to be flattened
+        unflattened_size (Union[tuple, torch.Size]): Size of the output tensor
+
+    Examples:
+        >>> input = torch.randn(2, 50)
+        >>> # With tuple of ints
+        >>> m = nn.Sequential(
+        >>>     nn.Linear(50, 50),
+        >>>     nn.Unflatten(1, (2, 5, 5))
+        >>> )
+        >>> output = m(output)
+        >>> output.size()
+        torch.Size([2, 2, 5, 5])
+        >>> # With torch.Size
+        >>> m = nn.Sequential(
+        >>>     nn.Linear(50, 50),
+        >>>     nn.Unflatten(1, torch.Size([2, 5, 5]))
+        >>> )
+        >>> output = m(output)
+        >>> output.size()
+        torch.Size([2, 2, 5, 5])
+        >>> # With namedshape (tuple of tuples)
+        >>> m = nn.Sequential(
+        >>>     nn.Linear(50, 50),
+        >>>     nn.Unflatten('features', (('C', 2), ('H', 50), ('W',50)))
+        >>> )
+        >>> output = m(output)
+        >>> output.size()
+        torch.Size([2, 2, 5, 5])
+    """
+    __constants__ = ['dim', 'unflattened_size']
+    dim: Union[int, str]
+    unflattened_size: Union[tuple, Size]
+
+    def __init__(self, dim: Union[int, str], unflattened_size: Union[tuple, Size]) -> None:
+        super(Unflatten, self).__init__()
+        if isinstance(dim, int):
+            self._require_tuple_int(unflattened_size)
+            self.named = False
+        else:
+            self._require_tuple_tuple(unflattened_size)
+            self.named = True
+
+        self.dim = dim
+        self.unflattened_size = unflattened_size
+
+    def _require_tuple_tuple(self, input):
+        if (isinstance(input, tuple)):
+            for idx, elem in enumerate(input):
+                if not isinstance(elem, tuple):
+                    raise TypeError("unflattened_size must be tuple of tuples, " + 
+                                    "but found element of type {} at pos {}".format(type(elem).__name__, idx))
+            return
+        raise TypeError("unflattened_size must be a tuple of tuples, " +
+                        "but found type {}".format(type(input).__name__))
+
+    def _require_tuple_int(self, input):
+        if (isinstance(input, tuple)):
+            for idx, elem in enumerate(input):
+                if not isinstance(elem, int):
+                    raise TypeError("unflattened_size must be tuple of ints, " + 
+                                    "but found element of type {} at pos {}".format(type(elem).__name__, idx))
+            return
+        raise TypeError("unflattened_size must be a tuple of ints, but found type {}".format(type(input).__name__))
+
+    def forward(self, input: Tensor) -> Tensor:
+        if self.named:
+            return input.unflatten(self.dim, self.unflattened_size)
+        else:
+            dim = int(self.dim)
+            if dim < 0:
+                dim += input.dim()
+            inp_size = list(input.size())
+            new_size = inp_size[:dim] + list(self.unflattened_size) + inp_size[dim + 1:]
+            return input.view(new_size)
+
+    def extra_repr(self) -> str:
+        return 'dim={}, unflattened_size={}'.format(
+            self.dim, self.unflattened_size
+        )
