{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "z6uizglMgQ26"
      },
      "source": [
        "# Training a Custom TensorFlow.js Audio Model\n",
        "\n",
        "In this notebook, we show how to train a custom audio model based on the model topology of the\n",
        "[TensorFlow.js Speech Commands model](https://www.npmjs.com/package/@tensorflow-models/speech-commands).\n",
        "The training is done in Python by using a set of audio examples stored as .wav files.\n",
        "The trained model is convertible to the\n",
        "[TensorFlow.js LayersModel](https://js.tensorflow.org/api/latest/#loadLayersModel) format for\n",
        "inference and further fine-tuning in the browser.\n",
        "It may also be converted to the [TFLite](https://www.tensorflow.org/lite) format\n",
        "for inference on mobile devices.\n",
        "\n",
        "This example uses a small subset of the\n",
        "[Speech Commands v0.02](https://arxiv.org/abs/1804.03209) dataset, and builds\n",
        "a model that detects two English words (\"yes\" and \"no\") against background noises. But the methodology demonstrated here is general and can be applied to\n",
        "other sounds, as long as they are stored in the same .wav file format as in this example.\n",
        "\n",
        "## Data format\n",
        "\n",
        "The training procedure in this notebook makes the following assumption about the raw audio data:\n",
        "\n",
        "1. The root data directory contains a number of folders. The name of each folder is the name\n",
        "   of the audio class. You can select any subset of the folders (i.e., classes) to train the\n",
        "   model on.\n",
        "2. Within each folder, there are a number of .wav files. Each .wav file corresponds to an\n",
        "   example. Each .wav file is mono (single-channel) and has the typical pulse-code modulation\n",
        "   (PCM) encoding. The duration of each wave file should be 1 second or slightly longer. \n",
        "3. There can be a special folder called \"_background_noise_\" that contains .wav files for\n",
        "   audio samples that fall into the background noise class. Each of these .wav files can be\n",
        "   much longer than 1 second in duration. This notebook contains code that extracts 1-second\n",
        "   snippets from these .wav files\n",
        "   \n",
        "The Speech Commands v0.3 dataset used in this notebook meets these data format requirements."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Grv5UK5rHxyY"
      },
      "outputs": [],
      "source": [
        "!pip install librosa tensorflowjs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3BbsTxcuCwYO"
      },
      "outputs": [],
      "source": [
        "import glob\n",
        "import json\n",
        "import os\n",
        "import random\n",
        "\n",
        "import librosa\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "from scipy import signal\n",
        "from scipy.io import wavfile\n",
        "import tensorflow as tf\n",
        "import tensorflowjs as tfjs\n",
        "import tqdm\n",
        "\n",
        "print(tf.__version__)\n",
        "print(tfjs.__version__)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 119
        },
        "colab_type": "code",
        "id": "wkPnHDHITAJH",
        "outputId": "8c64930f-b03e-48df-fc93-7ab894307ee7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "sc_preproc_model/\n",
            "sc_preproc_model/assets/\n",
            "sc_preproc_model/variables/\n",
            "sc_preproc_model/variables/variables.data-00000-of-00001\n",
            "sc_preproc_model/variables/variables.index\n",
            "sc_preproc_model/saved_model.pb\n"
          ]
        }
      ],
      "source": [
        "# Download the TensorFlow.js Speech Commands model and the associated\n",
        "# preprocesssing model.\n",
        "!mkdir -p /tmp/tfjs-sc-model\n",
        "!curl -o /tmp/tfjs-sc-model/metadata.json -fsSL https://storage.googleapis.com/tfjs-models/tfjs/speech-commands/v0.3/browser_fft/18w/metadata.json\n",
        "!curl -o /tmp/tfjs-sc-model/model.json -fsSL https://storage.googleapis.com/tfjs-models/tfjs/speech-commands/v0.3/browser_fft/18w/model.json\n",
        "!curl -o /tmp/tfjs-sc-model/group1-shard1of2 -fSsL https://storage.googleapis.com/tfjs-models/tfjs/speech-commands/v0.3/browser_fft/18w/group1-shard1of2\n",
        "!curl -o /tmp/tfjs-sc-model/group1-shard2of2 -fsSL https://storage.googleapis.com/tfjs-models/tfjs/speech-commands/v0.3/browser_fft/18w/group1-shard2of2\n",
        "!curl -o /tmp/tfjs-sc-model/sc_preproc_model.tar.gz -fSsL https://storage.googleapis.com/tfjs-models/tfjs/speech-commands/conversion/sc_preproc_model.tar.gz\n",
        "!cd /tmp/tfjs-sc-model/ && tar xzvf sc_preproc_model.tar.gz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ioTiCDp4HO_V"
      },
      "outputs": [],
      "source": [
        "# Download Speech Commands v0.02 dataset. The dataset contains 30+ word and\n",
        "# sound categories, but we will only use a subset of them\n",
        "\n",
        "!mkdir -p /tmp/speech_commands_v0.02\n",
        "!curl -o /tmp/speech_commands_v0.02/speech_commands_v0.02.tar.gz -fSsL http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz\n",
        "!cd  /tmp/speech_commands_v0.02 && tar xzf speech_commands_v0.02.tar.gz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 221
        },
        "colab_type": "code",
        "id": "TqnjnnPoTR8E",
        "outputId": "447dd1df-edc1-4829-ce31-4e8f2f9bb328"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.\n",
            "Model: \"audio_preproc\"\n",
            "_________________________________________________________________\n",
            "Layer (type)                 Output Shape              Param #   \n",
            "=================================================================\n",
            "audio_preprocessing_layer (A (None, None, None, 1)     2048      \n",
            "=================================================================\n",
            "Total params: 2,048\n",
            "Trainable params: 0\n",
            "Non-trainable params: 2,048\n",
            "_________________________________________________________________\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "(None, 44032)"
            ]
          },
          "execution_count": 7,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Load the preprocessing model, which transforms audio waveform into \n",
        "# spectrograms (2D image-like representation of sound).\n",
        "# This preprocessing model replicates WebAudio's AnalyzerNode.getFloatFrequencyData\n",
        "# (https://developer.mozilla.org/en-US/docs/Web/API/AnalyserNode/getFloatFrequencyData).\n",
        "# It performs short-time Fourier transform (STFT) using a length-2048 Blackman\n",
        "# window. It opeartes on mono audio at the 44100-Hz sample rate.\n",
        "\n",
        "preproc_model_path = '/tmp/tfjs-sc-model/sc_preproc_model'\n",
        "preproc_model = tf.keras.models.load_model(preproc_model_path)\n",
        "preproc_model.summary()\n",
        "preproc_model.input_shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mihiFei-dE3u"
      },
      "outputs": [],
      "source": [
        "# Create some constants to be used later.\n",
        "\n",
        "# Target sampling rate. It is required by the audio preprocessing model.\n",
        "TARGET_SAMPLE_RATE = 44100\n",
        "# The specific audio tensor length expected by the preprocessing model.\n",
        "EXPECTED_WAVEFORM_LEN = preproc_model.input_shape[-1]\n",
        "\n",
        "# Where the Speech Commands v0.02 dataset has been downloaded.\n",
        "DATA_ROOT = \"/tmp/speech_commands_v0.02\"\n",
        "\n",
        "WORDS = (\"_background_noise_snippets_\", \"no\", \"yes\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 153
        },
        "colab_type": "code",
        "id": "0jl1n0SCNYUj",
        "outputId": "1a524c46-ddd4-4162-f60a-c2511fd14626"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Extracting snippets from /tmp/speech_commands_v0.02/_background_noise_/pink_noise.wav...\n",
            "Extracting snippets from /tmp/speech_commands_v0.02/_background_noise_/exercise_bike.wav...\n",
            "Extracting snippets from /tmp/speech_commands_v0.02/_background_noise_/dude_miaowing.wav...\n",
            "Extracting snippets from /tmp/speech_commands_v0.02/_background_noise_/running_tap.wav...\n",
            "Extracting snippets from /tmp/speech_commands_v0.02/_background_noise_/doing_the_dishes.wav...\n",
            "Extracting snippets from /tmp/speech_commands_v0.02/_background_noise_/white_noise.wav...\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:13: WavFileWarning: Chunk (non-data) not understood, skipping it.\n",
            "  del sys.path[0]\n"
          ]
        }
      ],
      "source": [
        "# Unlike word examples, the noise samples in the Speech Commands v0.02 dataset\n",
        "# are not divided into 1-second snippets. Instead, they are stored as longer\n",
        "# recordings. Therefore we need to cut them up in to 1-second snippet .wav\n",
        "# files.\n",
        "\n",
        "noise_wav_paths = glob.glob(os.path.join(DATA_ROOT, \"_background_noise_\", \"*.wav\"))\n",
        "snippets_dir = os.path.join(DATA_ROOT, \"_background_noise_snippets_\")\n",
        "os.makedirs(snippets_dir, exist_ok=True)\n",
        "\n",
        "\n",
        "def extract_snippets(wav_path, snippet_duration_sec=1.0):\n",
        "  basename = os.path.basename(os.path.splitext(wav_path)[0])\n",
        "  sample_rate, xs = wavfile.read(wav_path)\n",
        "  assert xs.dtype == np.int16\n",
        "  n_samples_per_snippet = int(snippet_duration_sec * sample_rate)\n",
        "  i = 0\n",
        "  while i + n_samples_per_snippet < len(xs):\n",
        "    snippet_wav_path = os.path.join(snippets_dir, \"%s_%.5d.wav\" % (basename, i))\n",
        "    snippet = xs[i : i + n_samples_per_snippet].astype(np.int16)\n",
        "    wavfile.write(snippet_wav_path, sample_rate, snippet)\n",
        "    i += n_samples_per_snippet\n",
        "\n",
        "for noise_wav_path in noise_wav_paths:\n",
        "  print(\"Extracting snippets from %s...\" % noise_wav_path)\n",
        "  extract_snippets(noise_wav_path, snippet_duration_sec=1.0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "FSjiThysWrTx"
      },
      "outputs": [],
      "source": [
        "def resample_wavs(dir_path, target_sample_rate=44100):\n",
        "  \"\"\"Resample the .wav files in an input directory to given sampling rate.\n",
        "  \n",
        "  The resampled waveforms are written to .wav files in the same directory with\n",
        "  file names that ends in \"_44100hz.wav\".\n",
        "\n",
        "  44100 Hz is the sample rate required by the preprocessing model. It is also\n",
        "  the most widely supported sample rate among web browsers and mobile devices.\n",
        "  For example, see:\n",
        "  https://developer.mozilla.org/en-US/docs/Web/API/AudioContextOptions/sampleRate\n",
        "  https://developer.android.com/ndk/guides/audio/sampling-audio\n",
        "\n",
        "  Args:\n",
        "    dir_path: Path to a directory that contains .wav files.\n",
        "    target_sapmle_rate: Target sampling rate in Hz.\n",
        "  \"\"\"\n",
        "  wav_paths = glob.glob(os.path.join(dir_path, \"*.wav\"))\n",
        "  resampled_suffix = \"_%shz.wav\" % target_sample_rate\n",
        "  for i, wav_path in tqdm.tqdm(enumerate(wav_paths)):\n",
        "    if wav_path.endswith(resampled_suffix):\n",
        "      continue\n",
        "    sample_rate, xs = wavfile.read(wav_path)\n",
        "    xs = xs.astype(np.float32)\n",
        "    xs = librosa.resample(xs, sample_rate, TARGET_SAMPLE_RATE).astype(np.int16)\n",
        "    resampled_path = os.path.splitext(wav_path)[0] + resampled_suffix\n",
        "    wavfile.write(resampled_path, target_sample_rate, xs)\n",
        "\n",
        "\n",
        "for word in WORDS:\n",
        "  word_dir = os.path.join(DATA_ROOT, word)\n",
        "  assert os.path.isdir(word_dir)\n",
        "  resample_wavs(word_dir, target_sample_rate=TARGET_SAMPLE_RATE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aFA-TSmpK935"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def read_wav(filepath):\n",
        "  file_contents = tf.io.read_file(filepath)\n",
        "  return tf.expand_dims(tf.squeeze(tf.audio.decode_wav(\n",
        "      file_contents, \n",
        "      desired_channels=-1,\n",
        "      desired_samples=TARGET_SAMPLE_RATE).audio, axis=-1), 0)\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def filter_by_waveform_length(waveform, label):\n",
        "  return tf.size(waveform) > EXPECTED_WAVEFORM_LEN\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def crop_and_convert_to_spectrogram(waveform, label):\n",
        "  cropped = tf.slice(waveform, begin=[0, 0], size=[1, EXPECTED_WAVEFORM_LEN])\n",
        "  return tf.squeeze(preproc_model(cropped), axis=0), label\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def spectrogram_elements_finite(spectrogram, label):\n",
        "  return tf.math.reduce_all(tf.math.is_finite(spectrogram))\n",
        "\n",
        "\n",
        "def get_dataset(input_wav_paths, labels):\n",
        "  \"\"\"Get a tf.data.Dataset given input .wav files and their labels.\n",
        "\n",
        "  The returned dataset emits 2-tuples of `(spectrogram, label)`, wherein\n",
        "  - `spectrogram` is a tensor of dtype tf.float32 and shape [43, 232, 1].\n",
        "    It is z-normalized (i.e., have a mean of ~0.0 and variance of ~1.0).\n",
        "  - `label` is a tensor of dtype tf.int32 and shape [] (scalar).\n",
        "  \n",
        "  Args:\n",
        "    input_wav_paths: Input audio .wav file paths as a list of string.\n",
        "    labels: integer labels (class indices) of the input .wav files. Must have\n",
        "      the same lengh as `input_wav_paths`.\n",
        "\n",
        "  Returns:\n",
        "    A tf.data.Dataset object as described above.\n",
        "  \"\"\"\n",
        "  ds = tf.data.Dataset.from_tensor_slices(input_wav_paths)\n",
        "  # Read audio waveform from the .wav files.\n",
        "  ds = ds.map(read_wav)\n",
        "  ds = tf.data.Dataset.zip((ds, tf.data.Dataset.from_tensor_slices(labels)))\n",
        "  # Keep only the waveforms longer than `EXPECTED_WAVEFORM_LEN`.\n",
        "  ds = ds.filter(filter_by_waveform_length)\n",
        "  # Crop the waveforms to `EXPECTED_WAVEFORM_LEN` and convert them to\n",
        "  # spectrograms using the preprocessing layer.\n",
        "  ds = ds.map(crop_and_convert_to_spectrogram)\n",
        "  # Discard examples that contain infinite or NaN elements.\n",
        "  ds = ds.filter(spectrogram_elements_finite)\n",
        "  return ds"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 333
        },
        "colab_type": "code",
        "id": "tU6gho3nuvQl",
        "outputId": "07dea6f2-9b7c-477e-8814-f69772fb39da"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Found 396 examples for class _background_noise_snippets_\n",
            "Found 3941 examples for class no\n",
            "Found 4044 examples for class yes\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 707632 characters]\n",
            "text/plain": [
              "<Figure size 2880x7200 with 10 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light",
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "input_wav_paths_and_labels = []\n",
        "for i, word in enumerate(WORDS):\n",
        "  wav_paths = glob.glob(os.path.join(DATA_ROOT, word, \"*_%shz.wav\" % TARGET_SAMPLE_RATE))\n",
        "  print(\"Found %d examples for class %s\" % (len(wav_paths), word))\n",
        "  labels = [i] * len(wav_paths)\n",
        "  input_wav_paths_and_labels.extend(zip(wav_paths, labels))\n",
        "random.shuffle(input_wav_paths_and_labels)\n",
        "  \n",
        "input_wav_paths, labels = ([t[0] for t in input_wav_paths_and_labels],\n",
        "                           [t[1] for t in input_wav_paths_and_labels])\n",
        "dataset = get_dataset(input_wav_paths, labels)\n",
        "\n",
        "# Show some example spectrograms for inspection.\n",
        "fig = plt.figure(figsize=(40, 100))\n",
        "dataset_iter = iter(dataset)\n",
        "num_spectrograms_to_show = 10\n",
        "for i in range(num_spectrograms_to_show):\n",
        "  ax = fig.add_subplot(1, num_spectrograms_to_show, i + 1)\n",
        "  spectrogram, label = next(dataset_iter)\n",
        "  spectrogram = spectrogram.numpy()\n",
        "  label = label.numpy()\n",
        "  plt.imshow(np.flipud(np.squeeze(spectrogram, -1).T), aspect=0.2)\n",
        "  ax.set_title(\"Example of \\\"%s\\\"\" % WORDS[label])\n",
        "  ax.set_xlabel(\"Time frame #\")\n",
        "  if i == 0:\n",
        "    ax.set_ylabel(\"Frequency bin #\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "colab_type": "code",
        "id": "ZvWr-_R2ym7d",
        "outputId": "4784a0d3-d0a3-406a-9199-2c626efe454f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading dataset and converting data to numpy arrays. This may take a few minutes...\n",
            "Done.\n"
          ]
        }
      ],
      "source": [
        "# The amount of data we have is relatively small. It fits into typical host RAM\n",
        "# or GPU memory. For better training performance, we preload the data and\n",
        "# put it into numpy arrays:\n",
        "# - xs: The audio features (normalized spectrograms).\n",
        "# - ys: The labels (class indices).\n",
        "print(\n",
        "    \"Loading dataset and converting data to numpy arrays. \"\n",
        "    \"This may take a few minutes...\")\n",
        "xs_and_ys = list(dataset)\n",
        "xs = np.stack([item[0] for item in xs_and_ys])\n",
        "ys = np.stack([item[1] for item in xs_and_ys])\n",
        "print(\"Done.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 595
        },
        "colab_type": "code",
        "id": "o6sV5t2Kwi7p",
        "outputId": "d127e5ce-24ef-4ced-bd2e-19f8ca31e586"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Model: \"TransferLearnedModel\"\n",
            "_________________________________________________________________\n",
            "Layer (type)                 Output Shape              Param #   \n",
            "=================================================================\n",
            "conv2d_1 (Conv2D)            (None, 42, 225, 8)        136       \n",
            "_________________________________________________________________\n",
            "max_pooling2d_1 (MaxPooling2 (None, 21, 112, 8)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_2 (Conv2D)            (None, 20, 109, 32)       2080      \n",
            "_________________________________________________________________\n",
            "max_pooling2d_2 (MaxPooling2 (None, 10, 54, 32)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_3 (Conv2D)            (None, 9, 51, 32)         8224      \n",
            "_________________________________________________________________\n",
            "max_pooling2d_3 (MaxPooling2 (None, 4, 25, 32)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_4 (Conv2D)            (None, 3, 22, 32)         8224      \n",
            "_________________________________________________________________\n",
            "max_pooling2d_4 (MaxPooling2 (None, 2, 11, 32)         0         \n",
            "_________________________________________________________________\n",
            "flatten_1 (Flatten)          (None, 704)               0         \n",
            "_________________________________________________________________\n",
            "dropout_1 (Dropout)          (None, 704)               0         \n",
            "_________________________________________________________________\n",
            "dense_1 (Dense)              (None, 2000)              1410000   \n",
            "_________________________________________________________________\n",
            "dropout_2 (Dropout)          (None, 2000)              0         \n",
            "_________________________________________________________________\n",
            "dense (Dense)                (None, 3)                 6003      \n",
            "=================================================================\n",
            "Total params: 1,434,667\n",
            "Trainable params: 6,003\n",
            "Non-trainable params: 1,428,664\n",
            "_________________________________________________________________\n"
          ]
        }
      ],
      "source": [
        "tfjs_model_json_path = '/tmp/tfjs-sc-model/model.json'\n",
        "\n",
        "# Load the Speech Commands model. Weights are loaded along with the topology,\n",
        "# since we train the model from scratch. Instead, we will perform transfer\n",
        "# learning based on the model.\n",
        "orig_model = tfjs.converters.load_keras_model(tfjs_model_json_path, load_weights=True)\n",
        "\n",
        "# Remove the top Dense layer and add a new Dense layer of which the output\n",
        "# size fits the number of sound classes we care about.\n",
        "model = tf.keras.Sequential(name=\"TransferLearnedModel\")\n",
        "for layer in orig_model.layers[:-1]:\n",
        "  model.add(layer)\n",
        "model.add(tf.keras.layers.Dense(units=len(WORDS), activation=\"softmax\"))\n",
        "\n",
        "# Freeze all but the last layer of the model. The last layer will be fine-tuned\n",
        "# during transfer learning.\n",
        "for layer in model.layers[:-1]:\n",
        "  layer.trainable = False\n",
        "\n",
        "model.compile(optimizer=\"sgd\", loss=\"sparse_categorical_crossentropy\", metrics=[\"acc\"])\n",
        "model.summary()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "colab_type": "code",
        "id": "zRNtKeuPFIaq",
        "outputId": "7cd176f4-db38-448f-a153-7d78dcbd24fc"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch 1/50\n",
            "21/21 [==============================] - 11s 538ms/step - loss: 0.4393 - acc: 0.8521 - val_loss: 0.0732 - val_acc: 0.9756\n",
            "Epoch 2/50\n",
            "21/21 [==============================] - 11s 527ms/step - loss: 0.1562 - acc: 0.9438 - val_loss: 0.0571 - val_acc: 0.9823\n",
            "Epoch 3/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.1207 - acc: 0.9584 - val_loss: 0.0536 - val_acc: 0.9810\n",
            "Epoch 4/50\n",
            "21/21 [==============================] - 11s 532ms/step - loss: 0.0987 - acc: 0.9653 - val_loss: 0.0471 - val_acc: 0.9827\n",
            "Epoch 5/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0894 - acc: 0.9662 - val_loss: 0.0455 - val_acc: 0.9836\n",
            "Epoch 6/50\n",
            "21/21 [==============================] - 11s 527ms/step - loss: 0.0818 - acc: 0.9730 - val_loss: 0.0442 - val_acc: 0.9841\n",
            "Epoch 7/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0791 - acc: 0.9738 - val_loss: 0.0422 - val_acc: 0.9849\n",
            "Epoch 8/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0749 - acc: 0.9757 - val_loss: 0.0407 - val_acc: 0.9858\n",
            "Epoch 9/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0748 - acc: 0.9736 - val_loss: 0.0399 - val_acc: 0.9863\n",
            "Epoch 10/50\n",
            "21/21 [==============================] - 11s 532ms/step - loss: 0.0698 - acc: 0.9761 - val_loss: 0.0400 - val_acc: 0.9858\n",
            "Epoch 11/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0670 - acc: 0.9774 - val_loss: 0.0389 - val_acc: 0.9867\n",
            "Epoch 12/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0676 - acc: 0.9768 - val_loss: 0.0385 - val_acc: 0.9876\n",
            "Epoch 13/50\n",
            "21/21 [==============================] - 11s 527ms/step - loss: 0.0655 - acc: 0.9770 - val_loss: 0.0378 - val_acc: 0.9872\n",
            "Epoch 14/50\n",
            "21/21 [==============================] - 11s 527ms/step - loss: 0.0615 - acc: 0.9799 - val_loss: 0.0360 - val_acc: 0.9876\n",
            "Epoch 15/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0651 - acc: 0.9765 - val_loss: 0.0362 - val_acc: 0.9876\n",
            "Epoch 16/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0477 - acc: 0.9818 - val_loss: 0.0361 - val_acc: 0.9880\n",
            "Epoch 17/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0561 - acc: 0.9801 - val_loss: 0.0357 - val_acc: 0.9889\n",
            "Epoch 18/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0521 - acc: 0.9816 - val_loss: 0.0350 - val_acc: 0.9885\n",
            "Epoch 19/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0659 - acc: 0.9803 - val_loss: 0.0345 - val_acc: 0.9885\n",
            "Epoch 20/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0544 - acc: 0.9812 - val_loss: 0.0358 - val_acc: 0.9872\n",
            "Epoch 21/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0573 - acc: 0.9787 - val_loss: 0.0333 - val_acc: 0.9885\n",
            "Epoch 22/50\n",
            "21/21 [==============================] - 11s 530ms/step - loss: 0.0442 - acc: 0.9839 - val_loss: 0.0332 - val_acc: 0.9889\n",
            "Epoch 23/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0528 - acc: 0.9820 - val_loss: 0.0335 - val_acc: 0.9889\n",
            "Epoch 24/50\n",
            "21/21 [==============================] - 11s 530ms/step - loss: 0.0411 - acc: 0.9861 - val_loss: 0.0343 - val_acc: 0.9894\n",
            "Epoch 25/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0535 - acc: 0.9812 - val_loss: 0.0333 - val_acc: 0.9889\n",
            "Epoch 26/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0418 - acc: 0.9848 - val_loss: 0.0341 - val_acc: 0.9889\n",
            "Epoch 27/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0484 - acc: 0.9820 - val_loss: 0.0334 - val_acc: 0.9885\n",
            "Epoch 28/50\n",
            "21/21 [==============================] - 11s 527ms/step - loss: 0.0474 - acc: 0.9823 - val_loss: 0.0339 - val_acc: 0.9885\n",
            "Epoch 29/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0452 - acc: 0.9831 - val_loss: 0.0334 - val_acc: 0.9889\n",
            "Epoch 30/50\n",
            "21/21 [==============================] - 11s 531ms/step - loss: 0.0465 - acc: 0.9833 - val_loss: 0.0325 - val_acc: 0.9889\n",
            "Epoch 31/50\n",
            "21/21 [==============================] - 11s 530ms/step - loss: 0.0419 - acc: 0.9860 - val_loss: 0.0327 - val_acc: 0.9889\n",
            "Epoch 32/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0415 - acc: 0.9854 - val_loss: 0.0326 - val_acc: 0.9889\n",
            "Epoch 33/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0450 - acc: 0.9842 - val_loss: 0.0320 - val_acc: 0.9898\n",
            "Epoch 34/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0419 - acc: 0.9863 - val_loss: 0.0325 - val_acc: 0.9889\n",
            "Epoch 35/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0447 - acc: 0.9841 - val_loss: 0.0318 - val_acc: 0.9898\n",
            "Epoch 36/50\n",
            "21/21 [==============================] - 11s 530ms/step - loss: 0.0417 - acc: 0.9848 - val_loss: 0.0317 - val_acc: 0.9894\n",
            "Epoch 37/50\n",
            "21/21 [==============================] - 11s 532ms/step - loss: 0.0464 - acc: 0.9850 - val_loss: 0.0318 - val_acc: 0.9894\n",
            "Epoch 38/50\n",
            "21/21 [==============================] - 11s 530ms/step - loss: 0.0355 - acc: 0.9878 - val_loss: 0.0320 - val_acc: 0.9898\n",
            "Epoch 39/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0452 - acc: 0.9829 - val_loss: 0.0326 - val_acc: 0.9898\n",
            "Epoch 40/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0373 - acc: 0.9875 - val_loss: 0.0320 - val_acc: 0.9898\n",
            "Epoch 41/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0389 - acc: 0.9858 - val_loss: 0.0312 - val_acc: 0.9894\n",
            "Epoch 42/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0416 - acc: 0.9867 - val_loss: 0.0310 - val_acc: 0.9898\n",
            "Epoch 43/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0400 - acc: 0.9863 - val_loss: 0.0315 - val_acc: 0.9898\n",
            "Epoch 44/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0420 - acc: 0.9835 - val_loss: 0.0313 - val_acc: 0.9898\n",
            "Epoch 45/50\n",
            "21/21 [==============================] - 11s 528ms/step - loss: 0.0423 - acc: 0.9856 - val_loss: 0.0305 - val_acc: 0.9894\n",
            "Epoch 46/50\n",
            "21/21 [==============================] - 11s 530ms/step - loss: 0.0353 - acc: 0.9871 - val_loss: 0.0311 - val_acc: 0.9894\n",
            "Epoch 47/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0439 - acc: 0.9837 - val_loss: 0.0329 - val_acc: 0.9898\n",
            "Epoch 48/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0409 - acc: 0.9858 - val_loss: 0.0310 - val_acc: 0.9894\n",
            "Epoch 49/50\n",
            "21/21 [==============================] - 11s 527ms/step - loss: 0.0386 - acc: 0.9858 - val_loss: 0.0298 - val_acc: 0.9898\n",
            "Epoch 50/50\n",
            "21/21 [==============================] - 11s 529ms/step - loss: 0.0357 - acc: 0.9886 - val_loss: 0.0306 - val_acc: 0.9894\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "<tensorflow.python.keras.callbacks.History at 0x7f17cc7e1908>"
            ]
          },
          "execution_count": 20,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Train the model.\n",
        "model.fit(xs, ys, batch_size=256, validation_split=0.3, shuffle=True, epochs=50)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 139
        },
        "colab_type": "code",
        "id": "LFHnTUroi_3u",
        "outputId": "8a353c2d-b154-41fe-a53c-2d818bebe074"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/tensorflowjs/converters/keras_h5_conversion.py:123: H5pyDeprecationWarning: The default file mode will change to 'r' (read-only) in h5py 3.0. To suppress this warning, pass the mode you need to h5py.File(), or set the global default h5.get_config().default_file_mode, or set the environment variable H5PY_DEFAULT_READONLY=1. Available modes are: 'r', 'r+', 'w', 'w-'/'x', 'a'. See the docs for details.\n",
            "  return h5py.File(h5file)\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "total 5.5M\n",
            "-rw-r--r-- 1 root root 4.0M Aug 10 13:34 group1-shard1of2.bin\n",
            "-rw-r--r-- 1 root root 1.5M Aug 10 13:34 group1-shard2of2.bin\n",
            "-rw-r--r-- 1 root root   65 Aug 10 13:34 metadata.json\n",
            "-rw-r--r-- 1 root root 6.2K Aug 10 13:34 model.json\n"
          ]
        }
      ],
      "source": [
        "# Convert the model to TensorFlow.js Layers model format.\n",
        "\n",
        "tfjs_model_dir = \"/tmp/tfjs-model\"\n",
        "tfjs.converters.save_keras_model(model, tfjs_model_dir)\n",
        "\n",
        "# Create the metadata.json file.\n",
        "metadata = {\"words\": [\"_background_noise_\"] + WORDS[1:], \"frameSize\": model.input_shape[-2]}\n",
        "with open(os.path.join(tfjs_model_dir, \"metadata.json\"), \"w\") as f:\n",
        "  json.dump(metadata, f)\n",
        "\n",
        "!ls -lh /tmp/tfjs_model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "To deploy this model to the web, you can use the\n",
        "[speech-commands NPM package](https://www.npmjs.com/package/@tensorflow-models/speech-commands).\n",
        "\n",
        "The model.json and metadata.json should be hosted together with the two weights (.bin) files in the same HTTP/HTTPS directory.\n",
        "\n",
        "Then the custom model can be loaded in JavaScript with:\n",
        "\n",
        "```js\n",
        "import * as tf from '@tensorflow/tfjs-core';\n",
        "import * as tfl from '@tensorflow/tfjs-layers';\n",
        "import * as speechCommands from '@tensorflow-models/speech-commands';\n",
        "\n",
        "const recognizer = speechCommands.create(\n",
        "    'BROWSER_FFT',\n",
        "    null,\n",
        "    'http://test.com/my-audio-model/model.json',  // URL to the custom model's model.json\n",
        "    'http://test.com/my-audio-model/metadata.json'  // URL to the custom model's metadata.json\n",
        ");\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 357
        },
        "colab_type": "code",
        "id": "AZDMFkxulS8C",
        "outputId": "cbd28071-1a09-43f5-a16f-2b793e3fc4e1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Model: \"CombinedModel\"\n",
            "_________________________________________________________________\n",
            "Layer (type)                 Output Shape              Param #   \n",
            "=================================================================\n",
            "audio_preproc (Sequential)   (None, None, None, 1)     2048      \n",
            "_________________________________________________________________\n",
            "TransferLearnedModel (Sequen (None, 3)                 1434667   \n",
            "=================================================================\n",
            "Total params: 1,436,715\n",
            "Trainable params: 6,003\n",
            "Non-trainable params: 1,430,712\n",
            "_________________________________________________________________\n",
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/tracking.py:111: Model.state_updates (from tensorflow.python.keras.engine.training) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "This property should not be used in TensorFlow 2.0, as updates are applied automatically.\n",
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/tracking.py:111: Layer.updates (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "This property should not be used in TensorFlow 2.0, as updates are applied automatically.\n",
            "INFO:tensorflow:Assets written to: /tmp/tmp865tq81o/assets\n",
            "Saved tflite file at: /tmp/tfjs-sc-model/combined_model.tflite\n"
          ]
        }
      ],
      "source": [
        "# Convert the model to TFLite. \n",
        "\n",
        "# We need to combine the preprocessing model and the newly trained 3-class model\n",
        "# so that the resultant model will be able to preform STFT and spectrogram \n",
        "# calculation on mobile devices (i.e., without web browser's WebAudio).\n",
        "\n",
        "combined_model = tf.keras.Sequential(name='CombinedModel')\n",
        "combined_model.add(preproc_model)\n",
        "combined_model.add(model)\n",
        "combined_model.build([None, EXPECTED_WAVEFORM_LEN])\n",
        "combined_model.summary()\n",
        "\n",
        "tflite_output_path = '/tmp/tfjs-sc-model/combined_model.tflite'\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model(combined_model)\n",
        "converter.target_spec.supported_ops = [\n",
        "    tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS\n",
        "]\n",
        "with open(tflite_output_path, 'wb') as f:\n",
        "    f.write(converter.convert())\n",
        "print(\"Saved tflite file at: %s\" % tflite_output_path)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "training_custom_audio_model_in_python.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
