PyTorch Datasets¶
Caution
Lhotse datasets are still very much in the works and are subject to breaking changes.
We supply subclasses of the torch.data.Dataset
for various audio/speech tasks.
These datasets are created from CutSet
objects and load the features from disk into memory on-the-fly.
Each dataset accepts an optional root_dir
argument which is used as a prefix for the paths to features and audio.
Currently, we provide the following:
-
class
lhotse.dataset.diarization.
DiarizationDataset
(cuts, min_speaker_dim=None, global_speaker_ids=False) A PyTorch Dataset for the speaker diarization task. Our assumptions about speaker diarization are the following:
- we assume a single channel input (for now), which could be either a true mono signal
or a beamforming result from a microphone array.
- we assume that the supervision used for model training is a speech activity matrix, with one
row dedicated to each speaker (either in the current cut or the whole dataset, depending on the settings). The columns correspond to feature frames. Each row is effectively a Voice Activity Detection supervision for a single speaker. This setup is somewhat inspired by the TS-VAD paper: https://arxiv.org/abs/2005.07272
Each item in this dataset is a dict of:
{ 'features': (T x F) tensor 'speaker_activity': (num_speaker x T) tensor }
Constructor arguments:
- Parameters
cuts (
CutSet
) – aCutSet
used to create the dataset object.min_speaker_dim (
Optional
[int
]) – optional int, when specified it will enforce that the matrix shape is at least that value (useful for datasets like CHiME 6 where the number of speakers is always 4, but some cuts might have less speakers than that).global_speaker_ids (
bool
) – a bool, indicates whether the same speaker should always retain the same row index in the speaker activity matrix (useful for speaker-dependent systems)root_dir – a prefix path to be attached to the feature files paths.
-
__init__
(cuts, min_speaker_dim=None, global_speaker_ids=False) Initialize self. See help(type(self)) for accurate signature.
-
class
lhotse.dataset.unsupervised.
UnsupervisedDataset
(cuts) Dataset that contains no supervision - it only provides the features extracted from recordings. The returned features are a
torch.Tensor
of shape(T x F)
, where T is the number of frames, and F is the feature dimension.-
__init__
(cuts) Initialize self. See help(type(self)) for accurate signature.
-
-
class
lhotse.dataset.unsupervised.
UnsupervisedWaveformDataset
(cuts) A variant of UnsupervisedDataset that provides waveform samples instead of features. The output is a tensor of shape (C, T), with C being the number of channels and T the number of audio samples. In this implemenation, there will always be a single channel.
-
class
lhotse.dataset.unsupervised.
DynamicUnsupervisedDataset
(feature_extractor, cuts, augment_fn=None) An example dataset that shows how to use on-the-fly feature extraction in Lhotse. It accepts two additional inputs - a FeatureExtractor and an optional WavAugmenter for time-domain data augmentation.. The output is approximately the same as that of the
UnsupervisedDataset
- there might be slight differences forMixedCut``s, because this dataset mixes them in the time domain, and ``UnsupervisedDataset
does that in the feature domain. Cuts that are not mixed will yield identical results in both dataset classes.-
__init__
(feature_extractor, cuts, augment_fn=None) Initialize self. See help(type(self)) for accurate signature.
-
-
class
lhotse.dataset.speech_recognition.
SpeechRecognitionDataset
(cuts) The PyTorch Dataset for the speech recognition task. Each item in this dataset is a dict of:
{ 'features': (T x F) tensor, 'text': string, 'supervisions_mask': (T) tensor }
The
supervisions_mask
field is a mask that specifies which frames are covered by a supervision by assigning a value of 1 (in this case: segments with transcribed speech contents), and which are not by asigning a value of 0 (in this case: padding, contextual noise, or in general the acoustic context without transcription).In the future, will be extended by graph supervisions.
-
__init__
(cuts) Initialize self. See help(type(self)) for accurate signature.
-
-
class
lhotse.dataset.speech_recognition.
K2SpeechRecognitionIterableDataset
(cuts, max_frames=26000, max_cuts=None, shuffle=False, concat_cuts=True, concat_cuts_gap=1.0, concat_cuts_duration_factor=2) The PyTorch Dataset for the speech recognition task using K2 library.
This dataset internally batches and collates the Cuts and should be used with PyTorch DataLoader with argument batch_size=None to work properly. The batch size is determined automatically to satisfy the constraints of
max_frames
andmax_cuts
.This dataset will automatically partition itself when used with a multiprocessing DataLoader (i.e. the same cut will not appear twice in the same epoch).
By default, we “pack” the batches to minimize the amount of padding - we achieve that by concatenating the cuts’ feature matrices with a small amount of silence (padding) in between.
Each item in this dataset is a dict of:
{ 'features': float tensor of shape (B, T, F) 'supervisions': [ { 'cut_id': List[str] of len S 'sequence_idx': Tensor[int] of shape (S,) 'text': List[str] of len S 'start_frame': Tensor[int] of shape (S,) 'num_frames': Tensor[int] of shape (S,) } ] }
Dimension symbols legend: *
B
- batch size (number of Cuts) *S
- number of supervision segments (greater or equal to B, as each Cut may have multiple supervisions) *T
- number of frames of the longest Cut *F
- number of featuresThe ‘sequence_idx’ field is the index of the Cut used to create the example in the Dataset.
-
__init__
(cuts, max_frames=26000, max_cuts=None, shuffle=False, concat_cuts=True, concat_cuts_gap=1.0, concat_cuts_duration_factor=2) K2 ASR IterableDataset constructor.
- Parameters
cuts (
CutSet
) – theCutSet
to sample data from.max_frames (
int
) – The maximum number of feature frames that we’re going to put in a single batch. The padding frames do not contribute to that limit, since we pack the batch by default to minimze the amount of padding.max_cuts (
Optional
[int
]) – The maximum number of cuts sampled to form a mini-batch. By default, this constraint is off.shuffle (
bool
) – WhenTrue
, the cuts will be shuffled at the start of iteration. Convenient when mini-batch loop is inside an outer epoch-level loop, e.g.: for epoch in range(10): for batch in dataset: … as every epoch will see a different cuts order.concat_cuts (
bool
) – WhenTrue
, we will concatenate the cuts to minimize the total amount of padding; e.g. instead of creating a batch with 40 examples, we will merge some of the examples together adding some silence between them to avoid a large number of padding frames that waste the computation. Enabled by default.concat_cuts_gap (
float
) – The duration of silence in seconds that is inserted between the cuts; it’s goal is to let the model “know” that there are separate utterances in a single example.concat_cuts_duration_factor (
float
) – Determines the maximum duration of the concatenated cuts; by default it’s twice the duration of the longest cut in the batch.
-
-
lhotse.dataset.speech_recognition.
concat_cuts
(cuts, gap=1.0, max_duration=None) We’re going to concatenate the cuts to minimize the amount of total padding frames used. This is actually solving a knapsack problem. In this initial implementation we’re using a greedy approach: going from the back (i.e. the shortest cuts) we’ll try to concat them to the longest cut that still has some “space” at the end.
- Parameters
cuts (
List
[Union
[Cut
,MixedCut
,PaddingCut
]]) – a list of cuts to pack.gap (
float
) – the duration of silence inserted between concatenated cuts.max_duration (
Optional
[float
]) – the maximum duration for the concatenated cuts (by default set to the duration of the first cut).
:return a list of packed cuts.
- Return type
List
[Union
[Cut
,MixedCut
,PaddingCut
]]
-
class
lhotse.dataset.speech_recognition.
K2SpeechRecognitionDataset
(cuts) The PyTorch Dataset for the speech recognition task using K2 library. Each item in this dataset is a dict of:
{ 'features': (T x F) tensor, 'supervisions': List[Dict] -> [ { 'sequence_idx': int 'text': string, 'start_frame': int, 'num_frames': int } (multiplied N times, for each of the N supervisions present in the Cut) ] }
The ‘sequence_idx’ field is the index of the Cut used to create the example in the Dataset. It is mapped to the batch index later in the DataLoader.
-
__init__
(cuts) Initialize self. See help(type(self)) for accurate signature.
-
-
class
lhotse.dataset.speech_recognition.
K2DataLoader
(*args, **kwds) A PyTorch DataLoader that has a custom collate_fn that complements the K2SpeechRecognitionDataset.
The ‘features’ tensor is collated in a standard way to return a tensor of shape (B, T, F).
The ‘supervisions’ dict contains the same fields as in
K2SpeechRecognitionDataset
, except that each sub-field (like ‘start_frame’) is a 1D PyTorch tensor with shape (B,). The ‘text’ sub-field is an exception - it’s a list of strings with length equal to batch size.The ‘sequence_idx’ sub-field in ‘supervisions’, which originally points to index of the example in the Dataset, is remapped to the index of the corresponding features matrix in the collated ‘features’. Multiple supervisions coming from the same cut will share the same ‘sequence_idx’.
For an example, see
test/dataset/test_speech_recognition_dataset.py::test_k2_dataloader()
.-
__init__
(*args, **kwargs) Initialize self. See help(type(self)) for accurate signature.
-
dataset
-
batch_size
-
num_workers
-
pin_memory
-
drop_last
-
timeout
-
sampler
-
prefetch_factor
-
-
lhotse.dataset.speech_recognition.
multi_supervision_collate_fn
(batch) Custom collate_fn for K2SpeechRecognitionDataset.
It merges the items provided by K2SpeechRecognitionDataset into the following structure:
{ 'features': float tensor of shape (B, T, F) 'supervisions': [ { 'sequence_idx': Tensor[int] of shape (S,) 'text': List[str] of len S 'start_frame': Tensor[int] of shape (S,) 'num_frames': Tensor[int] of shape (S,) } ] }
Dimension symbols legend: *
B
- batch size (number of Cuts), *S
- number of supervision segments (greater or equal to B, as each Cut may have multiple supervisions), *T
- number of frames of the longest Cut *F
- number of features- Return type
Dict
-
lhotse.dataset.
speech_synthesis
alias of
lhotse.dataset.speech_synthesis
-
class
lhotse.dataset.source_separation.
DynamicallyMixedSourceSeparationDataset
(sources_set, mixtures_set, nonsources_set=None) A PyTorch Dataset for the source separation task. It’s created from a number of CutSets:
sources_set
: provides the audio cuts for the sources that (the targets of source separation),mixtures_set
: provides the audio cuts for the signal mix (the input of source separation),nonsources_set
: (optional) provides the audio cuts for other signals that are in the mix, but are not the targets of source separation. Useful for adding noise.
When queried for data samples, it returns a dict of:
{ 'sources': (N x T x F) tensor, 'mixture': (T x F) tensor, 'real_mask': (N x T x F) tensor, 'binary_mask': (T x F) tensor }
This Dataset performs on-the-fly feature-domain mixing of the sources. It expects the mixtures_set to contain MixedCuts, so that it knows which Cuts should be mixed together.
-
__init__
(sources_set, mixtures_set, nonsources_set=None) Initialize self. See help(type(self)) for accurate signature.
-
class
lhotse.dataset.source_separation.
PreMixedSourceSeparationDataset
(sources_set, mixtures_set) A PyTorch Dataset for the source separation task. It’s created from two CutSets - one provides the audio cuts for the sources, and the other one the audio cuts for the signal mix. When queried for data samples, it returns a dict of:
{ 'sources': (N x T x F) tensor, 'mixture': (T x F) tensor, 'real_mask': (N x T x F) tensor, 'binary_mask': (T x F) tensor }
It expects both CutSets to return regular Cuts, meaning that the signals were mixed in the time domain. In contrast to DynamicallyMixedSourceSeparationDataset, no on-the-fly feature-domain-mixing is performed.
-
__init__
(sources_set, mixtures_set) Initialize self. See help(type(self)) for accurate signature.
-
-
class
lhotse.dataset.vad.
VadDataset
(cuts) The PyTorch Dataset for the voice activity detection task. Each item in this dataset is a dict of:
{ 'features': (T x F) tensor 'is_voice': (T x 1) tensor }
-
__init__
(cuts) Initialize self. See help(type(self)) for accurate signature.
-