Example pipeline with wav2letter (#632)

* example pipeline, initial commit.

* removing notebook conversion artifacts.

* remove extra comments. lint.

* addressing some feedback.

* main function.

* defining args in function.

* refactor.

* lint.

* checkpoint.

* clean version to start with.

* adding more parameters.

* lint.

* cleaning full version.

* check for not None.

* cleaning.

* back -l 160

* black.

* fix runtime error.

* removing some print statements.

* add help to command line. add progress bar option.

* grouping librispeech-specific transform in subclass.

* typo.

* fix concatenation.

* typo.

* black. tqdm.

* missing transpose.

* renaming variables.

* sum cer and wer

* clip norm.

* second signal handler removed.

* cosmetic.

* default to no checkpoint.

* remove non_blocking.

* adadelta works better than sgd.

* anomaly detection.

* moving dataset to separate file.

* lint.

* move to separate module: languagemodel, decoder, metric.

* flush=True.

* renaming decoder.

* CTC Decoders.

* flush=True.

* pass length for viterbi decoder.

* progress bar. relative path.

* generalize transition matrix to n-gram. progress bar.

* choice of decoder.

* collate func.

* remove signal handling.

* adding distributed.

* lint.

* normalize w/r to length of dataset, and w/r to total number characters.

* relative cer/wer.

* clip grad parameter. momentum back but not yet used.

* Switch to SGD.

* choice of optimizer.

* scheduler.

* move to utils file.

* metric log, and utils file.

* rename metric_logger.

* stderr and stdout. simpler metric logger.

* replace by logging.

* adding time measurement in metric logger.

* fix duplicate name. remove tqdm. keep track of epoch instead and iteration instead.

* rename main file. and add readme.

* refactor distributed.

* swap example and output in readme.

* remove time from logger.

* check non-empty tensor input.

* typo in variable name and log update.

* typo.

* compute cer/wer in training too.

* typo.

* add back slurm signal capture to resubmit job.

* update levinstein distance.

* adding tests for levenstein distance.

* record error rate during iteration.

* metric logger using setitem.

* moving signal break to end of loop and return loss so far.

* typo.

* add citation.

* change default to best run.

* adding other experiment with decoders.

* remove other decoders than greedy.

* Revert "remove other decoders than greedy."

This reverts commit fb11437.

* changing name of folfder.

* remove other decoders, and unused dataset class.

* rename functions to align with other pipeline.

* pick which parts to train with.

* adding specaugment to validation. note that caching prevents randomization from happening in validation.

* updating readme.

* typo in metric logging.

* Revert "typo in metric logging."

This reverts commit acac245.

* Revert "Revert "typo in metric logging.""

This reverts commit 2c80d96.

* update metric logger.

* simplify metric logger implementation.

* use json dumps instead.

* group metric together.

* move function.

* lint.

* quick summary of files in folder.

* pass clip_grad explictly.

* typo in default dataset name.

* option to disable logger.

* ergonomics for distributed.

* reminder about signal handler.

* minor refactor of main in main.

* replace by not_main_rank.

* raising error if parameter not supported.

* move model before invoking DDP.

* changing log level. using python 2 style string for logging.

* dynamic augmentations.

* update metric log.

batch cer/wer metric. correct typo in time. adding other dimensions in metric.

* save learning rate even if function not available.

* add type option to model.

* add adamw.

* reduce lr on validation step or training step.

* specify hop-length and win-length.

* normalize option.

* rename parameter.

* add dropout and tweak to number of channels.

* copy model in pipeline folder for experimentation.

* fix scheduler stepping.

* fix input_type and num_features.

* waveform mode changes shape more.

* adding best character error rate with current implementation of model with mfcc.

* comment update.

* remove signal. remove custom wav2letter model.

* remove comment.

* simpler import with pandas.

Author: vincentqb

examples/pipeline_wav2letter/README.md

@@ -0,0 +1,45 @@
+This is an example pipeline for speech recognition using a greedy or Viterbi CTC decoder, along with the Wav2Letter model trained on LibriSpeech, see [Wav2Letter: an End-to-End ConvNet-based Speech Recognition System](https://arxiv.org/pdf/1609.03193.pdf). Wav2Letter and LibriSpeech are available in torchaudio.
+
+### Usage
+
+More information about each command line parameters is available with the `--help` option. An example can be invoked as follows.
+```
+python main.py \
+    --reduce-lr-valid \
+    --dataset-train train-clean-100 train-clean-360 train-other-500 \
+    --dataset-valid dev-clean \
+    --batch-size 128 \
+    --learning-rate .6 \
+    --momentum .8 \
+    --weight-decay .00001 \
+    --clip-grad 0. \
+    --gamma .99 \
+    --hop-length 160 \
+    --n-hidden-channels 2000 \
+    --win-length 400 \
+    --n-bins 13 \
+    --normalize \
+    --optimizer adadelta \
+    --scheduler reduceonplateau \
+    --epochs 30
+```
+With these default parameters, we get a character error rate of 13.8% on dev-clean after 30 epochs.
+
+### Output
+
+The information reported at each iteration and epoch (e.g. loss, character error rate, word error rate) is printed to standard output in the form of one json per line, e.g.
+```python
+{"name": "train", "epoch": 0, "cer over target length": 1.0, "cumulative cer": 23317.0, "total chars": 23317.0, "cer": 0.0, "cumulative cer over target length": 0.0, "wer over target length": 1.0, "cumulative wer": 4446.0, "total words": 4446.0, "wer": 0.0, "cumulative wer over target length": 0.0, "lr": 0.6, "batch size": 128, "n_channel": 13, "n_time": 2453, "dataset length": 128.0, "iteration": 1.0, "loss": 8.712121963500977, "cumulative loss": 8.712121963500977, "average loss": 8.712121963500977, "iteration time": 41.46276903152466, "epoch time": 41.46276903152466}
+{"name": "train", "epoch": 0, "cer over target length": 1.0, "cumulative cer": 46005.0, "total chars": 46005.0, "cer": 0.0, "cumulative cer over target length": 0.0, "wer over target length": 1.0, "cumulative wer": 8762.0, "total words": 8762.0, "wer": 0.0, "cumulative wer over target length": 0.0, "lr": 0.6, "batch size": 128, "n_channel": 13, "n_time": 1703, "dataset length": 256.0, "iteration": 2.0, "loss": 8.918599128723145, "cumulative loss": 17.63072109222412, "average loss": 8.81536054611206, "iteration time": 1.2905676364898682, "epoch time": 42.753336668014526}
+{"name": "train", "epoch": 0, "cer over target length": 1.0, "cumulative cer": 70030.0, "total chars": 70030.0, "cer": 0.0, "cumulative cer over target length": 0.0, "wer over target length": 1.0, "cumulative wer": 13348.0, "total words": 13348.0, "wer": 0.0, "cumulative wer over target length": 0.0, "lr": 0.6, "batch size": 128, "n_channel": 13, "n_time": 1713, "dataset length": 384.0, "iteration": 3.0, "loss": 8.550191879272461, "cumulative loss": 26.180912971496582, "average loss": 8.726970990498861, "iteration time": 1.2109291553497314, "epoch time": 43.96426582336426}
+```
+One way to import the output in python with pandas is by saving the standard output to a file, and then using `pandas.read_json(filename, lines=True)`.
+
+## Structure of pipeline
+
+* `main.py` -- the entry point
+* `ctc_decoders.py` -- the greedy CTC decoder
+* `datasets.py` -- the function to split and process librispeech, a collate factory function
+* `languagemodels.py` -- a class to encode and decode strings
+* `metrics.py` -- the levenshtein edit distance
+* `utils.py` -- functions to log metrics, save checkpoint, and count parameters

examples/pipeline_wav2letter/ctc_decoders.py

@@ -0,0 +1,15 @@
+from torch import topk
+
+
+class GreedyDecoder:
+    def __call__(self, outputs):
+        """Greedy Decoder. Returns highest probability of class labels for each timestep
+
+        Args:
+            outputs (torch.Tensor): shape (input length, batch size, number of classes (including blank))
+
+        Returns:
+            torch.Tensor: class labels per time step.
+        """
+        _, indices = topk(outputs, k=1, dim=-1)
+        return indices[..., 0]

examples/pipeline_wav2letter/datasets.py

@@ -0,0 +1,113 @@
+import torch
+from torchaudio.datasets import LIBRISPEECH
+
+
+class MapMemoryCache(torch.utils.data.Dataset):
+    """
+    Wrap a dataset so that, whenever a new item is returned, it is saved to memory.
+    """
+
+    def __init__(self, dataset):
+        self.dataset = dataset
+        self._cache = [None] * len(dataset)
+
+    def __getitem__(self, n):
+        if self._cache[n] is not None:
+            return self._cache[n]
+
+        item = self.dataset[n]
+        self._cache[n] = item
+
+        return item
+
+    def __len__(self):
+        return len(self.dataset)
+
+
+class Processed(torch.utils.data.Dataset):
+    def __init__(self, dataset, transforms, encode):
+        self.dataset = dataset
+        self.transforms = transforms
+        self.encode = encode
+
+    def __getitem__(self, key):
+        item = self.dataset[key]
+        return self.process_datapoint(item)
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def process_datapoint(self, item):
+        transformed = item[0]
+        target = item[2].lower()
+
+        transformed = self.transforms(transformed)
+        transformed = transformed[0, ...].transpose(0, -1)
+
+        target = self.encode(target)
+        target = torch.tensor(target, dtype=torch.long, device=transformed.device)
+
+        return transformed, target
+
+
+def split_process_librispeech(
+    datasets, transforms, language_model, root, folder_in_archive,
+):
+    def create(tags, cache=True):
+
+        if isinstance(tags, str):
+            tags = [tags]
+        if isinstance(transforms, list):
+            transform_list = transforms
+        else:
+            transform_list = [transforms]
+
+        data = torch.utils.data.ConcatDataset(
+            [
+                Processed(
+                    LIBRISPEECH(
+                        root, tag, folder_in_archive=folder_in_archive, download=False,
+                    ),
+                    transform,
+                    language_model.encode,
+                )
+                for tag, transform in zip(tags, transform_list)
+            ]
+        )
+
+        data = MapMemoryCache(data)
+        return data
+
+    # For performance, we cache all datasets
+    return tuple(create(dataset) for dataset in datasets)
+
+
+def collate_factory(model_length_function, transforms=None):
+
+    if transforms is None:
+        transforms = torch.nn.Sequential()
+
+    def collate_fn(batch):
+
+        tensors = [transforms(b[0]) for b in batch if b]
+
+        tensors_lengths = torch.tensor(
+            [model_length_function(t) for t in tensors],
+            dtype=torch.long,
+            device=tensors[0].device,
+        )
+
+        tensors = torch.nn.utils.rnn.pad_sequence(tensors, batch_first=True)
+        tensors = tensors.transpose(1, -1)
+
+        targets = [b[1] for b in batch if b]
+        target_lengths = torch.tensor(
+            [target.shape[0] for target in targets],
+            dtype=torch.long,
+            device=tensors.device,
+        )
+        targets = torch.nn.utils.rnn.pad_sequence(targets, batch_first=True)
+
+        return tensors, targets, tensors_lengths, target_lengths
+
+    return collate_fn

examples/pipeline_wav2letter/languagemodels.py

@@ -0,0 +1,38 @@
+import collections
+import itertools
+
+
+class LanguageModel:
+    def __init__(self, labels, char_blank, char_space):
+
+        self.char_space = char_space
+        self.char_blank = char_blank
+
+        labels = [l for l in labels]
+        self.length = len(labels)
+        enumerated = list(enumerate(labels))
+        flipped = [(sub[1], sub[0]) for sub in enumerated]
+
+        d1 = collections.OrderedDict(enumerated)
+        d2 = collections.OrderedDict(flipped)
+        self.mapping = {**d1, **d2}
+
+    def encode(self, iterable):
+        if isinstance(iterable, list):
+            return [self.encode(i) for i in iterable]
+        else:
+            return [self.mapping[i] + self.mapping[self.char_blank] for i in iterable]
+
+    def decode(self, tensor):
+        if len(tensor) > 0 and isinstance(tensor[0], list):
+            return [self.decode(t) for t in tensor]
+        else:
+            # not idempotent, since clean string
+            x = (self.mapping[i] for i in tensor)
+            x = "".join(i for i, _ in itertools.groupby(x))
+            x = x.replace(self.char_blank, "")
+            # x = x.strip()
+            return x
+
+    def __len__(self):
+        return self.length