File size: 3,015 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
from typing import Tuple, Union
from torch import Tensor, Size
def fold_batch(x: Tensor, nonbatch_ndims: int = 1) -> Tuple[Tensor, Size]:
"""
Overview:
:math:`(T, B, X) \leftarrow (T*B, X)`\
Fold the first (ndim - nonbatch_ndims) dimensions of a tensor as batch dimension.\
This operation is similar to `torch.flatten` but provides an inverse function
`unfold_batch` to restore the folded dimensions.
Arguments:
- x (:obj:`torch.Tensor`): the tensor to fold
- nonbatch_ndims (:obj:`int`): the number of dimensions that is not folded as
batch dimension.
Returns:
- x (:obj:`torch.Tensor`): the folded tensor
- batch_dims: the folded dimensions of the original tensor, which can be used to
reverse the operation
Examples:
>>> x = torch.ones(10, 20, 5, 4, 8)
>>> x, batch_dim = fold_batch(x, 2)
>>> x.shape == (1000, 4, 8)
>>> batch_dim == (10, 20, 5)
"""
if nonbatch_ndims > 0:
batch_dims = x.shape[:-nonbatch_ndims]
x = x.view(-1, *(x.shape[-nonbatch_ndims:]))
return x, batch_dims
else:
batch_dims = x.shape
x = x.view(-1)
return x, batch_dims
def unfold_batch(x: Tensor, batch_dims: Union[Size, Tuple]) -> Tensor:
"""
Overview:
Unfold the batch dimension of a tensor.
Arguments:
- x (:obj:`torch.Tensor`): the tensor to unfold
- batch_dims (:obj:`torch.Size`): the dimensions that are folded
Returns:
- x (:obj:`torch.Tensor`): the original unfolded tensor
Examples:
>>> x = torch.ones(10, 20, 5, 4, 8)
>>> x, batch_dim = fold_batch(x, 2)
>>> x.shape == (1000, 4, 8)
>>> batch_dim == (10, 20, 5)
>>> x = unfold_batch(x, batch_dim)
>>> x.shape == (10, 20, 5, 4, 8)
"""
return x.view(*batch_dims, *x.shape[1:])
def unsqueeze_repeat(x: Tensor, repeat_times: int, unsqueeze_dim: int = 0) -> Tensor:
"""
Overview:
Squeeze the tensor on `unsqueeze_dim` and then repeat in this dimension for `repeat_times` times.\
This is useful for preproprocessing the input to an model ensemble.
Arguments:
- x (:obj:`torch.Tensor`): the tensor to squeeze and repeat
- repeat_times (:obj:`int`): the times that the tensor is repeatd
- unsqueeze_dim (:obj:`int`): the unsqueezed dimension
Returns:
- x (:obj:`torch.Tensor`): the unsqueezed and repeated tensor
Examples:
>>> x = torch.ones(64, 6)
>>> x = unsqueeze_repeat(x, 4)
>>> x.shape == (4, 64, 6)
>>> x = torch.ones(64, 6)
>>> x = unsqueeze_repeat(x, 4, -1)
>>> x.shape == (64, 6, 4)
"""
assert -1 <= unsqueeze_dim <= len(x.shape), f'unsqueeze_dim should be from {-1} to {len(x.shape)}'
x = x.unsqueeze(unsqueeze_dim)
repeats = [1] * len(x.shape)
repeats[unsqueeze_dim] *= repeat_times
return x.repeat(*repeats)
|