Add Semi Sorted batch sampler (SSB). 
State: closed | Pull request: yes
Labels: name: core

name: ASR

# What does this PR do ?

**Helps to speed up ASR models training without quality degradation.**

The experiments below showed a **19-12% training speed up with the same quality** for the some models.

**Collection**: ASR, core

# Changelog 
- `nemo/collections/asr/parts/utils/asr_batching.py`

    Add a semi sorted batch sampler and its preparation function to a new file (in the future it might be some other samplers or new techniques like dynamic batching in the file).
    
- `nemo/collections/asr/models/ctc_models.py`
   `nemo/collections/asr/models/ctc_bpe_models.py`
   `nemo/collections/asr/models/rnnt_models.py`
   `nemo/collections/asr/models/rnnt_bpe_models.py`
   `nemo/collections/asr/models/hybrid_rnnt_ctc_bpe_models.py`

    Add changes to the data loader preparation function in the models class.

- `nemo/core/optim/lr_scheduler.py`

    Add changes to the scheduler preparation function because the batch sampler is passed to the dataloader as a sampler and we need to parse the batch size from it. 

- `nemo/collections/asr/data/audio_to_text.py`

    Add 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.

# Usage
The sampler is currently added to the CTC, Transducer and Hybrid models.

**Need just to add one line to the 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
```
Paper default value for randomization_factor is `0.2`.

# 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?

@titu1994, @redoctopus, @jbalam-nv, @okuchaiev

# Additional Information

* Semi Sorted Batching, was proposed in "Speed up training with variable length inputs by efficient batching strategies.", Zhenhao Ge et al. (2021). ([link](https://www.isca-speech.org/archive/pdfs/interspeech_2021/ge21_interspeech.pdf) to the paper)

* **Idea**: Get the duration of the samples, add pseudo noise to them and sort them. Sample the batches from the sorted list. This way we avoid strict sorting and retain some level of randomness with a significant reduction in padding in batches.

* Experiment for Conformer char CTC:
    * train script: `examples/asr/asr_ctc/speech_to_text_ctc.py`
    * train config: `asr/examples/asr/conf/conformer/conformer_ctc_char.yaml`
    * train data: libri-speech `train-clean-100` `train-clean-360`
    * test sata: libri-speech `test-clean`
    * overrided parameters:
      * `model.validation_ds.batch_size=32`
      * `trainer.devices=2`
      * `trainer.max_epochs=50`
      * `trainer.accumulate_grad_batches=10`
      * `trainer.gradient_clip_val=1.0 `
      * `++model.train_ds.use_semi_sorted_batching=true` (only experiment with ssb)
      * `++model.train_ds.randomization_factor=0.1` (only experiment with ssb)
    * devices: Nvidia Tesla V100 x2
    * tensorboard logs:
      * time
        <img src="https://github.com/fedorovgv/buf/blob/main/char_ctc_time.png?raw=true" alt="alt text" width="400" height="200">
      * wer
        <img src="https://github.com/fedorovgv/buf/blob/main/char_ctc_wer.png?raw=true" alt="alt text" width="400" height="200">
    * average, time per epoch / total time / final wer metric
      * `ssb`:  **22 min 28 sec** / **18 hours 44 min 14 sec** / 9,371 wer
      * `without ssb`: 27 min 44 sec /  23 hours 7 min 31 sec / **9,08** wer
  
**Model training speed up by 18,9% with 3% loss of quality.**