{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "838ffa12",
   "metadata": {},
   "source": [
    "This notebook demonstrates how the speech embedding model from Google (https://www.kaggle.com/models/google/speech-embedding/frameworks/tensorFlow1/variations/speech-embedding/versions/1) is re-implemented in Keras manually, which can then be converted to ONNX and tflite formats for use in openWakeWord.\n",
    "\n",
    "Note that Keras was used here, but in theory other deep learning frameworks (e.g., PyTorch) could work as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "893d29dc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-18T00:26:11.649261Z",
     "start_time": "2024-01-18T00:26:10.190666Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-01-17 19:26:10.372628: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2024-01-17 19:26:10.372640: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n"
     ]
    }
   ],
   "source": [
    "# Imports\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import scipy\n",
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub  # install with `pip install tensorflow_hub`\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe3054ae",
   "metadata": {},
   "source": [
    "# Load Orignal Model from TFHub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fa2bc0d3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-18T00:26:12.257919Z",
     "start_time": "2024-01-18T00:26:11.650661Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-01-17 19:26:11.857817: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:936] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2024-01-17 19:26:11.858167: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2024-01-17 19:26:11.858193: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublas.so.11'; dlerror: libcublas.so.11: cannot open shared object file: No such file or directory\n",
      "2024-01-17 19:26:11.858215: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublasLt.so.11'; dlerror: libcublasLt.so.11: cannot open shared object file: No such file or directory\n",
      "2024-01-17 19:26:11.858237: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcufft.so.10'; dlerror: libcufft.so.10: cannot open shared object file: No such file or directory\n",
      "2024-01-17 19:26:11.858258: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcurand.so.10'; dlerror: libcurand.so.10: cannot open shared object file: No such file or directory\n",
      "2024-01-17 19:26:11.858278: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusolver.so.11'; dlerror: libcusolver.so.11: cannot open shared object file: No such file or directory\n",
      "2024-01-17 19:26:11.858299: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusparse.so.11'; dlerror: libcusparse.so.11: cannot open shared object file: No such file or directory\n",
      "2024-01-17 19:26:11.858320: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudnn.so.8'; dlerror: libcudnn.so.8: cannot open shared object file: No such file or directory\n",
      "2024-01-17 19:26:11.858325: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1850] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n",
      "Skipping registering GPU devices...\n",
      "2024-01-17 19:26:11.858458: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "# Load the original speech embedding model (now hosted on Kaggle) as a KerasLayer object\n",
    "\n",
    "embedding_model_url = \"https://tfhub.dev/google/speech_embedding/1\"\n",
    "embedding_model = hub.KerasLayer(embedding_model_url, trainable=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6769931f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-18T00:26:12.375204Z",
     "start_time": "2024-01-18T00:26:12.259632Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedding Output Shape: (1, 1, 1, 96)\n"
     ]
    }
   ],
   "source": [
    "# Get predictions from the embedding model for a chunk of ~775 ms audio data (at 16khz)\n",
    "# This is the minimum input size for the model per the documentation here: https://www.kaggle.com/models/google/speech-embedding/frameworks/tensorFlow1/variations/speech-embedding/versions/1\n",
    "\n",
    "# Load sample clip, and select a 775 ms chunk and normalize between -1 and 1\n",
    "sr, sample_data = scipy.io.wavfile.read(\"../tests/data/hey_mycroft_test.wav\")\n",
    "sample_data = (sample_data[0:12400][None,]/32767).astype(np.float32)\n",
    "embeddings = embedding_model(sample_data)\n",
    "print(\"Embedding Output Shape:\", embeddings.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8acfba9",
   "metadata": {},
   "source": [
    "# Convert original model to tflite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b4ee1693",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-18T00:26:14.554068Z",
     "start_time": "2024-01-18T00:26:12.376427Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.\n",
      "2024-01-17 19:26:12.970115: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: google_speech_embedding_fixed_input/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: google_speech_embedding_fixed_input/assets\n",
      "2024-01-17 19:26:14.061242: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:357] Ignored output_format.\n",
      "2024-01-17 19:26:14.061261: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:360] Ignored drop_control_dependency.\n",
      "2024-01-17 19:26:14.061688: I tensorflow/cc/saved_model/reader.cc:43] Reading SavedModel from: google_speech_embedding_fixed_input\n",
      "2024-01-17 19:26:14.066816: I tensorflow/cc/saved_model/reader.cc:78] Reading meta graph with tags { serve }\n",
      "2024-01-17 19:26:14.066828: I tensorflow/cc/saved_model/reader.cc:119] Reading SavedModel debug info (if present) from: google_speech_embedding_fixed_input\n",
      "2024-01-17 19:26:14.076781: I tensorflow/cc/saved_model/loader.cc:228] Restoring SavedModel bundle.\n",
      "2024-01-17 19:26:14.173155: I tensorflow/cc/saved_model/loader.cc:212] Running initialization op on SavedModel bundle at path: google_speech_embedding_fixed_input\n",
      "2024-01-17 19:26:14.236537: I tensorflow/cc/saved_model/loader.cc:301] SavedModel load for tags { serve }; Status: success: OK. Took 174850 microseconds.\n",
      "2024-01-17 19:26:14.288965: I tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc:237] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.\n",
      "2024-01-17 19:26:14.405541: W tensorflow/compiler/mlir/lite/flatbuffer_export.cc:1881] Graph contains the following resource op(s), that use(s) resource type. Currently, the resource type is not natively supported in TFLite. Please consider not using the resource type if there are issues with either TFLite converter or TFLite runtime:\n",
      "Resource ops: TensorArrayGatherV3, TensorArrayReadV3, TensorArrayScatterV3, TensorArrayV3, TensorArrayWriteV3\n",
      "Details:\n",
      "\ttf.TensorArrayGatherV3(tensor<2x!tf_type.resource<tensor<*xf32>>>, tensor<?xi32>, tensor<f32>) -> (tensor<?x?x1x32xf32>) : {device = \"\", element_shape = #tf_type.shape<?x1x32>}\n",
      "\ttf.TensorArrayReadV3(tensor<2x!tf_type.resource<tensor<*xf32>>>, tensor<i32>, tensor<f32>) -> (tensor<*xf32>) : {device = \"\"}\n",
      "\ttf.TensorArrayScatterV3(tensor<2x!tf_type.resource<tensor<*xf32>>>, tensor<?xi32>, tensor<?x12400xf32>, tensor<f32>) -> (tensor<f32>) : {device = \"\"}\n",
      "\ttf.TensorArrayV3(tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*xf32>>>, tensor<f32>) : {clear_after_read = true, device = \"\", dtype = f32, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = \"\"}\n",
      "\ttf.TensorArrayWriteV3(tensor<2x!tf_type.resource<tensor<*xf32>>>, tensor<i32>, tensor<?x?x?xf32>, tensor<f32>) -> (tensor<f32>) : {device = \"\"}\n",
      "2024-01-17 19:26:14.405561: W tensorflow/compiler/mlir/lite/flatbuffer_export.cc:1892] TFLite interpreter needs to link Flex delegate in order to run the model since it contains the following Select TFop(s):\n",
      "Flex ops: FlexTensorArrayGatherV3, FlexTensorArrayReadV3, FlexTensorArrayScatterV3, FlexTensorArrayV3, FlexTensorArrayWriteV3\n",
      "Details:\n",
      "\ttf.TensorArrayGatherV3(tensor<2x!tf_type.resource<tensor<*xf32>>>, tensor<?xi32>, tensor<f32>) -> (tensor<?x?x1x32xf32>) : {device = \"\", element_shape = #tf_type.shape<?x1x32>}\n",
      "\ttf.TensorArrayReadV3(tensor<2x!tf_type.resource<tensor<*xf32>>>, tensor<i32>, tensor<f32>) -> (tensor<*xf32>) : {device = \"\"}\n",
      "\ttf.TensorArrayScatterV3(tensor<2x!tf_type.resource<tensor<*xf32>>>, tensor<?xi32>, tensor<?x12400xf32>, tensor<f32>) -> (tensor<f32>) : {device = \"\"}\n",
      "\ttf.TensorArrayV3(tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*xf32>>>, tensor<f32>) : {clear_after_read = true, device = \"\", dtype = f32, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = \"\"}\n",
      "\ttf.TensorArrayWriteV3(tensor<2x!tf_type.resource<tensor<*xf32>>>, tensor<i32>, tensor<?x?x?xf32>, tensor<f32>) -> (tensor<f32>) : {device = \"\"}\n",
      "See instructions: https://www.tensorflow.org/lite/guide/ops_select\n"
     ]
    }
   ],
   "source": [
    "# Build model with specific input size, and save\n",
    "inputs = tf.keras.Input((12400,))\n",
    "x = embedding_model(inputs)\n",
    "model = tf.keras.Model(inputs=inputs, outputs=x)\n",
    "model.save(\"google_speech_embedding_fixed_input\")\n",
    "\n",
    "speech_embedding_dir = \"google_speech_embedding_fixed_input\"\n",
    "# speech_embedding_dir = \"google_speech_embedding_savedmodel/\"\n",
    "\n",
    "converter = tf.lite.TFLiteConverter.from_saved_model(speech_embedding_dir)#, tags=[\"train\"])\n",
    "# convert = tf.lite.TFLiteConverter.from_keras_model(embedding_model)\n",
    "converter.target_spec.supported_ops = [\n",
    "  tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS\n",
    "]\n",
    "# converter.allow_custom_ops = True\n",
    "\n",
    "tflite_model = converter.convert()\n",
    "with open(speech_embedding_dir + '/speech_embeddings.tflite', 'wb') as f:\n",
    "    f.write(tflite_model)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "927ebfda",
   "metadata": {},
   "source": [
    "# Comparing Log-Mel Features"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48c7138b",
   "metadata": {},
   "source": [
    "The speech embedding model from Google computes it's own input features from raw audio, which is convenient, but not ideal as it combines pre-processing with the model in a way that makes the model less understandable. In particular, this is (to my knowledge) the total information provided about the feature creation:\n",
    "\n",
    "From the model page [here:](https://www.kaggle.com/models/google/speech-embedding/frameworks/tensorFlow1/variations/speech-embedding/versions/1)\n",
    "```\n",
    "The module computes its own 32 dimensional log-mel features from the provided audio samples using the following parameters:\n",
    "\n",
    "    stft window size: 25ms\n",
    "    stft window step: 10ms\n",
    "    mel band limits: 60Hz - 3800Hz\n",
    "    mel frequency bins: 32\n",
    "```\n",
    "\n",
    "And then this excerpt from the corresponding [paper](https://arxiv.org/abs/2002.01322):\n",
    "\n",
    "```\n",
    "Our model is designed for deployment in an environment\n",
    "where both memory and compute power are very limited,\n",
    "such as on a digital signal processor (DSP). It runs on top of a\n",
    "low footprint feature extractor that provides a 32 dimensional\n",
    "log mel feature vector covering the frequency range from\n",
    "60 Hz to 3800 Hz, quantized to 8 bits every 10 ms\n",
    "```\n",
    "\n",
    "It seems likely that this implementation is simply a [spectrogram](https://librosa.org/doc/main/generated/librosa.feature.melspectrogram.html) with [log scaling](https://librosa.org/doc/main/generated/librosa.power_to_db.html), but the investigation below shows that this may note be the case.\n",
    "\n",
    "If you have a theory as to what the original model is doing, or why a standard log-mel spectrogram does not match, please open an issue on the [openWakeWord](https://github.com/dscripka/openWakeWord), I would love learn more about this!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6599a6a0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-18T00:26:14.703602Z",
     "start_time": "2024-01-18T00:26:14.555114Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO: Created TensorFlow Lite delegate for select TF ops.\n",
      "2024-01-17 19:26:14.557849: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:936] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2024-01-17 19:26:14.558198: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1850] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n",
      "Skipping registering GPU devices...\n",
      "INFO: TfLiteFlexDelegate delegate: 4 nodes delegated out of 76 nodes with 2 partitions.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedding model features shape: (32, 76)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 23024 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Use the converted tflite model and get intermediate outputs to extract the log-mel features\n",
    "\n",
    "interpreter = tf.lite.Interpreter(\n",
    "    model_path=os.path.join(speech_embedding_dir, \"speech_embeddings.tflite\"),\n",
    "    num_threads=1,\n",
    "    experimental_preserve_all_tensors=True\n",
    ")\n",
    "interpreter.allocate_tensors()\n",
    "\n",
    "# Get input and output tensors\n",
    "input_details = interpreter.get_input_details()\n",
    "output_details = interpreter.get_output_details()\n",
    "interpreter.set_tensor(input_details[0]['index'], sample_data)\n",
    "interpreter.invoke()\n",
    "\n",
    "spec = interpreter.get_tensor(65)  # This index is the log-mel features, to my knowledge\n",
    "spec = spec.squeeze().T  # transform for visualization\n",
    "print(\"Embedding model features shape:\", spec.shape)\n",
    "\n",
    "_ = plt.imshow(spec)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc88eeaa",
   "metadata": {},
   "source": [
    "This certainly *looks* like a log-mel spectrogram, and we can compute the same from the reference Librosa implementation for comparison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5363ba83",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-18T00:26:26.195376Z",
     "start_time": "2024-01-18T00:26:25.418886Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Librosa features shape: (32, 76)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 22476 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import librosa\n",
    "\n",
    "S = librosa.feature.melspectrogram(y=sample_data, win_length=int(0.025*16000), \n",
    "                                   hop_length=int(0.010*16000), n_fft=512, center=True,\n",
    "                                   sr=16000, n_mels=32, fmin=60, fmax=3800, power=2)#, norm=None)\n",
    "\n",
    "S = librosa.power_to_db(S).squeeze()[:, 1:-1] # convert to logmel and remove edge columns from center=True\n",
    "\n",
    "print(\"Librosa features shape:\", spec.shape)\n",
    "_ = plt.imshow(S)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f03308c2",
   "metadata": {},
   "source": [
    "Visually, these mel-spectrograms are very similar, but on closer inspection there are differences. Plotting at a single time slice better shows the difference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d80091c0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-18T00:26:30.675702Z",
     "start_time": "2024-01-18T00:26:30.579560Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48908 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot single time slice from both melspectrograms\n",
    "_ = plt.plot(spec[:, 33])\n",
    "_ = plt.plot(S[:, 33]/10 + 2)  # apply simple scalar transformation to better align the points\n",
    "_ = plt.xlabel(\"Frequency Bins\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98dce06c",
   "metadata": {},
   "source": [
    "While the overall trend and specific frequency features are very similar between the two, they are not exact. After some time investigating this difference, eventually I moved on to other tasks with openWakeWord, and assumed that the similarity of the spectrograms would mean the downstream model performance would be relatively unnaffected. This assumption seems to have been largely true, and typically the performance difference between the openWakeWord implementation and the original Google embedding model is small.\n",
    "\n",
    "For completeness, below is the implementation of a melspectrogram using just PyTorch, so that it can be converted to ONNX/tflite for more efficient computation on a wide range of devices. This code was based on the implementation from [torchlibrosa](https://github.com/qiuqiangkong/torchlibrosa) and is identical to the librosa reference implementation to within rounding error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07e24374",
   "metadata": {},
   "outputs": [],
   "source": [
    "# torchlibrosa version of melspectrogram\n",
    "\n",
    "import torch\n",
    "import torchlibrosa as tl\n",
    "import numpy as np\n",
    "\n",
    "batch_size = 1\n",
    "sample_rate = 16000\n",
    "win_length = 400\n",
    "hop_length = 160\n",
    "n_mels = 32\n",
    "nfft=512\n",
    "\n",
    "batch_audio = torch.empty(batch_size, 32000).uniform_(-1, 1)  # (batch_size, sample_rate)\n",
    "\n",
    "def f(self, input):\n",
    "    r\"\"\"Power to db, this function is the pytorch implementation of \n",
    "    librosa.power_to_lb.\n",
    "    \"\"\"\n",
    "    ref_value = self.ref\n",
    "    log_spec = 10.0 * torch.log(torch.clamp(input, min=self.amin, max=np.inf))/torch.log(torch.tensor(10))\n",
    "    log_spec -= 10.0 * torch.log(torch.maximum(torch.tensor(self.amin), torch.tensor(ref_value)))/torch.log(torch.tensor(10))\n",
    "\n",
    "    if self.top_db is not None:\n",
    "        if self.top_db < 0:\n",
    "            raise librosa.util.exceptions.ParameterError('top_db must be non-negative')\n",
    "        log_spec = torch.clamp(log_spec, min=log_spec.max() - self.top_db, max=np.inf)\n",
    "\n",
    "    return log_spec\n",
    "\n",
    "tl.stft.LogmelFilterBank.power_to_db = f\n",
    "\n",
    "# TorchLibrosa feature extractor the same as librosa.feature.melspectrogram()\n",
    "feature_extractor = torch.nn.Sequential(\n",
    "    tl.Spectrogram(\n",
    "        center=False,\n",
    "        n_fft=nfft,\n",
    "        hop_length=hop_length,\n",
    "        win_length=win_length,\n",
    "    ), tl.LogmelFilterBank(\n",
    "        n_fft=nfft,\n",
    "        sr=sample_rate,\n",
    "        n_mels=n_mels,\n",
    "        fmin=60,\n",
    "        fmax=3800,\n",
    "        is_log=True, # Default is true\n",
    "    ))\n",
    "\n",
    "# export to onnx\n",
    "torch.onnx.export(feature_extractor, batch_audio, \"torchlibrosa_onnx_melspectrogram.onnx\",\n",
    "                  opset_version=12, input_names = ['input'], output_names = ['output'], \n",
    "                  dynamic_axes={\"input\": {0: 'batch_size', 1: 'samples'}, \"output\": {0: 'time'}})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbf87b19",
   "metadata": {},
   "source": [
    "# Create New Model with Keras"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bda8bd3",
   "metadata": {},
   "source": [
    "After separting the log-mel feature calculution from the embedding model, we can now re-produce the rest of the model manually in Keras.\n",
    "\n",
    "Note that for many of the layers below, the hard-coded values and parameters were obtained by inspecting the tflite version of the original embedding model, using a tool like [Netron](http://www.netron.app)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4fca0fa3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-18T00:27:56.152072Z",
     "start_time": "2024-01-18T00:27:55.876595Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model_1\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_2 (InputLayer)           [(None, 76, 32, 1)]  0           []                               \n",
      "                                                                                                  \n",
      " zero_padding2d (ZeroPadding2D)  (None, 76, 34, 1)   0           ['input_2[0][0]']                \n",
      "                                                                                                  \n",
      " conv2d (Conv2D)                (None, 74, 32, 24)   216         ['zero_padding2d[0][0]']         \n",
      "                                                                                                  \n",
      " batch_normalization (BatchNorm  (None, 74, 32, 24)  96          ['conv2d[0][0]']                 \n",
      " alization)                                                                                       \n",
      "                                                                                                  \n",
      " tf.math.multiply (TFOpLambda)  (None, 74, 32, 24)   0           ['batch_normalization[0][0]']    \n",
      "                                                                                                  \n",
      " tf.math.truediv (TFOpLambda)   (None, 74, 32, 24)   0           ['tf.math.multiply[0][0]']       \n",
      "                                                                                                  \n",
      " tf.math.maximum (TFOpLambda)   (None, 74, 32, 24)   0           ['tf.math.truediv[0][0]',        \n",
      "                                                                  'batch_normalization[0][0]']    \n",
      "                                                                                                  \n",
      " tf.math.maximum_1 (TFOpLambda)  (None, 74, 32, 24)  0           ['tf.math.maximum[0][0]']        \n",
      "                                                                                                  \n",
      " conv2d_1 (Conv2D)              (None, 74, 32, 24)   1728        ['tf.math.maximum_1[0][0]']      \n",
      "                                                                                                  \n",
      " batch_normalization_1 (BatchNo  (None, 74, 32, 24)  96          ['conv2d_1[0][0]']               \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " tf.math.multiply_1 (TFOpLambda  (None, 74, 32, 24)  0           ['batch_normalization_1[0][0]']  \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.truediv_1 (TFOpLambda)  (None, 74, 32, 24)  0           ['tf.math.multiply_1[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.maximum_2 (TFOpLambda)  (None, 74, 32, 24)  0           ['tf.math.truediv_1[0][0]',      \n",
      "                                                                  'batch_normalization_1[0][0]']  \n",
      "                                                                                                  \n",
      " tf.math.maximum_3 (TFOpLambda)  (None, 74, 32, 24)  0           ['tf.math.maximum_2[0][0]']      \n",
      "                                                                                                  \n",
      " conv2d_2 (Conv2D)              (None, 72, 32, 24)   1728        ['tf.math.maximum_3[0][0]']      \n",
      "                                                                                                  \n",
      " batch_normalization_2 (BatchNo  (None, 72, 32, 24)  96          ['conv2d_2[0][0]']               \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " tf.math.multiply_2 (TFOpLambda  (None, 72, 32, 24)  0           ['batch_normalization_2[0][0]']  \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.truediv_2 (TFOpLambda)  (None, 72, 32, 24)  0           ['tf.math.multiply_2[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.maximum_4 (TFOpLambda)  (None, 72, 32, 24)  0           ['tf.math.truediv_2[0][0]',      \n",
      "                                                                  'batch_normalization_2[0][0]']  \n",
      "                                                                                                  \n",
      " tf.math.maximum_5 (TFOpLambda)  (None, 72, 32, 24)  0           ['tf.math.maximum_4[0][0]']      \n",
      "                                                                                                  \n",
      " max_pooling2d (MaxPooling2D)   (None, 36, 16, 24)   0           ['tf.math.maximum_5[0][0]']      \n",
      "                                                                                                  \n",
      " conv2d_3 (Conv2D)              (None, 36, 16, 48)   3456        ['max_pooling2d[0][0]']          \n",
      "                                                                                                  \n",
      " batch_normalization_3 (BatchNo  (None, 36, 16, 48)  192         ['conv2d_3[0][0]']               \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " tf.math.multiply_3 (TFOpLambda  (None, 36, 16, 48)  0           ['batch_normalization_3[0][0]']  \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.truediv_3 (TFOpLambda)  (None, 36, 16, 48)  0           ['tf.math.multiply_3[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.maximum_6 (TFOpLambda)  (None, 36, 16, 48)  0           ['tf.math.truediv_3[0][0]',      \n",
      "                                                                  'batch_normalization_3[0][0]']  \n",
      "                                                                                                  \n",
      " tf.math.maximum_7 (TFOpLambda)  (None, 36, 16, 48)  0           ['tf.math.maximum_6[0][0]']      \n",
      "                                                                                                  \n",
      " conv2d_4 (Conv2D)              (None, 34, 16, 48)   6912        ['tf.math.maximum_7[0][0]']      \n",
      "                                                                                                  \n",
      " batch_normalization_4 (BatchNo  (None, 34, 16, 48)  192         ['conv2d_4[0][0]']               \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " tf.math.multiply_4 (TFOpLambda  (None, 34, 16, 48)  0           ['batch_normalization_4[0][0]']  \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.truediv_4 (TFOpLambda)  (None, 34, 16, 48)  0           ['tf.math.multiply_4[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.maximum_8 (TFOpLambda)  (None, 34, 16, 48)  0           ['tf.math.truediv_4[0][0]',      \n",
      "                                                                  'batch_normalization_4[0][0]']  \n",
      "                                                                                                  \n",
      " tf.math.maximum_9 (TFOpLambda)  (None, 34, 16, 48)  0           ['tf.math.maximum_8[0][0]']      \n",
      "                                                                                                  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " conv2d_5 (Conv2D)              (None, 34, 16, 48)   6912        ['tf.math.maximum_9[0][0]']      \n",
      "                                                                                                  \n",
      " batch_normalization_5 (BatchNo  (None, 34, 16, 48)  192         ['conv2d_5[0][0]']               \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " tf.math.multiply_5 (TFOpLambda  (None, 34, 16, 48)  0           ['batch_normalization_5[0][0]']  \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.truediv_5 (TFOpLambda)  (None, 34, 16, 48)  0           ['tf.math.multiply_5[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.maximum_10 (TFOpLambda  (None, 34, 16, 48)  0           ['tf.math.truediv_5[0][0]',      \n",
      " )                                                                'batch_normalization_5[0][0]']  \n",
      "                                                                                                  \n",
      " tf.math.maximum_11 (TFOpLambda  (None, 34, 16, 48)  0           ['tf.math.maximum_10[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv2d_6 (Conv2D)              (None, 32, 16, 48)   6912        ['tf.math.maximum_11[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_6 (BatchNo  (None, 32, 16, 48)  192         ['conv2d_6[0][0]']               \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " tf.math.multiply_6 (TFOpLambda  (None, 32, 16, 48)  0           ['batch_normalization_6[0][0]']  \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.truediv_6 (TFOpLambda)  (None, 32, 16, 48)  0           ['tf.math.multiply_6[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.maximum_12 (TFOpLambda  (None, 32, 16, 48)  0           ['tf.math.truediv_6[0][0]',      \n",
      " )                                                                'batch_normalization_6[0][0]']  \n",
      "                                                                                                  \n",
      " tf.math.maximum_13 (TFOpLambda  (None, 32, 16, 48)  0           ['tf.math.maximum_12[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " max_pooling2d_1 (MaxPooling2D)  (None, 32, 8, 48)   0           ['tf.math.maximum_13[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_7 (Conv2D)              (None, 32, 8, 72)    10368       ['max_pooling2d_1[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_7 (BatchNo  (None, 32, 8, 72)   288         ['conv2d_7[0][0]']               \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " tf.math.multiply_7 (TFOpLambda  (None, 32, 8, 72)   0           ['batch_normalization_7[0][0]']  \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.truediv_7 (TFOpLambda)  (None, 32, 8, 72)   0           ['tf.math.multiply_7[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.maximum_14 (TFOpLambda  (None, 32, 8, 72)   0           ['tf.math.truediv_7[0][0]',      \n",
      " )                                                                'batch_normalization_7[0][0]']  \n",
      "                                                                                                  \n",
      " tf.math.maximum_15 (TFOpLambda  (None, 32, 8, 72)   0           ['tf.math.maximum_14[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv2d_8 (Conv2D)              (None, 30, 8, 72)    15552       ['tf.math.maximum_15[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_8 (BatchNo  (None, 30, 8, 72)   288         ['conv2d_8[0][0]']               \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " tf.math.multiply_8 (TFOpLambda  (None, 30, 8, 72)   0           ['batch_normalization_8[0][0]']  \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.truediv_8 (TFOpLambda)  (None, 30, 8, 72)   0           ['tf.math.multiply_8[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.maximum_16 (TFOpLambda  (None, 30, 8, 72)   0           ['tf.math.truediv_8[0][0]',      \n",
      " )                                                                'batch_normalization_8[0][0]']  \n",
      "                                                                                                  \n",
      " tf.math.maximum_17 (TFOpLambda  (None, 30, 8, 72)   0           ['tf.math.maximum_16[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv2d_9 (Conv2D)              (None, 30, 8, 72)    15552       ['tf.math.maximum_17[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_9 (BatchNo  (None, 30, 8, 72)   288         ['conv2d_9[0][0]']               \n",
      " rmalization)                                                                                     \n",
      "                                                                                                  \n",
      " tf.math.multiply_9 (TFOpLambda  (None, 30, 8, 72)   0           ['batch_normalization_9[0][0]']  \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.truediv_9 (TFOpLambda)  (None, 30, 8, 72)   0           ['tf.math.multiply_9[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.maximum_18 (TFOpLambda  (None, 30, 8, 72)   0           ['tf.math.truediv_9[0][0]',      \n",
      " )                                                                'batch_normalization_9[0][0]']  \n",
      "                                                                                                  \n",
      " tf.math.maximum_19 (TFOpLambda  (None, 30, 8, 72)   0           ['tf.math.maximum_18[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv2d_10 (Conv2D)             (None, 28, 8, 72)    15552       ['tf.math.maximum_19[0][0]']     \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                                                                  \n",
      " batch_normalization_10 (BatchN  (None, 28, 8, 72)   288         ['conv2d_10[0][0]']              \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " tf.math.multiply_10 (TFOpLambd  (None, 28, 8, 72)   0           ['batch_normalization_10[0][0]'] \n",
      " a)                                                                                               \n",
      "                                                                                                  \n",
      " tf.math.truediv_10 (TFOpLambda  (None, 28, 8, 72)   0           ['tf.math.multiply_10[0][0]']    \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.maximum_20 (TFOpLambda  (None, 28, 8, 72)   0           ['tf.math.truediv_10[0][0]',     \n",
      " )                                                                'batch_normalization_10[0][0]'] \n",
      "                                                                                                  \n",
      " tf.math.maximum_21 (TFOpLambda  (None, 28, 8, 72)   0           ['tf.math.maximum_20[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " max_pooling2d_2 (MaxPooling2D)  (None, 14, 4, 72)   0           ['tf.math.maximum_21[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_11 (Conv2D)             (None, 14, 4, 96)    20736       ['max_pooling2d_2[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_11 (BatchN  (None, 14, 4, 96)   384         ['conv2d_11[0][0]']              \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " tf.math.multiply_11 (TFOpLambd  (None, 14, 4, 96)   0           ['batch_normalization_11[0][0]'] \n",
      " a)                                                                                               \n",
      "                                                                                                  \n",
      " tf.math.truediv_11 (TFOpLambda  (None, 14, 4, 96)   0           ['tf.math.multiply_11[0][0]']    \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.maximum_22 (TFOpLambda  (None, 14, 4, 96)   0           ['tf.math.truediv_11[0][0]',     \n",
      " )                                                                'batch_normalization_11[0][0]'] \n",
      "                                                                                                  \n",
      " tf.math.maximum_23 (TFOpLambda  (None, 14, 4, 96)   0           ['tf.math.maximum_22[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv2d_12 (Conv2D)             (None, 12, 4, 96)    27648       ['tf.math.maximum_23[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_12 (BatchN  (None, 12, 4, 96)   384         ['conv2d_12[0][0]']              \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " tf.math.multiply_12 (TFOpLambd  (None, 12, 4, 96)   0           ['batch_normalization_12[0][0]'] \n",
      " a)                                                                                               \n",
      "                                                                                                  \n",
      " tf.math.truediv_12 (TFOpLambda  (None, 12, 4, 96)   0           ['tf.math.multiply_12[0][0]']    \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.maximum_24 (TFOpLambda  (None, 12, 4, 96)   0           ['tf.math.truediv_12[0][0]',     \n",
      " )                                                                'batch_normalization_12[0][0]'] \n",
      "                                                                                                  \n",
      " tf.math.maximum_25 (TFOpLambda  (None, 12, 4, 96)   0           ['tf.math.maximum_24[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv2d_13 (Conv2D)             (None, 12, 4, 96)    27648       ['tf.math.maximum_25[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_13 (BatchN  (None, 12, 4, 96)   384         ['conv2d_13[0][0]']              \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " tf.math.multiply_13 (TFOpLambd  (None, 12, 4, 96)   0           ['batch_normalization_13[0][0]'] \n",
      " a)                                                                                               \n",
      "                                                                                                  \n",
      " tf.math.truediv_13 (TFOpLambda  (None, 12, 4, 96)   0           ['tf.math.multiply_13[0][0]']    \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.maximum_26 (TFOpLambda  (None, 12, 4, 96)   0           ['tf.math.truediv_13[0][0]',     \n",
      " )                                                                'batch_normalization_13[0][0]'] \n",
      "                                                                                                  \n",
      " tf.math.maximum_27 (TFOpLambda  (None, 12, 4, 96)   0           ['tf.math.maximum_26[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv2d_14 (Conv2D)             (None, 10, 4, 96)    27648       ['tf.math.maximum_27[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_14 (BatchN  (None, 10, 4, 96)   384         ['conv2d_14[0][0]']              \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " tf.math.multiply_14 (TFOpLambd  (None, 10, 4, 96)   0           ['batch_normalization_14[0][0]'] \n",
      " a)                                                                                               \n",
      "                                                                                                  \n",
      " tf.math.truediv_14 (TFOpLambda  (None, 10, 4, 96)   0           ['tf.math.multiply_14[0][0]']    \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.maximum_28 (TFOpLambda  (None, 10, 4, 96)   0           ['tf.math.truediv_14[0][0]',     \n",
      " )                                                                'batch_normalization_14[0][0]'] \n",
      "                                                                                                  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " tf.math.maximum_29 (TFOpLambda  (None, 10, 4, 96)   0           ['tf.math.maximum_28[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " max_pooling2d_3 (MaxPooling2D)  (None, 10, 2, 96)   0           ['tf.math.maximum_29[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_15 (Conv2D)             (None, 10, 2, 96)    27648       ['max_pooling2d_3[0][0]']        \n",
      "                                                                                                  \n",
      " batch_normalization_15 (BatchN  (None, 10, 2, 96)   384         ['conv2d_15[0][0]']              \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " tf.math.multiply_15 (TFOpLambd  (None, 10, 2, 96)   0           ['batch_normalization_15[0][0]'] \n",
      " a)                                                                                               \n",
      "                                                                                                  \n",
      " tf.math.truediv_15 (TFOpLambda  (None, 10, 2, 96)   0           ['tf.math.multiply_15[0][0]']    \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.maximum_30 (TFOpLambda  (None, 10, 2, 96)   0           ['tf.math.truediv_15[0][0]',     \n",
      " )                                                                'batch_normalization_15[0][0]'] \n",
      "                                                                                                  \n",
      " tf.math.maximum_31 (TFOpLambda  (None, 10, 2, 96)   0           ['tf.math.maximum_30[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv2d_16 (Conv2D)             (None, 8, 2, 96)     27648       ['tf.math.maximum_31[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_16 (BatchN  (None, 8, 2, 96)    384         ['conv2d_16[0][0]']              \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " tf.math.multiply_16 (TFOpLambd  (None, 8, 2, 96)    0           ['batch_normalization_16[0][0]'] \n",
      " a)                                                                                               \n",
      "                                                                                                  \n",
      " tf.math.truediv_16 (TFOpLambda  (None, 8, 2, 96)    0           ['tf.math.multiply_16[0][0]']    \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.maximum_32 (TFOpLambda  (None, 8, 2, 96)    0           ['tf.math.truediv_16[0][0]',     \n",
      " )                                                                'batch_normalization_16[0][0]'] \n",
      "                                                                                                  \n",
      " tf.math.maximum_33 (TFOpLambda  (None, 8, 2, 96)    0           ['tf.math.maximum_32[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv2d_17 (Conv2D)             (None, 8, 2, 96)     27648       ['tf.math.maximum_33[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_17 (BatchN  (None, 8, 2, 96)    384         ['conv2d_17[0][0]']              \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " tf.math.multiply_17 (TFOpLambd  (None, 8, 2, 96)    0           ['batch_normalization_17[0][0]'] \n",
      " a)                                                                                               \n",
      "                                                                                                  \n",
      " tf.math.truediv_17 (TFOpLambda  (None, 8, 2, 96)    0           ['tf.math.multiply_17[0][0]']    \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.maximum_34 (TFOpLambda  (None, 8, 2, 96)    0           ['tf.math.truediv_17[0][0]',     \n",
      " )                                                                'batch_normalization_17[0][0]'] \n",
      "                                                                                                  \n",
      " tf.math.maximum_35 (TFOpLambda  (None, 8, 2, 96)    0           ['tf.math.maximum_34[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " conv2d_18 (Conv2D)             (None, 6, 2, 96)     27648       ['tf.math.maximum_35[0][0]']     \n",
      "                                                                                                  \n",
      " batch_normalization_18 (BatchN  (None, 6, 2, 96)    384         ['conv2d_18[0][0]']              \n",
      " ormalization)                                                                                    \n",
      "                                                                                                  \n",
      " tf.math.multiply_18 (TFOpLambd  (None, 6, 2, 96)    0           ['batch_normalization_18[0][0]'] \n",
      " a)                                                                                               \n",
      "                                                                                                  \n",
      " tf.math.truediv_18 (TFOpLambda  (None, 6, 2, 96)    0           ['tf.math.multiply_18[0][0]']    \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.maximum_36 (TFOpLambda  (None, 6, 2, 96)    0           ['tf.math.truediv_18[0][0]',     \n",
      " )                                                                'batch_normalization_18[0][0]'] \n",
      "                                                                                                  \n",
      " tf.math.maximum_37 (TFOpLambda  (None, 6, 2, 96)    0           ['tf.math.maximum_36[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " max_pooling2d_4 (MaxPooling2D)  (None, 3, 1, 96)    0           ['tf.math.maximum_37[0][0]']     \n",
      "                                                                                                  \n",
      " conv2d_19 (Conv2D)             (None, 1, 1, 96)     27648       ['max_pooling2d_4[0][0]']        \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 332,088\n",
      "Trainable params: 329,448\n",
      "Non-trainable params: 2,640\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# Recreate the embedding model after the melspectrogram layers\n",
    "# That is, have the melspectrogram of the audio as the input instead of the raw audio\n",
    "\n",
    "# A custom function for the leaky relu activation function, to make exporting to ONNX/tflite easier\n",
    "def MyLeakyReLU(alpha = 0.20000000298023224*2):\n",
    "    return lambda x : tf.keras.backend.maximum(alpha * x/2, x)\n",
    "\n",
    "# Define convolutional block helper functions\n",
    "def batch_norm_and_activation(x):\n",
    "    x = tf.keras.layers.BatchNormalization()(x)\n",
    "    x = MyLeakyReLU()(x)\n",
    "    x = tf.maximum(x, -0.4000000059604645)\n",
    "    return x\n",
    "\n",
    "# Define contraint for zero mean conv2d layer\n",
    "class CenterAround(tf.keras.constraints.Constraint):\n",
    "    \"\"\"Constrains weight tensors to be centered around `ref_value`.\"\"\"\n",
    "    def __init__(self, ref_value):\n",
    "        self.ref_value = ref_value\n",
    "\n",
    "    def __call__(self, w):\n",
    "        mean = tf.reduce_mean(w, axis=(0,1))\n",
    "        return w - mean + self.ref_value\n",
    "\n",
    "\n",
    "# Contruct inputs\n",
    "inputs = tf.keras.Input((76, 32, 1)) # melspectrogram shape when provided with 12400 samples at 16 khz\n",
    "\n",
    "# Input conv block\n",
    "x = tf.keras.layers.ZeroPadding2D((0,1))(inputs)\n",
    "x = tf.keras.layers.Conv2D(24, (3,3), use_bias=False, kernel_constraint=CenterAround(0.0),\n",
    "                           activation='relu', padding='valid')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "\n",
    "# Conv block #1\n",
    "x = tf.keras.layers.Conv2D(24, (1,3), use_bias=False, padding='same')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.Conv2D(24, (3,1), use_bias=False, padding='valid')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.MaxPool2D((2,2), (2,2), padding='valid')(x)\n",
    "x = tf.keras.layers.Conv2D(48, (1,3), use_bias=False, padding='same')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.Conv2D(48, (3,1), use_bias=False, padding='valid')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "\n",
    "# Conv block #2\n",
    "x = tf.keras.layers.Conv2D(48, (1,3), use_bias=False, padding='same')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.Conv2D(48, (3,1), use_bias=False, padding='valid')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.MaxPool2D((1,2), (1,2), padding='same')(x)\n",
    "x = tf.keras.layers.Conv2D(72, (1,3), use_bias=False, padding='same')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.Conv2D(72, (3,1), use_bias=False, padding='valid')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "\n",
    "# Conv block #3\n",
    "x = tf.keras.layers.Conv2D(72, (1,3), use_bias=False, padding='same')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.Conv2D(72, (3,1), use_bias=False, padding='valid')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.MaxPool2D((2,2), (2,2), padding='valid')(x)\n",
    "x = tf.keras.layers.Conv2D(96, (1,3), use_bias=False, padding='same')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.Conv2D(96, (3,1), use_bias=False, padding='valid')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "\n",
    "# Conv block #4\n",
    "x = tf.keras.layers.Conv2D(96, (1,3), use_bias=False, padding='same')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.Conv2D(96, (3,1), use_bias=False, padding='valid')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.MaxPool2D((1,2), (1,2), padding='valid')(x)\n",
    "x = tf.keras.layers.Conv2D(96, (1,3), use_bias=False, padding='same')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.Conv2D(96, (3,1), use_bias=False, padding='valid')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "\n",
    "# Conv block #5\n",
    "x = tf.keras.layers.Conv2D(96, (1,3), use_bias=False, padding='same')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.Conv2D(96, (3,1), use_bias=False, padding='valid')(x)\n",
    "x = batch_norm_and_activation(x)\n",
    "x = tf.keras.layers.MaxPool2D((2,2), (2,2), padding=\"valid\")(x)\n",
    "x = tf.keras.layers.Conv2D(96, (3,1), use_bias=False, padding='valid')(x)\n",
    "\n",
    "# Build the keras model\n",
    "reimplemented_model = tf.keras.Model(inputs=inputs, outputs=x)\n",
    "reimplemented_model.summary()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3a83d707",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-18T00:28:05.030126Z",
     "start_time": "2024-01-18T00:28:05.008907Z"
    }
   },
   "outputs": [],
   "source": [
    "# Manually set the weights of the new Keras model with those from the original embedding model\n",
    "\n",
    "# Set weights for all layers\n",
    "reimplemented_model.set_weights(embedding_model.get_weights())\n",
    "\n",
    "# Adjust weights of specific layer that needs to be centered around 0.0\n",
    "reimplemented_model.layers[2].set_weights([CenterAround(0.0)(reimplemented_model.layers[2].weights[0])])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eba224b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the new keras model to tflite format (optional for this notebook)\n",
    "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
    "tflite_model = converter.convert()\n",
    "\n",
    "# Save the model.\n",
    "with open('embedding_model.tflite', 'wb') as f:\n",
    "    f.write(tflite_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cfc246b",
   "metadata": {},
   "source": [
    "# Compare Predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4907f757",
   "metadata": {},
   "source": [
    "Now that we have a re-implemented embedding model, we can verify that the predictions are the same as the original. Note that as discussed previously, the log-mel feature calculation is different, so we will start from the original audio features obtained via tflite and calculate the final embeddings from there."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6d6ba4b5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-18T00:29:27.288771Z",
     "start_time": "2024-01-18T00:29:27.165734Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 78520 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Get original embedding model prediction\n",
    "original_embeddings = embedding_model(sample_data)\n",
    "\n",
    "# Reshape original log-mel inputs from tflite model above and pass to re-implemented model\n",
    "reimplemented_embeddings = reimplemented_model(spec.T[None, ..., None])\n",
    "\n",
    "# Plot final output embeddings for the sample data\n",
    "_ = plt.plot(original_embeddings.numpy().flatten(), label=\"Original Google Model\")\n",
    "_ = plt.plot(reimplemented_embeddings.numpy().flatten(), label=\"Reimplemented Model\")\n",
    "_ = plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ea6f9936",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-18T00:29:29.894151Z",
     "start_time": "2024-01-18T00:29:29.882876Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.00010585785"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check maximum absolute difference in the output embeddings to confirm practical equivalence\n",
    "np.abs(original_embeddings.numpy().flatten() - reimplemented_embeddings.numpy().flatten()).max()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "openwakeword_dev",
   "language": "python",
   "name": "openwakeword_dev"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "384px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
