dataloading

Dataloaders for the neural networks. Uses pytorch dataloaders and converts to jax format

env: XLA_PYTHON_CLIENT_ALLOCATOR=platform

Dummy Dataset

source

DummyDataset

 DummyDataset (n_samples=100, n_features=5, n_targets=1, seed=42)

Dummy dataset for testing

	Type	Default	Details
n_samples	int	100	number of samples
n_features	int	5	number of features
n_targets	int	1	number of targets
seed	int	42	random seed

source

NumpyLoader

 NumpyLoader (dataset, batch_size=1, shuffle=False, sampler=None,
              batch_sampler=None, num_workers=0, pin_memory=False,
              drop_last=False, timeout=0, worker_init_fn=None)

A dataloader that uses numpy_collate to allow numpy arrays instead of torch tensors

source

numpy_collate

 numpy_collate (batch)

dummy_dataset = DummyDataset(n_samples=250, n_features=4, n_targets=4, seed=42)
dataloader = NumpyLoader(dummy_dataset, batch_size=32, shuffle=True)

for x, y in dataloader:
    print(x.shape, y.shape)
    break

test_eq(x.shape, (32, 4))
test_eq(y.shape, (32, 4))

(32, 4) (32, 4)

Real Pendulum

source

RealPendulumDataset

 RealPendulumDataset (root_path:str, train_split:float, train:bool=True)

An abstract class representing a :class:Dataset.

All datasets that represent a map from keys to data samples should subclass it. All subclasses should overwrite :meth:__getitem__, supporting fetching a data sample for a given key. Subclasses could also optionally overwrite :meth:__len__, which is expected to return the size of the dataset by many :class:~torch.utils.data.Sampler implementations and the default options of :class:~torch.utils.data.DataLoader.

.. note:: :class:~torch.utils.data.DataLoader by default constructs a index sampler that yields integral indices. To make it work with a map-style dataset with non-integral indices/keys, a custom sampler must be provided.

	Type	Default	Details
root_path	str		path to data
train_split	float		fraction of data to use for training
train	bool	True	whether to use train or test data

Henon Heiles

source

HHDataset

 HHDataset (root_path, train)

An abstract class representing a :class:Dataset.

	Details
root_path
train	path to data # whether to use train or test data

Swinging Sticks

source

SwingingSticksDataset

 SwingingSticksDataset (root_path, train_split, train=True)

An abstract class representing a :class:Dataset.

	Type	Default	Details
root_path			path to data
train_split			fraction of data to use for training
train	bool	True	whether to use train or test data