Diarization Inference - Runs out of memory, bad performance
State: closed | Pull request: no
Labels: name: stale

Hi all,

Thank you to the creators of this tool! 

### Running out of Memory
Unfortunately, while running a stripped down version of [this tutorial](https://github.com/NVIDIA/NeMo/blob/main/tutorials/speaker_tasks/Speaker_Diarization_Inference.ipynb) on my dataset of radio conversations (which are 1 min to 10 minutes in length), the diarization inference occasionally runs out of memory for files longer than ~180s, although sometimes it only runs outof memory for files longer than ~200s. Is there any way to change *any* configurations at all for inference to prevent running out of memory? This **[NVIDIA guide](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/results.html#inference-on-long-audio)** suggests some sort of alterations, but I don't understand what it's explicitly recommending in the context of my code. Would changing the number of workers help? My GPU is admittedly not the greatest (NVIDIA GeForce RTX 3060 ~6GB GPU memory), but I'd like to avoid needing to use Google Colab, as my hundreds of gB of files are offline.


### My code
This is my code - basically, just loads up pre-calculated VAD manifests, sets up the path and configures Omega correctly, then launches the diarizer. At the end it saves a pyplo.
```python
def diarize(wav_path, rttm_path):
    #initial imports
    import os, wget, json, librosa, torch
    from nemo.collections.asr.parts.utils.speaker_utils import rttm_to_labels, labels_to_pyannote_object
    from nemo.collections.asr.models import ClusteringDiarizer
    from omegaconf import OmegaConf
    import matplotlib.pyplot as plt
    from pyannote.core import notebook
    from pytorch_memlab import profile
    import re

    
    output_dir = "/optimize_diarization/4_outputRTTMs"
    yml_dir = "/optimize_diarization/3_to_4_scripts/yml"
    manifests_dir = "/optimize_diarization/3b_speechManifests"
    graphics_dir = "optimize_diarization/4_results/graphics"
    intermediate_inputmanifest_dir = "/optimize_diarization/3_to_4_scripts/input_manifest"


    # signal, sr = librosa.load(wav_path,sr=16000)
    duration = librosa.get_duration(path=wav_path)

   # configure manifest path stuff
    filename = re.search(r'.*/(.*)', wav_path).group(1)[0:-4]
    print(f'\n\n FILENAME IS: {filename} \n\n WAV_PATH is {wav_path} \n\n RTTM_PATH is {rttm_path} \n\n')
    external_vad_manifest = manifests_dir + '/' + filename + 'vad_out.json'
    input_manifest = intermediate_inputmanifest_dir + '/input_manifest.json'

    #make manifest
    meta = {
        'audio_filepath': wav_path, 
        'offset': 0, 
        'duration': duration, 
        'label': 'infer', 
        'text': '-', 
        'num_speakers': None, 
        'rttm_filepath': rttm_path, 
        'uem_filepath' : None,
        'uniq_id' : filename
    }
    with open(f'{intermediate_inputmanifest_dir}/input_manifest.json','w') as fp:
        json.dump(meta,fp)
        fp.write('\n')

    #set up model/config
    MODEL_CONFIG = os.path.join(yml_dir,'diar_infer_telephonic.yaml')
    if not os.path.exists(MODEL_CONFIG):
        config_url = "https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/speaker_tasks/diarization/conf/inference/diar_infer_telephonic.yaml"
        MODEL_CONFIG = wget.download(config_url,yml_dir)
    config = OmegaConf.load(MODEL_CONFIG)

    config.num_workers = 1 # Workaround for multiprocessing hanging with ipython issue 
    config.diarizer.oracle_vad = True
    config.diarizer.vad.model_path = None

    config.diarizer.vad.external_vad_manifest = external_vad_manifest
    config.diarizer.manifest_filepath = input_manifest

    config.diarizer.out_dir = output_dir # Directory to store intermediate files and prediction outputs
    config.diarizer.speaker_embeddings.model_path = 'titanet_large'
    config.diarizer.speaker_embeddings.parameters.window_length_in_sec = [1.5,1.25,1.0,0.75,0.5] 
    config.diarizer.speaker_embeddings.parameters.shift_length_in_sec = [0.75,0.625,0.5,0.375,0.1] 
    config.diarizer.speaker_embeddings.parameters.multiscale_weights= [1,1,1,1,1] 
    config.diarizer.clustering.parameters.oracle_num_speakers = False
    config.diarizer.msdd_model.model_path = 'diar_msdd_telephonic' # Telephonic speaker diarization model 
    config.diarizer.msdd_model.parameters.sigmoid_threshold = [0.7, 1.0] # Evaluate with T=0.7 and T=1.0

    
    # Run the diarization
    sd_model = ClusteringDiarizer(cfg=config)
    sd_model.diarize()


    #output as a graph
    pred_labels_neural = rttm_to_labels(f'{output_dir}/pred_rttms/{filename}.rttm')
    hypothesis_neural = labels_to_pyannote_object(pred_labels_neural)
    fig, ax = plt.subplots()
    notebook.plot_annotation(hypothesis_neural, ax=ax, time=True, legend=True)
    fig.savefig(f'{graphics_dir}/{filename}diarization_results.png')
 
```

### Poor Results
The performance is quite poor. All files are resampled to 16kHz, and while a few achieve 10-20% diarization ER, the majority have diarization errors from  60-70%. I'm not sure what the data environment is for Titanet and MSDD, but my conversations feature usually 3-6 different individuals. I'm experimenting with preprocessing and trying other models, but any suggestions would be welcome.

I'm unfortunately unable to finetune the Titanet model to improve performance, because it, also, consumes insane amounts of GPU memory [(1)](https://github.com/NVIDIA/NeMo/issues/6031), [(2)](https://github.com/NVIDIA/NeMo/issues/4533). Presumably it wants to be trained on small audio files, but mine are quite large. Google Colab's 16GB T4 is not up to the task, even with small batch sizes (4 / 4).
```[tasklist]
### Tasks
- [ ] memory
```