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.

scatter_dataset(dataset[, permute_fn])

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

Functional

linear_sum_assignment(cost_matrix[, maximize])

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

Learning Rate Schedulers

SequentialLR(optimizer, schedulers_fn, ...)

Build a sequential learning rate scheduler from scheduler factory functions.

Inference Pipeline Helpers

BasePipeline([allow_tf32, unk_args])

Base class for inference pipelines.

reorder_fields(cls, fields)

Reorder attrs fields such that fields without default values come first, then fields with default values.

Lightning Utilities

OptimizerLightningModule(optimizer_config)

LightningModule subclass for models that use custom optimizers and schedulers.

LRSchedulerConfig(scheduler_generator[, ...])

Configuration for a learning rate scheduler in Lightning.

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.

load_checkpoint(model_path[, config_name, ...])

Load a PyTorch Lightning model from a pre-trained checkpoint.

ABCIEnvironment()

Environment class for ABCI.

Lightning Datamodules

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

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

Lightning Callbacks

SaveMetricCallback(metric_name, output_path)

Lightning Callback for save metric in fit ends.

LoadPretrainedCallback(model_path, ...[, ...])

Initialize a model from a pretrained checkpoint before training or validation.

NoBufferWeightAveraging([device])

Weight averaging callback that ignores buffers during averaging and swapping.

NoBufferEMAWeightAveraging([device, decay, ...])

Exponential moving average (EMA) callback that ignores buffers.

PrintUnusedParam()

Warn once when trainable parameters do not receive gradients.

H5py Utilities

HDF5Writer()

Abstract base class for writing data to an HDF5 file.