PyTorch/Lightning Toolkit

Datasets

CachedDataset(dataset)	A dataset wrapper that caches the samples to improve performance.
FileCachedDataset(dataset, cache_path)	A dataset wrapper that caches samples to disk to reduce memory usage.
HDF5Dataset(dataset_path[, grp_list])	A dataset class for loading data from an HDF5 file.
RawHDF5Dataset(dataset_path[, grp_list])	A dataset class for reading data from HDF5 files.

Dataset Utilities

scatter_dataset(dataset[, permute_fn])

Splits a dataset into subsets and returns the subset corresponding to the current process rank.

Lightning Utilities

ABCIEnvironment()	Environment class for ABCI.
OptimizerLightningModule(optimizer_config)	LightningModule subclass for models that use custom optimizers and schedulers.
OptimizerConfig(optimizer_generator[, ...])	Configuration for the optimizer and scheduler in a LightningModule.
build_param_groups(named_params, groups)	Build parameter groups for the optimizer based on the provided patterns.

Lightning Datamodules

SingleDataModule(train_dataset_fn, ...[, ...])

A PyTorch Lightning DataModule designed to handle training and validation datasets with support for caching and dataset scattering.

H5py Utilities

HDF5Writer()

Abstract base class for writing data to an HDF5 file.

Functional

linear_sum_assignment(cost_matrix[, maximize])

Solve the linear sum assignment problem for a batch of cost matrices.