speechbrain.dataio.sampler module¶
PyTorch compatible samplers.
These determine the order of iteration through a dataset.
- Authors:
Aku Rouhe 2020
Samuele Cornell 2020
Ralf Leibold 2020
Summary¶
Classes:
This sampler is built to work with a standard Pytorch ConcatDataset. |
|
This wrapper allows using any sampler with Distributed Data Parallel (DDP) correctly. |
|
This BatchSampler batches examples together by grouping them by their length. |
|
A modification of RandomSampler which always returns the same values. |
|
A reproducible modification of WeightedRandomSampler. |
Reference¶
- class speechbrain.dataio.sampler.ReproducibleRandomSampler(data_source, seed=563375142, epoch=0, **kwargs)[source]¶
Bases:
torch.utils.data.sampler.Sampler[int]A modification of RandomSampler which always returns the same values.
Also look at torch.utils.data.RandomSampler. This has mostly the same behaviour and arguments, except for adding ‘seed’ and ‘epoch’ and not supporting ‘generator’.
Note
Call set_epoch before every epoch. Otherwise, the sampler will produce the same sequence of indices every epoch.
- Parameters
Example
>>> import torch >>> from speechbrain.utils.checkpoints import Checkpointer >>> from speechbrain.dataio.dataloader import SaveableDataLoader >>> # An example "dataset" >>> dataset = torch.arange(10).unsqueeze(1) >>> # Create the random sampler: >>> sampler = ReproducibleRandomSampler(dataset) >>> dataloader = SaveableDataLoader(dataset, sampler = sampler, ... num_workers = 3) >>> # Setup the checkpointer. >>> # Note that the sampler doesn't need to be saved itself. >>> tmpdir = getfixture('tmpdir') >>> checkpointer = Checkpointer(tmpdir, {"dataloader": dataloader}) >>> # Iterate: >>> subset = [] >>> for i, data_point in enumerate(dataloader): ... # Say you save a checkpoint on the fourth batch: ... if i == 3: ... _ = checkpointer.save_checkpoint(end_of_epoch = False) ... # So let's save the numbers you would get if you continue ... if i >= 4: ... subset.append(data_point.item()) >>> # What if instead you had to restart the experiment? >>> new_sampler = ReproducibleRandomSampler(dataset) >>> new_dataloader = SaveableDataLoader(dataset, sampler = new_sampler, ... num_workers = 3) >>> new_checkpointer = Checkpointer(tmpdir, {"dataloader": new_dataloader}) >>> _ = new_checkpointer.recover_if_possible() >>> # You'll get the same random order again: >>> new_subset = [data_point.item() for data_point in new_dataloader] >>> assert subset == new_subset
- set_epoch(epoch)[source]¶
You can also just access self.epoch, but we maintain this interface to mirror torch.utils.data.distributed.DistributedSampler
- data_source: Sized¶
- class speechbrain.dataio.sampler.ReproducibleWeightedRandomSampler(weights, num_samples, replacement, seed=129491412, epoch=0, **kwargs)[source]¶
Bases:
torch.utils.data.sampler.Sampler[int]A reproducible modification of WeightedRandomSampler.
Also look at torch.utils.data.WeightedRandomSampler. This has the the same behaviour and arguments, except for adding ‘seed’ and ‘epoch’ and not supporting ‘generator’.
Note
Call set_epoch before every epoch. Otherwise, the sampler will produce the same sequence of indices every epoch.
- Parameters
weights (sequence of float) – Weights for each index. Doesn’t need to sum to one.
num_samples (int) – Number of samples to draw
replacement (bool) – To draw with replacement or not (within an epoch of num_samples).
seed (int) – The base seed to use for the random number generator. It is recommended to use a value which has a good mix of 0 and 1 bits.
epoch (int) – The epoch to start at.
Example
>>> a = ReproducibleWeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True) >>> b = ReproducibleWeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True) >>> list(a) [3, 1, 4, 4, 4] >>> list(b) [3, 1, 4, 4, 4] >>> a.set_epoch(1) >>> list(a) [4, 5, 4, 4, 3] >>> b.set_epoch(1) >>> list(b) [4, 5, 4, 4, 3]
- set_epoch(epoch)[source]¶
You can also just access self.epoch, but we maintain this interface to mirror torch.utils.data.distributed.DistributedSampler
- weights: torch.Tensor¶
- class speechbrain.dataio.sampler.ConcatDatasetBatchSampler(samplers, batch_sizes: (<class 'tuple'>, <class 'list'>), epoch=0)[source]¶
Bases:
Generic[torch.utils.data.sampler.T_co]This sampler is built to work with a standard Pytorch ConcatDataset.
It is used to retrieve elements from the different concatenated datasets placing them in the same batch with proportion specified by batch_sizes, e.g 8, 16 means each batch will be of 24 elements with the first 8 belonging to the first dataset in ConcatDataset object and the last 16 to the second. More than two datasets are supported, in that case you need to provide 3 batch sizes.
Note
Batched are drawn from the datasets till the one with smallest length is exhausted. Thus number of examples in your training epoch is dictated by the dataset whose length is the smallest.
- Parameters
Example
>>> import torch >>> from speechbrain.dataio.sampler import ConcatDatasetBatchSampler, ReproducibleRandomSampler >>> from speechbrain.dataio.sampler import ReproducibleRandomSampler >>> from speechbrain.dataio.dataloader import SaveableDataLoader >>> # example "datasets" >>> dataset1 = torch.arange(0, 10).unsqueeze(1) >>> dataset2 = torch.arange(20, 40).unsqueeze(1) >>> tot_dataset = torch.utils.data.ConcatDataset([dataset1, dataset2]) >>> sampler1 = ReproducibleRandomSampler(dataset1) >>> sampler2 = ReproducibleRandomSampler(dataset2) >>> tot_sampler = ConcatDatasetBatchSampler([sampler1, sampler2], [2, 4]) >>> dataloader = SaveableDataLoader(tot_dataset, batch_sampler = tot_sampler, ... num_workers = 3) >>> for data_point in dataloader: ... assert len(data_point) == 6 ... for i in range(2): ... assert data_point[i] in [x for x in range(0, 10)] ... for i in range(2, 4): ... assert data_point[i] in [x for x in range(10, 40)]
- class speechbrain.dataio.sampler.DynamicBatchSampler(dataset, max_batch_length: int, left_bucket_length: int, bucket_length_multiplier: float = 1.1, length_func=<function DynamicBatchSampler.<lambda>>, shuffle: bool = True, bucket_boundaries: List[int] = [], lengths_list: Optional[List[int]] = None, seed: int = 42, epoch: int = 0, drop_last: bool = False)[source]¶
Bases:
Generic[torch.utils.data.sampler.T_co]This BatchSampler batches examples together by grouping them by their length.
Every example in the batch have approximately the same length and thus padding is minimized. This enables faster training on datasets where length of examples can vary significantly (e.g Librispeech). Inspired by: https://www.tensorflow.org/api_docs/python/tf/data/experimental/bucket_by_sequence_length
Dynamic batching is performed by specifying a max_batch_length which is the upper limit for the sum of the length of examples in a batch: e.g., if ex1 has length 4, ex2 length 5 andn if max_batch_length is set to 6 ex1 and ex2 will be placed, alone, in two distinct batches.
Length for each example can be obtained in two manners. If the input dataset is a DynamicItemDataset it can be obtained by specifying a length_func. Default assumes a “duration” entry is in the annotation. Length for each example can also be passed to this class upon instantiation by specifying a list containing the length for each example and passing it to lengths_list.
Examples are grouped together by defining a set of possible discrete intervals (buckets) multiple of a left_bucket_length. A bucket_length_multiplier is used to specify the number of possible buckets. E.g., if max_batch_length = 32 and left_bucket_length = 10, bucket_length_multiplier = 2 there will be 3 buckets: [0, 10), [10, 20), [20, 40). A common choice would be setting left_bucket_length to approximately the length of your shortest example in the dataset. Decreasing bucket_length_multiplier creates more buckets in the whole interval of [left_bucket_length, max_batch_size]: e.g. if max_batch_length = 32 and left_bucket_length = 10, bucket_length_multiplier = 1.5 the number of buckets increases to 8. With right boundaries: [10 12 14 17 21 25 30 36]. Thus examples with length less than 10 are all grouped together but more buckets are created for longer examples. Note that the bucket boundary grows exponentially using the multiplier.
The buckets can also be specified by passing a list to the bucket_boundaries argument instead of specifying a left_bucket_length and a bucket_length_multiplier.
Example
>>> import torch >>> import speechbrain as sb >>> from speechbrain.dataio.sampler import DynamicBatchSampler >>> from speechbrain.dataio.dataset import DynamicItemDataset >>> from speechbrain.dataio.dataloader import SaveableDataLoader >>> from speechbrain.dataio.batch import PaddedBatch >>> import numpy as np >>> item_lengths = sorted([np.random.randint(10, 100) for x in range(20)]) >>> dataset = {"ex_{}".format(x) : {"wav" :torch.randn(x)} for x in item_lengths} >>> dataset = DynamicItemDataset(dataset) >>> dataset.set_output_keys(["wav"]) >>> length_func = lambda x : len(x) # trivial in this example >>> bsampler = DynamicBatchSampler(dataset, 20, 10, 1.1, length_func, shuffle=False) >>> dataloader = SaveableDataLoader(dataset, batch_sampler=bsampler, collate_fn=PaddedBatch) >>> for i, b in enumerate(dataloader): ... data, length = b["wav"] >>> assert data.shape[-1] == max(item_lengths)
- Parameters
dataset (torch.utils.data.Dataset) – Pytorch Dataset from which elements will be sampled.
max_batch_length (int) – Upper limit for the sum of the length of examples in a batch. Should be chosen based on your GPU memory.
left_bucket_length (int) – Minimum length of a bucket. Specifies resolution of buckets and thus this sampler stochasticity. A common choice is to set this to length of your shortest example.
bucket_length_multiplier (float) – Multiplier for bucket length, specifies number of buckets from left_bucket_length to max_batch_length.
length_func (callable) – Function used to get length of each example from the dataset. This argument can be used only when the dataset is a Speechbrain DynamicItemDataset object. Can be anything: e.g. lambda x: x[“duration”]*16000 returns number of samples if duration key in the annotation is in seconds and the file has 16kHz sampling freq.
shuffle (bool) – Whether or not shuffle examples between each epoch.
bucket_boundaries (list) – Overrides bucket_length_multiplier and left_bucket_length by specifying manually the buckets right boundaries.
lengths_list (list) – Overrides length_func by passing a list containing the length of each example in the dataset. This argument must be set when the dataset is a plain Pytorch Dataset object and not a DynamicItemDataset object as length_func cannot be used on Pytorch Datasets.
epoch (int) – The epoch to start at.
drop_last (bool) – If
True, the sampler will drop the last examples which have not been grouped.
- class speechbrain.dataio.sampler.DistributedSamplerWrapper(sampler, *args, **kwargs)[source]¶
Bases:
torch.utils.data.sampler.Sampler[torch.utils.data.distributed.T_co]This wrapper allows using any sampler with Distributed Data Parallel (DDP) correctly.
Passing blindly the sampler to each DDP process will cause to have access within each process to all the data in the dataset instead of only a subset of it which is unique to each process. This wrapper prevents this and allows to use only a subset of the original data for each process.
Note
This is is automatically applied to any sampler in the Brain class when DDP training is used.