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.

Neural Network Modules#

SinusoidalPositionalEmbedding([d_emb, tau, ...])

Generates sinusoidal positional embeddings for input tensors.