Add Semi Sorted Batching.
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.**

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
- [ x] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

## Who can review?

@nithinraok

# 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.**

--- 

- Experiment for Conformer BPE CTC:
    * train script: `examples/asr/asr_ctc/speech_to_text_ctc_bpe.py`
    * train config: `asr/examples/asr/conf/conformer/conformer_ctc_bpe.yaml`
    * train data: libri-speech `train-clean-100` `train-clean-360`
    * test sata: libri-speech `test-clean`
    * size of bpe vocabulary `512`
    * 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/bpe_ctc_time.png?raw=true" alt="alt text" width="400" height="200">
      * wer
        <img src="https://github.com/fedorovgv/buf/blob/main/bpe_ctc_wer.png?raw=true" alt="alt text" width="400" height="200">
    * average, time per epoch / total time / final wer metric
      * `ssb`:  **50 min 1 sec** / **1 day 17 hours 42 min 52 sec** / **6,417** wer
      * `without ssb`: 1 hour 1 min 0 sec /  2 day 2 hours 46 min 59 sec / 6,682 wer

**Model training  speed up by 17,8% with improved quality.**

--- 

- Experiment for Conformer Char RNNT:
    * train script: `examples/asr/asr_transducer/speech_to_text_rnnt.py`
    * train config: `examples/asr/conf/conformer/conformer_transducer_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_rnnt_rime.png?raw=true" alt="alt text" width="400" height="200">
      * wer
        <img src="https://github.com/fedorovgv/buf/blob/main/char_rnnt_wer.png?raw=true" alt="alt text" width="400" height="200">
    * average, time per epoch / total time / final wer metric
      * `ssb`:  **48 min 54 sec** / **1 day 16 hours 45 min 33 sec** / 4,913 wer
      * `without ssb`: 56 min 0 sec /  1 day 22 hours 10 min 40 sec / **4,766** wer

**Model training  speed up by 13,9% with 3% loss of quality.**

---

- Experiment for Conformer BPE RNNT:
    * train script: `examples/asr/asr_transducer/speech_to_text_rnnt_bpe.py`
    * train config: `examples/asr/conf/conformer/conformer_transducer_bpe.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/bpe_rnnt_time.png?raw=true" alt="alt text" width="400" height="200">
      * wer
        <img src="https://github.com/fedorovgv/buf/blob/main/bpe_rnnt_wer.png?raw=true" alt="alt text" width="400" height="200">
    * average, time per epoch / total time / final wer metric
      * `ssb`:  **42 min 3 sec** / **1 day 11 hours 2 min 30 sec** / **4,894** wer
      * `without ssb`:  47 min 52 sec /  1 day 15 hours 53 min 44 sec / 4,991 wer

**Model training  speed up by 12,1% with 2% improved of quality.**

---

Picture for docs is available via this link : https://github.com/NVIDIA/NeMo/releases/download/v1.22.0/ssb.png

Closed PR : https://github.com/NVIDIA/NeMo/pull/7307

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

This PR was closed because it has been inactive for 7 days since being marked as stale.

@nithinraok Hi, can you please approve new PR and re-open it.

@fedorovgv thanks, is this a new PR? Did you close earlier one? 

this is new one with the same code. yes, i did because of the problems with commits signs.

@nithinraok Hi! All minor comments have been corrected. Could you view it please?

@nithinraok Sorry for the ping, but how long does Jenkins take? is it finite?)

jenkins

I have initiated the jenkins run, it should be done in couple of hrs

@nithinraok so, can we merge it? or do I need to update the branch before your merge? 

> @nithinraok so, can we merge it? or do I need to update the branch before your merge? 

Thanks a lot for your contribution!

> > @nithinraok so, can we merge it? or do I need to update the branch before your merge?
> 
> Thanks a lot for your contribution!

Thanks a lot for great framework!