Module

torch_utils.py

import goldener.torch_utils
View source

Utilities to work with PyTorch tensors and datasets.

This module gathers helpers to collate batches, convert tensors from/to numpy vectors and access the samples of datasets, as well as ResetableTorchIterableDataset, an iterable dataset that can be re-initialized.

On this page 1 class8 functions
collate_keeping_sequences_as_sequences() make_2d_tensor() torch_tensor_to_numpy_vectors() numpy_vectors_to_torch_tensor() np_transform_from_torch() ResetableTorchIterableDataset get_dataset_sample_dict() get_unique_values_in_tensor() shuffle_list()
logger = <Logger goldener.torch_utils (INFO)>
def collate_keeping_sequences_as_sequences(batch: list[dict[str, typing.Any]]) -> dict[str, typing.Any]:

Collate dictionary samples while preserving sequence-valued features.

Arguments:
  • batch: A list of samples, where each sample is a dictionary.
Returns:

A dictionary with sequences kept per sample and all other values collated with PyTorch's default collate function.

Raises:
  • KeyError: If a sample is missing a key present in the first sample.
def make_2d_tensor(x: torch.Tensor) -> torch.Tensor:

Convert a torch tensor to a 2D tensor.

If the input tensor is 0D, it is converted to a (1, 1) tensor. If the input tensor is 1D, it is converted to a (N, 1) tensor. If the input tensor is 2D, it is kept as (N, M) tensor. If the input tensor is 3D or higher, the second dimension is moved to the last dimension, and the first dimensions are flattened to form a 2D tensor.

Arguments:
  • x: Input tensor convert to 2D tensor.
Returns:

A 2D tensor of shape (N, M)

def torch_tensor_to_numpy_vectors(x: torch.Tensor) -> numpy.ndarray:

Convert a torch tensor to a numpy array of vectors.

If the input tensor is 0D, it is converted to a (1, 1) array. If the input tensor is 1D, it is converted to a (N, 1) array. If the input tensor is 2D, it is kept as (N, M) array. If the input tensor is 3D or higher, the second dimension is moved to the last dimension, and the first dimensions are flattened to form a 2D array of vectors.

Arguments:
  • x: Input tensor convert to numpy vectors.
Returns:

A 2D numpy array of shape (N, M)

def numpy_vectors_to_torch_tensor( x: numpy.ndarray, shape: tuple[int, ...], dtype: torch.dtype, device: torch.device) -> torch.Tensor:

Convert a numpy array to a torch tensor.

If the input array is 0D, it is converted to a (1, 1) tensor. If the input array is 1D, it is converted to a (N, 1) tensor. If the input array is 2D, it is kept as (N, M) tensor. If the input array is 3D or higher, the last dimension is moved to the second dimension, and the first dimensions are reshaped to form the desired shape.

Arguments:
  • x: Input numpy array to convert to torch tensor.
  • shape: Desired shape of the output tensor.
  • dtype: Desired data type of the output tensor.
  • device: Desired device of the output tensor.
Returns:

A torch tensor of the desired shape, dtype, and device.

Raises:
  • ValueError: If the desired shape is less than 2D.
def np_transform_from_torch( x: torch.Tensor, transform_np: Callable[[numpy.ndarray], numpy.ndarray]) -> torch.Tensor:

Apply a numpy transformation to a torch tensor.

Arguments:
  • x: Input tensor to transform.
  • transform_np: A callable that takes a numpy array and returns a transformed numpy array.
Returns:

A transformed torch tensor with the same dtype and device as the input tensor.

class ResetableTorchIterableDataset(torch.utils.data.dataset.Dataset[+_T_co], typing.Iterable[+_T_co]):

A resetable iterable dataset that can be re-initialized.

It is useful for accessing the first element (to set up a table for instance) of the dataset before looping over the full dataset.

Attributes:
  • data_iterable: The underlying iterable dataset.
  • _data_iterator: The current iterator over the dataset, or None once it has been exhausted (see __next__) and not yet re-created.
ResetableTorchIterableDataset(data_iterable: torch.utils.data.dataset.IterableDataset)

Initialize the resetable iterable dataset.

Arguments:
  • data_iterable: The underlying iterable dataset to wrap.
data_iterable
def reset(self) -> None:

Reset the iterator to the beginning of the dataset.

def get_dataset_sample_dict( dataset: torch.utils.data.dataset.Dataset, collate_fn: Optional[Callable] = None, expected: list[str] | None = None, excluded: list[str] | None = None) -> dict[str, typing.Any]:

Get a sample from a dataset as a dictionary and validate some of its keys.

Arguments:
  • dataset: The dataset to get a sample from.
  • collate_fn: An optional collate function to apply to the sample.
  • expected: An optional list of keys that must be present in the sample.
  • excluded: An optional list of keys that must not be present in the sample.
Returns:

A sample from the dataset as a dictionary.

Raises:
  • ValueError: If the sample is not a dictionary after applying the collate_fn.
  • ValueError: If the sample is missing any of the expected keys.
  • ValueError: If the sample contains any of the excluded keys.
def get_unique_values_in_tensor(tensor: torch.Tensor, dim: int = 1) -> torch.Tensor:

Get the unique values/vectors of a given dimension in a tensor.

Arguments:
  • tensor: The input tensor to get unique values from.
  • dim: The dimension along which to get unique values. Default is 1.
Returns:

A tensor with all the unique values of the input tensor for teh specified dimension.

def shuffle_list(items: list[~T], generator: torch._C.Generator) -> list[~T]:

Shuffle the given items in a list.

Arguments:
  • items: The list of items to shuffle.
  • generator: A torch generator to manage the random shuffling.

Returns: A list of shuffled items.