model.denoising.data package

Submodules

model.denoising.data.data_samplers module

Dataloaders.

class model.denoising.data.data_samplers.MegatronPretrainingRandomSampler(total_samples, consumed_samples, micro_batch_size, data_parallel_rank, data_parallel_size)

Bases: object

class model.denoising.data.data_samplers.MegatronPretrainingSampler(total_samples, consumed_samples, micro_batch_size, data_parallel_rank, data_parallel_size, drop_last=True)

Bases: object

get_start_end_idx()

model.denoising.data.data_samplers.build_pretraining_data_loader(dataset, consumed_samples)

Buld dataloader given an input dataset.

model.denoising.data.dataset_utils module

model.denoising.data.dataset_utils.build_train_valid_test_datasets(data_prefix, seq_length, segment_length)

model.denoising.data.gw_dataset module

gravitational waveform dataset.

class model.denoising.data.gw_dataset.GwDataset(*args: Any, **kwargs: Any)

Bases: Dataset

model.denoising.data.gw_dataset.get_samples(data_prefix, name)

Module contents