Shortcuts

Source code for ding.torch_utils.reshape_helper

from typing import Tuple, Union

from torch import Tensor, Size


[docs]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
[docs]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:])
[docs]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)