I was assigned a task to train a wake-word detection model. Basically, it's a binary sequence classification model on audio samples where it should be 1 if it recognizes the wake word being said (e.g. "Ok Google", "Hey Cortana" etc.) and zero otherwise.

The Model

Before starting, I took a quick look at GitHub and found an architecture that seemed simple enough to run on my personal computer. The architecture that follows was inspired from this repository. The tl;dr version is that it is a GRU with a MFCC filter for pre-processing. I divide the input by $2^{15}$ to ensure the audio is within the interval $[-1,1]$.

from torch import nn
import torch
import json
import torchaudio

class WakeWordDetectionModel(nn.Module):
    Wake Word Detection model based on Mycroft Precise
    GitHub: https://github.com/MycroftAI/mycroft-precise/blob/dev/precise/model.py

    def __init__(
        num_layers: int,
        input_dim: int,
        num_units: int,
        num_mels: int,
        num_fft: int,
        dropout: float,
        sample_rate: int
        Initialize the WakeWordDetectionModel.

            num_layers (int): Number of GRU layers in the model.
            input_dim (int): Dimension of the input features.
            num_units (int): Number of units in the GRU hidden layers.
            num_mels (int): Number of mel filter banks for MFCC computation.
            num_fft (int): Size of the FFT window for MFCC computation.
            dropout (float): Dropout probability for the GRU layers.
            sample_rate (int): Audio sample rate.
        super(WakeWordDetectionModel, self).__init__()

        self.gru = nn.GRU(

        self.mfcc = torchaudio.transforms.MFCC(
                "n_mels": num_mels,
                "n_fft": num_fft

        self.params = {
            "num_layers": num_layers,
            "input_dim": input_dim,
            "num_units": num_units,
            "num_mels": num_mels,
            "num_fft": num_fft,
            "dropout": dropout,
            "sample_rate": sample_rate

        self.linear_out = nn.Linear(num_units, 1)
    def dump_params(self, path: str):
        Dump the model parameters to a JSON file.

            path (str): Path to the JSON file.
        with open(path, "w") as f:
            json.dump(self.params, f, indent=4)

    def forward(self, x: torch.Tensor):
        Forward pass of the WakeWordDetectionModel.

            x (torch.Tensor): Input audio data.

            torch.Tensor: Model output.
        #Min-Max scale
        x = x/2**15
        x = self.mfcc(x).transpose(1, 2)
        #x = (x - x.mean())/x.std()
        gru_output, _ = self.gru(x)
        x = gru_output[:, -1, :]
        x = self.linear_out(x)
        return x

    def from_params(path: str):
        Load a WakeWordDetectionModel from a JSON file.

            path (str): Path to the JSON file.

            WakeWordDetectionModel: Loaded model.
        with open(path, "r") as f:
            params = json.load(f)

        return WakeWordDetectionModel(**params)

The Dataset

I received about 126 samples of people saying the wake word and recorded a few myself, totaling about 225 samples. Then, for the false samples I used the Mozilla Common Voice dataset, as anything there wouldn't align with the wake-word (I hope).

Besides that, I also made an augmentation pipeline where at each batch, the audio samples would be crossed with noise from the Urban Sounds 8K dataset and have their tempo sped up or slowed down by sampling a rate from the interval $[0.8, 1.2]$. In short, that means that at each epoch, the model sees slightly different data than before.

For the validation dataset, I used only samples recorded by my voice and the Common Voice dataset.

To make sure there wasn't an imbalance on the samples, as the Common Voice dataset is massively larger than the 225 samples that I have for the activation word, I ensured that only 225 audio files were sampled from it.

Model Selection

I used Optuna to search for hyperparameters. The search space composed of the dimension of the hidden and input tensors, dropout rate, number of layers in the GRU and learning rate. I used a mixture of the Matthews Correlation Coefficient and a custom metric as the optimization objective:

$$F(\text{model}) = \text{matthews}(y_{pred}, y_{real}) + \left(1 - \sigma\left(\frac{2\text{n_params}}{N}\right)\right)$$

Where $N$ is the total sample size. The idea of the second summand is to prioritize models whose number of parameters are o larger than $2N$.

As for the training, I am using BCELossWithLogits, which is why there's no sigmoid at the end of the forward method. The optimization algorithm that I'm using is the NAdam with a learning rate scheduler. I have also implemented an early stopper that will stop the training whenever the derivative of the moving average of the last 20 loss points is too close to zero. I also use L1 regularization with $\lambda$ on the order of $10^{-6}$.


Most of the time, I end up with a history that looks like this:

enter image description here

Very noisy with sharp peaks. When I do manage to get a high Matthews score for the validation, the model seems only to recognize whether there's human speech, regardless of the presence of the wake-word.

With all that said, I would like to ask for some suggestions on what else I could try to improve this model. I have been working on it for over a month and so far haven't got a satisfactory result. Thanks in advance for your time.



Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.