mlx-examples/speechcommands

Train a Keyword Spotting Transformer on Speech Commands

An example of training a Keyword Spotting Transformer¹ on the Speech Commands dataset² with MLX. All supervised only configurations from the paper are available. The example also illustrates how to use MLX Data to load and process an audio dataset.

Pre-requisites

Install the remaining python requirements:

pip install -r requirements.txt

Running the example

Run the example with:

python main.py

By default the example runs on the GPU. To run it on the CPU, use:

python main.py --cpu

For all available options, run:

python main.py --help

Results

After training with the kwt1 architecture for 100 epochs, you should see the following results:

Epoch: 99 | avg. Train loss 0.018 | avg. Train acc 0.996 | Throughput: 662.51 samples/sec
Epoch: 99 | Val acc 0.893 | Throughput: 3091.26 samples/sec
Testing best model from epoch 97
Test acc -> 0.882

For the kwt2 model, you should see:

Epoch: 99 | avg. Train loss 0.003 | avg. Train acc 1.000 | Throughput: 396.53 samples/sec
Epoch: 99 | Val acc 0.901 | Throughput: 1543.48 samples/sec
Testing best model from epoch 94
Test acc -> 0.893

Note that this was run on an M1 Macbook Pro with 16GB RAM.

At the time of writing, mlx doesn't have built-in cosine learning rate schedules, which is used along with the AdamW optimizer in the official implementation. We intend to update this example once these features are added, as well as with appropriate data augmentations.

---

1. Based on the paper Keyword Transformer: A Self-Attention Model for Keyword Spotting ↩︎

2. We use version 0.02. See the paper for more details. ↩︎