Speaker Diarization using GRU in PyTorch

Installation

Conda environment

To use this conda environment, you need to install Miniconda, then run this command

conda env create -f environment.yml

Install OpenVINO (Optional)

If you would like to use OpenVINO for inference. Please check OpenVINO official Documentation for the installation.

Step 1: Training 💪

Simply run the code in the Part 1 - Training.ipynb notebook and you are good to go.

Here is the training process if you would want to learn more:

• Dataset Preparation
  You should structure your data in these format:

  data
      train
          speaker_1
              audio_1.wav
              audio_2.wav
          speaker_2
              audio_1.wav
              audio_2.wav            
      valid
          speaker_1
              audio_1.wav
              audio_2.wav
          speaker_2
              audio_1.wav
              audio_2.wav
      test
          both_speaker_speaking.wav      
  

  You can use this sample dataset in Kaggle to help you understand how to structure your own data.

• Dataset & Dataloader
  I have created AudioFolder PyTorch Dataset Object so automatically utilize the dataset. Here is a quick look for you

  AudioFolder(folder_path, n_data=500, slice_dur=5, sr=16000, n_mel=128, n_fft=2048, win_length=2048, hop_length=1024, freq_range=(0, 8000)):
  

  In a nutshell, here is what it does:

  • Randomly choose an audio from folder_path, then sampled it using sr rate
  • Then it randomly slice a chunk of the audio for slice_dur seconds, quantile normalize, then converted to log melspectrogram using specified parameter (n_mel, n_fft, and so on)
  • Zero crossing rate is used as VAD to label the spectrogram frames. Each frame is labeled as the subfolder name or <NONE> if it is silent.
  • AudioFolder does these steps n_data times to create the dataset

• Architecture & Config
  Simply create the model architecture you would like to use. I have made one sample for you in src/model.py. You can directly use it by simple changing the hyperparameters in the config. Please follow the notebook if you are confused.

  ❗ Note: If you would like to use custom architecture in OpenVINO, make sure it is compatible with OpenVINO model optimizer and inference engine.

• Training Preparation
  Set up the model, criterion, optimizer, and callback here.

• Training
  As what it says, running the code will train the model

Step 2: Inference 🤖

I have prepared a PyTorchPredictor and OpenVINOPredictor object. Here is a sneak peek:

PyTorchPredictor(weights_path, config_path, chunk_dur)
OpenVINOPredictor(model_xml, model_bin, config_path, chunk_dur, CPU_EXT)

To use it, simply input the arguments and use .predict(wav_path) and it will return the diarization timestamp. For example:

p = PyTorchPredictor("model/weights_best.pth", "model/configs.pth", chunk_dur=2)
timestamp = p.predict("dataset/test/test.wav", plot=True)

setting plot=True provides you with the diarization visualization like this

If you would like to use OpenVINO, use .to_onnx(fname, outdir) to convert the model into onnx format.

p.to_onnx("speaker_diarization.onnx", "model")

End notes ❤️

I hope you like the projects. I made this repo for educational purposes so it might need further tweaking to reach production level. Feel free to create an issue if you need help, and I hope I'll have the time to help you. Thank you.