Torch:

Code related to using pytorch with neural data from aolab.

Note: pytorch must be installed to use the functions in this module using pip install torch.

API

class aopy.torch.TensorDataset(*args: Any, **kwargs: Any)[source]

Dataset wrapping tensors. Implements the pytorch Dataset parent class.

Each sample will be retrieved by indexing tensors along the first dimension. These samples are collected and returned in a list of tensors in the same order as their source tensors.

Parameters:

  • *tensors (Tensor) – pytorch tensors that have the same size of the first dimension. This first dimension indexes over individual trials. Standard convention has the 2nd dimension index over time samples while extra dimensions index over channels or spatial dimentions.

  • device (str) – Memory location to place input tensors. (default: ‘cpu’)

  • transform (torch.nn.Module) – Transform class for augmenting or transforming tensor samples. (default: None)

  • transform_mast ([bool]) – List of boolean values indicating which input tensors are transformed when sampled. (default: None)

Returns:

List of pytorch Tensors sampled from an input index (int).

Return type:

sample ([Tensor])

aopy.torch.recursive_assign_device(x, device: str)[source]

Recursively assign tensor elements in a nested list or tuple of arbitrary depth to a specified device memory location.

Parameters:

  • x ([Tensor,[Tensor],...]) – List of tensors or lists of tensors. Can be arbitrarily deep.

  • device (str) – Memory location on the current system (ex: ‘cpu’, ‘cuda:0’)

Returns:

Copy of input x, memory relocated to designated location.

Return type:

[Tensor,[Tensor],…]