145 lines
5.3 KiB
Python
145 lines
5.3 KiB
Python
from .module import Module
|
|
|
|
from typing import Tuple, Union
|
|
from torch import Tensor
|
|
from torch.types import _size
|
|
|
|
__all__ = ['Flatten', 'Unflatten']
|
|
|
|
class Flatten(Module):
|
|
r"""
|
|
Flattens a contiguous range of dims into a tensor.
|
|
|
|
For use with :class:`~nn.Sequential`, see :meth:`torch.flatten` for details.
|
|
|
|
Shape:
|
|
- Input: :math:`(*, S_{\text{start}},..., S_{i}, ..., S_{\text{end}}, *)`,'
|
|
where :math:`S_{i}` is the size at dimension :math:`i` and :math:`*` means any
|
|
number of dimensions including none.
|
|
- Output: :math:`(*, \prod_{i=\text{start}}^{\text{end}} S_{i}, *)`.
|
|
|
|
Args:
|
|
start_dim: first dim to flatten (default = 1).
|
|
end_dim: last dim to flatten (default = -1).
|
|
|
|
Examples::
|
|
>>> input = torch.randn(32, 1, 5, 5)
|
|
>>> # With default parameters
|
|
>>> m = nn.Flatten()
|
|
>>> output = m(input)
|
|
>>> output.size()
|
|
torch.Size([32, 25])
|
|
>>> # With non-default parameters
|
|
>>> m = nn.Flatten(0, 2)
|
|
>>> output = m(input)
|
|
>>> output.size()
|
|
torch.Size([160, 5])
|
|
"""
|
|
|
|
__constants__ = ['start_dim', 'end_dim']
|
|
start_dim: int
|
|
end_dim: int
|
|
|
|
def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None:
|
|
super().__init__()
|
|
self.start_dim = start_dim
|
|
self.end_dim = end_dim
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
return input.flatten(self.start_dim, self.end_dim)
|
|
|
|
def extra_repr(self) -> str:
|
|
return f'start_dim={self.start_dim}, end_dim={self.end_dim}'
|
|
|
|
|
|
class Unflatten(Module):
|
|
r"""
|
|
Unflattens a tensor dim expanding it to a desired shape. For use with :class:`~nn.Sequential`.
|
|
|
|
* :attr:`dim` specifies the dimension of the input tensor to be unflattened, and it can
|
|
be either `int` or `str` when `Tensor` or `NamedTensor` is used, respectively.
|
|
|
|
* :attr:`unflattened_size` is the new shape of the unflattened dimension of the tensor and it can be
|
|
a `tuple` of ints or a `list` of ints or `torch.Size` for `Tensor` input; a `NamedShape`
|
|
(tuple of `(name, size)` tuples) for `NamedTensor` input.
|
|
|
|
Shape:
|
|
- Input: :math:`(*, S_{\text{dim}}, *)`, where :math:`S_{\text{dim}}` is the size at
|
|
dimension :attr:`dim` and :math:`*` means any number of dimensions including none.
|
|
- Output: :math:`(*, U_1, ..., U_n, *)`, where :math:`U` = :attr:`unflattened_size` and
|
|
:math:`\prod_{i=1}^n U_i = S_{\text{dim}}`.
|
|
|
|
Args:
|
|
dim (Union[int, str]): Dimension to be unflattened
|
|
unflattened_size (Union[torch.Size, Tuple, List, NamedShape]): New shape of the unflattened dimension
|
|
|
|
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(input)
|
|
>>> 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(input)
|
|
>>> output.size()
|
|
torch.Size([2, 2, 5, 5])
|
|
>>> # With namedshape (tuple of tuples)
|
|
>>> input = torch.randn(2, 50, names=('N', 'features'))
|
|
>>> unflatten = nn.Unflatten('features', (('C', 2), ('H', 5), ('W', 5)))
|
|
>>> output = unflatten(input)
|
|
>>> output.size()
|
|
torch.Size([2, 2, 5, 5])
|
|
"""
|
|
|
|
NamedShape = Tuple[Tuple[str, int]]
|
|
|
|
__constants__ = ['dim', 'unflattened_size']
|
|
dim: Union[int, str]
|
|
unflattened_size: Union[_size, NamedShape]
|
|
|
|
def __init__(self, dim: Union[int, str], unflattened_size: Union[_size, NamedShape]) -> None:
|
|
super().__init__()
|
|
|
|
if isinstance(dim, int):
|
|
self._require_tuple_int(unflattened_size)
|
|
elif isinstance(dim, str):
|
|
self._require_tuple_tuple(unflattened_size)
|
|
else:
|
|
raise TypeError("invalid argument type for dim parameter")
|
|
|
|
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, " +
|
|
f"but found element of type {type(elem).__name__} at pos {idx}")
|
|
return
|
|
raise TypeError("unflattened_size must be a tuple of tuples, " +
|
|
f"but found type {type(input).__name__}")
|
|
|
|
def _require_tuple_int(self, input):
|
|
if (isinstance(input, (tuple, list))):
|
|
for idx, elem in enumerate(input):
|
|
if not isinstance(elem, int):
|
|
raise TypeError("unflattened_size must be tuple of ints, " +
|
|
f"but found element of type {type(elem).__name__} at pos {idx}")
|
|
return
|
|
raise TypeError(f"unflattened_size must be a tuple of ints, but found type {type(input).__name__}")
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
return input.unflatten(self.dim, self.unflattened_size)
|
|
|
|
def extra_repr(self) -> str:
|
|
return f'dim={self.dim}, unflattened_size={self.unflattened_size}'
|