Add semi sorted sampler.
State: closed | Pull request: yes
Labels: name: core

name: stale

name: ASR

# What does this PR do ?

Helps to speed up ASR models training without quality degradation.

**Collection**: ASR, Core

# Changelog 
- Adds a semi sorted batch sampler and its prepare function to a new file (in the future it might be some other samplers or new techniques like dynamic batching):
  `nemo/collections/asr/parts/utils/asr_batching.py`
- Added changes to the data loader preparation function in the following model classes:
  `nemo/collections/asr/models/ctc_models.py`
   `nemo/collections/asr/models/rnnt_bpe_models.py`
- Adds changes to the scheduler prepare function because the batch sampler is passed to the dataloader as a sampler and we need to parse the batch size from it. 
  `nemo/core/optim/lr_scheduler.py`
- Adds a change to the `__getitem__` dataset method, as the batch sampler becomes an iterator and returns a list of indices instead of a single index.
  `nemo/collections/asr/data/audio_to_text.py`

# Usage
The sampler is currently added to the following models:
    * `nemo.collections.asr.models.EncDecCTCModel`
    * `nemo.collections.asr.models.EncDecCTCModelBPE`
    * `nemo.collections.asr.models.EncDecRNNTModel`
    * `nemo.collections.asr.models.EncDecRNNTBPEModel`

Needs just to add one line in config.
```bash
++config.model.train_ds.use_semi_sorted_batching=true
```
and two line if we want to specify randomization factor (see more in additional information).
```bash
++config.model.train_ds.randomization_factor=0.1
```

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [x] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Semi Sorted Batching, was proposed in "Speed up training with variable length inputs by efficient batching strategies.", Zhenhao Ge et al. (2021).
* **Idea**: let's sort the samples by their duration, after adding pseudo noise to them, and then sample the batches from the sorted list. This way we avoid strict sorting and keep some level of randomness.
* Experiment for Conformer RNNT BPE:
   * Experiment setup
      The config is used from the [example configs](https://github.com/NVIDIA/NeMo/blob/main/examples/asr/conf/conformer/conformer_transducer_bpe.yaml) with overridden parameters of  ` trainer.epochs=50`, `trainer.devices=8` after the divergence of the first experiment `model.optim.lr=1.0`. For the train data LibriSpeech clean 100 and train clean 360, as well as for the test data, LibriSpeech test clean was used. For data processing, a BPE with a dictionary size of 512 was used.

   * first experiment
        <img src="https://github.com/fedorovgv/asr_ckpt/blob/main/step_1.png?raw=true" width="500" height="250">
        <img src="https://github.com/fedorovgv/asr_ckpt/blob/main/wer_1.png?raw=true" width="500" height="250">
   * second experiment
        <img src="https://github.com/fedorovgv/asr_ckpt/blob/main/step_2.png?raw=true" width="500" height="250">
        <img src="https://github.com/fedorovgv/asr_ckpt/blob/main/wer_2.png?raw=true" width="500" height="250">
        
   * Interesting, that in the first experiment, the training of the model broke up without ssb. 
   * In the second experiment, after reducing the learning rate (from default 5.0 to 1.0), it can be seen that the ssb model converges without quality degradation to ~4.4 WER in the first 50k steps and at **17.5%** faster in time.
   * In paper experiments, it was shown that with further training, the quality of training improves compared to the model without sbb for the same number of epochs with decreasing time.