custaudio_training.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "custaudio_training.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/phrasenmaeher/custom-audio-classification-tf/blob/main/custaudio_training.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BZojfbQmfs6J"
      },
      "source": [
        "# Imports\n",
        "We start with the usual packages"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rnCgGBHx9xzx"
      },
      "source": [
        "Code for section 4 of the post at\n",
        "[TDS/Medium](https://towardsdatascience.com/custom-audio-classification-with-tensorflow-af8c16c38689)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0rWN4Cqxf4xR",
        "outputId": "8de708aa-4439-4776-c960-cbe9b78576ac"
      },
      "source": [
        "!pip install --upgrade kapre wandb"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting kapre\n",
            "  Downloading https://files.pythonhosted.org/packages/5b/a3/a80d9e09b67a728c167b787917813f4d80e0ad033697f9db237030f9a454/kapre-0.3.4.tar.gz\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting wandb\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/33/ae/79374d2b875e638090600eaa2a423479865b7590c53fb78e8ccf6a64acb1/wandb-0.10.22-py2.py3-none-any.whl (2.0MB)\n",
            "\u001b[K     |████████████████████████████████| 2.0MB 5.7MB/s \n",
            "\u001b[?25hRequirement already satisfied, skipping upgrade: tensorflow>=2.0 in /usr/local/lib/python3.7/dist-packages (from kapre) (2.4.1)\n",
            "Requirement already satisfied, skipping upgrade: librosa>=0.7.2 in /usr/local/lib/python3.7/dist-packages (from kapre) (0.8.0)\n",
            "Requirement already satisfied, skipping upgrade: numpy>=1.18.5 in /usr/local/lib/python3.7/dist-packages (from kapre) (1.19.5)\n",
            "Requirement already satisfied, skipping upgrade: requests<3,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (2.23.0)\n",
            "Collecting docker-pycreds>=0.4.0\n",
            "  Downloading https://files.pythonhosted.org/packages/f5/e8/f6bd1eee09314e7e6dee49cbe2c5e22314ccdb38db16c9fc72d2fa80d054/docker_pycreds-0.4.0-py2.py3-none-any.whl\n",
            "Requirement already satisfied, skipping upgrade: python-dateutil>=2.6.1 in /usr/local/lib/python3.7/dist-packages (from wandb) (2.8.1)\n",
            "Requirement already satisfied, skipping upgrade: psutil>=5.0.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (5.4.8)\n",
            "Requirement already satisfied, skipping upgrade: PyYAML in /usr/local/lib/python3.7/dist-packages (from wandb) (3.13)\n",
            "Requirement already satisfied, skipping upgrade: protobuf>=3.12.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (3.12.4)\n",
            "Requirement already satisfied, skipping upgrade: Click>=7.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (7.1.2)\n",
            "Collecting subprocess32>=3.5.3\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/32/c8/564be4d12629b912ea431f1a50eb8b3b9d00f1a0b1ceff17f266be190007/subprocess32-3.5.4.tar.gz (97kB)\n",
            "\u001b[K     |████████████████████████████████| 102kB 8.7MB/s \n",
            "\u001b[?25hCollecting sentry-sdk>=0.4.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/f3/92/5a33be64990ba815364a8f2dd9e6f51de60d23dfddafb4f1fc5577d4dc64/sentry_sdk-1.0.0-py2.py3-none-any.whl (131kB)\n",
            "\u001b[K     |████████████████████████████████| 133kB 14.1MB/s \n",
            "\u001b[?25hCollecting pathtools\n",
            "  Downloading https://files.pythonhosted.org/packages/e7/7f/470d6fcdf23f9f3518f6b0b76be9df16dcc8630ad409947f8be2eb0ed13a/pathtools-0.1.2.tar.gz\n",
            "Collecting GitPython>=1.0.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/a6/99/98019716955ba243657daedd1de8f3a88ca1f5b75057c38e959db22fb87b/GitPython-3.1.14-py3-none-any.whl (159kB)\n",
            "\u001b[K     |████████████████████████████████| 163kB 17.1MB/s \n",
            "\u001b[?25hRequirement already satisfied, skipping upgrade: promise<3,>=2.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (2.3)\n",
            "Collecting shortuuid>=0.5.0\n",
            "  Downloading https://files.pythonhosted.org/packages/25/a6/2ecc1daa6a304e7f1b216f0896b26156b78e7c38e1211e9b798b4716c53d/shortuuid-1.0.1-py3-none-any.whl\n",
            "Requirement already satisfied, skipping upgrade: six>=1.13.0 in /usr/local/lib/python3.7/dist-packages (from wandb) (1.15.0)\n",
            "Collecting configparser>=3.8.1\n",
            "  Downloading https://files.pythonhosted.org/packages/fd/01/ff260a18caaf4457eb028c96eeb405c4a230ca06c8ec9c1379f813caa52e/configparser-5.0.2-py3-none-any.whl\n",
            "Requirement already satisfied, skipping upgrade: keras-preprocessing~=1.1.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (1.1.2)\n",
            "Requirement already satisfied, skipping upgrade: gast==0.3.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (0.3.3)\n",
            "Requirement already satisfied, skipping upgrade: flatbuffers~=1.12.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (1.12)\n",
            "Requirement already satisfied, skipping upgrade: wheel~=0.35 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (0.36.2)\n",
            "Requirement already satisfied, skipping upgrade: astunparse~=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (1.6.3)\n",
            "Requirement already satisfied, skipping upgrade: tensorflow-estimator<2.5.0,>=2.4.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (2.4.0)\n",
            "Requirement already satisfied, skipping upgrade: wrapt~=1.12.1 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (1.12.1)\n",
            "Requirement already satisfied, skipping upgrade: termcolor~=1.1.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (1.1.0)\n",
            "Requirement already satisfied, skipping upgrade: google-pasta~=0.2 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (0.2.0)\n",
            "Requirement already satisfied, skipping upgrade: typing-extensions~=3.7.4 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (3.7.4.3)\n",
            "Requirement already satisfied, skipping upgrade: grpcio~=1.32.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (1.32.0)\n",
            "Requirement already satisfied, skipping upgrade: tensorboard~=2.4 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (2.4.1)\n",
            "Requirement already satisfied, skipping upgrade: h5py~=2.10.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (2.10.0)\n",
            "Requirement already satisfied, skipping upgrade: absl-py~=0.10 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (0.10.0)\n",
            "Requirement already satisfied, skipping upgrade: opt-einsum~=3.3.0 in /usr/local/lib/python3.7/dist-packages (from tensorflow>=2.0->kapre) (3.3.0)\n",
            "Requirement already satisfied, skipping upgrade: resampy>=0.2.2 in /usr/local/lib/python3.7/dist-packages (from librosa>=0.7.2->kapre) (0.2.2)\n",
            "Requirement already satisfied, skipping upgrade: pooch>=1.0 in /usr/local/lib/python3.7/dist-packages (from librosa>=0.7.2->kapre) (1.3.0)\n",
            "Requirement already satisfied, skipping upgrade: scipy>=1.0.0 in /usr/local/lib/python3.7/dist-packages (from librosa>=0.7.2->kapre) (1.4.1)\n",
            "Requirement already satisfied, skipping upgrade: audioread>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from librosa>=0.7.2->kapre) (2.1.9)\n",
            "Requirement already satisfied, skipping upgrade: joblib>=0.14 in /usr/local/lib/python3.7/dist-packages (from librosa>=0.7.2->kapre) (1.0.1)\n",
            "Requirement already satisfied, skipping upgrade: numba>=0.43.0 in /usr/local/lib/python3.7/dist-packages (from librosa>=0.7.2->kapre) (0.51.2)\n",
            "Requirement already satisfied, skipping upgrade: decorator>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from librosa>=0.7.2->kapre) (4.4.2)\n",
            "Requirement already satisfied, skipping upgrade: soundfile>=0.9.0 in /usr/local/lib/python3.7/dist-packages (from librosa>=0.7.2->kapre) (0.10.3.post1)\n",
            "Requirement already satisfied, skipping upgrade: scikit-learn!=0.19.0,>=0.14.0 in /usr/local/lib/python3.7/dist-packages (from librosa>=0.7.2->kapre) (0.22.2.post1)\n",
            "Requirement already satisfied, skipping upgrade: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.0.0->wandb) (1.24.3)\n",
            "Requirement already satisfied, skipping upgrade: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.0.0->wandb) (2.10)\n",
            "Requirement already satisfied, skipping upgrade: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.0.0->wandb) (2020.12.5)\n",
            "Requirement already satisfied, skipping upgrade: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3,>=2.0.0->wandb) (3.0.4)\n",
            "Requirement already satisfied, skipping upgrade: setuptools in /usr/local/lib/python3.7/dist-packages (from protobuf>=3.12.0->wandb) (54.0.0)\n",
            "Collecting gitdb<5,>=4.0.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/48/11/d1800bca0a3bae820b84b7d813ad1eff15a48a64caea9c823fc8c1b119e8/gitdb-4.0.5-py3-none-any.whl (63kB)\n",
            "\u001b[K     |████████████████████████████████| 71kB 8.6MB/s \n",
            "\u001b[?25hRequirement already satisfied, skipping upgrade: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.4->tensorflow>=2.0->kapre) (3.3.4)\n",
            "Requirement already satisfied, skipping upgrade: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.4->tensorflow>=2.0->kapre) (1.0.1)\n",
            "Requirement already satisfied, skipping upgrade: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.4->tensorflow>=2.0->kapre) (1.8.0)\n",
            "Requirement already satisfied, skipping upgrade: google-auth<2,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.4->tensorflow>=2.0->kapre) (1.27.1)\n",
            "Requirement already satisfied, skipping upgrade: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard~=2.4->tensorflow>=2.0->kapre) (0.4.3)\n",
            "Requirement already satisfied, skipping upgrade: appdirs in /usr/local/lib/python3.7/dist-packages (from pooch>=1.0->librosa>=0.7.2->kapre) (1.4.4)\n",
            "Requirement already satisfied, skipping upgrade: packaging in /usr/local/lib/python3.7/dist-packages (from pooch>=1.0->librosa>=0.7.2->kapre) (20.9)\n",
            "Requirement already satisfied, skipping upgrade: llvmlite<0.35,>=0.34.0.dev0 in /usr/local/lib/python3.7/dist-packages (from numba>=0.43.0->librosa>=0.7.2->kapre) (0.34.0)\n",
            "Requirement already satisfied, skipping upgrade: cffi>=1.0 in /usr/local/lib/python3.7/dist-packages (from soundfile>=0.9.0->librosa>=0.7.2->kapre) (1.14.5)\n",
            "Collecting smmap<4,>=3.0.1\n",
            "  Downloading https://files.pythonhosted.org/packages/d5/1e/6130925131f639b2acde0f7f18b73e33ce082ff2d90783c436b52040af5a/smmap-3.0.5-py2.py3-none-any.whl\n",
            "Requirement already satisfied, skipping upgrade: importlib-metadata; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard~=2.4->tensorflow>=2.0->kapre) (3.7.0)\n",
            "Requirement already satisfied, skipping upgrade: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.4->tensorflow>=2.0->kapre) (4.2.1)\n",
            "Requirement already satisfied, skipping upgrade: rsa<5,>=3.1.4; python_version >= \"3.6\" in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.4->tensorflow>=2.0->kapre) (4.7.2)\n",
            "Requirement already satisfied, skipping upgrade: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<2,>=1.6.3->tensorboard~=2.4->tensorflow>=2.0->kapre) (0.2.8)\n",
            "Requirement already satisfied, skipping upgrade: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard~=2.4->tensorflow>=2.0->kapre) (1.3.0)\n",
            "Requirement already satisfied, skipping upgrade: pyparsing>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->pooch>=1.0->librosa>=0.7.2->kapre) (2.4.7)\n",
            "Requirement already satisfied, skipping upgrade: pycparser in /usr/local/lib/python3.7/dist-packages (from cffi>=1.0->soundfile>=0.9.0->librosa>=0.7.2->kapre) (2.20)\n",
            "Requirement already satisfied, skipping upgrade: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata; python_version < \"3.8\"->markdown>=2.6.8->tensorboard~=2.4->tensorflow>=2.0->kapre) (3.4.1)\n",
            "Requirement already satisfied, skipping upgrade: pyasn1>=0.1.3 in /usr/local/lib/python3.7/dist-packages (from rsa<5,>=3.1.4; python_version >= \"3.6\"->google-auth<2,>=1.6.3->tensorboard~=2.4->tensorflow>=2.0->kapre) (0.4.8)\n",
            "Requirement already satisfied, skipping upgrade: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard~=2.4->tensorflow>=2.0->kapre) (3.1.0)\n",
            "Building wheels for collected packages: kapre\n",
            "  Building wheel for kapre (PEP 517) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for kapre: filename=kapre-0.3.4-cp37-none-any.whl size=20608 sha256=06c50835f8d330fd4214939c873a299bbc0e374fcc5c7f05c7700007eaa23b74\n",
            "  Stored in directory: /root/.cache/pip/wheels/f8/9e/dc/cc6f4989c58eee2913abac4fab60723ff91aa2a06c6de4280f\n",
            "Successfully built kapre\n",
            "Building wheels for collected packages: subprocess32, pathtools\n",
            "  Building wheel for subprocess32 (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for subprocess32: filename=subprocess32-3.5.4-cp37-none-any.whl size=6489 sha256=b0da439024e5981d5b1bad0881d8d9718a076658c5b4e4bade284b16b84d72cb\n",
            "  Stored in directory: /root/.cache/pip/wheels/68/39/1a/5e402bdfdf004af1786c8b853fd92f8c4a04f22aad179654d1\n",
            "  Building wheel for pathtools (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pathtools: filename=pathtools-0.1.2-cp37-none-any.whl size=8786 sha256=f36db17f80023c010b597ed09cb0553954167e4b032b084c8bfce8d2183c8a53\n",
            "  Stored in directory: /root/.cache/pip/wheels/0b/04/79/c3b0c3a0266a3cb4376da31e5bfe8bba0c489246968a68e843\n",
            "Successfully built subprocess32 pathtools\n",
            "Installing collected packages: kapre, docker-pycreds, subprocess32, sentry-sdk, pathtools, smmap, gitdb, GitPython, shortuuid, configparser, wandb\n",
            "  Found existing installation: kapre 0.1.3.1\n",
            "    Uninstalling kapre-0.1.3.1:\n",
            "      Successfully uninstalled kapre-0.1.3.1\n",
            "Successfully installed GitPython-3.1.14 configparser-5.0.2 docker-pycreds-0.4.0 gitdb-4.0.5 kapre-0.3.4 pathtools-0.1.2 sentry-sdk-1.0.0 shortuuid-1.0.1 smmap-3.0.5 subprocess32-3.5.4 wandb-0.10.22\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9DMg9412dbot"
      },
      "source": [
        "import random, string, os, tqdm, pickle, glob, logging, sys, shutil, datetime, gc, logging, io, argparse, itertools\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import sklearn.metrics\n",
        "import numpy as np\n",
        "import wandb"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6wy1Qz4Mf-4c"
      },
      "source": [
        "Imports for tensorflow and the network's layers, using audio-specific implementations from the kapre package"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vS9vyptwf-DB"
      },
      "source": [
        "#------for the network-----------#\n",
        "import tensorflow as tf\n",
        "from tensorflow.keras import layers\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras.callbacks import Callback\n",
        "from wandb.keras import WandbCallback\n",
        "\n",
        "#------for kapre-----------------#\n",
        "import kapre\n",
        "from kapre import STFT, Magnitude, MagnitudeToDecibel, Delta, Frame\n",
        "from kapre.composed import get_melspectrogram_layer, get_log_frequency_spectrogram_layer, get_stft_magnitude_layer"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TjtUhmjnhtWd"
      },
      "source": [
        "# Helper functions\n",
        "Let's implement some helper functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vCVQ4VQCtLI-"
      },
      "source": [
        "def set_all_seeds(seed):\n",
        "  random.seed(seed)\n",
        "  os.environ['PYTHONHASHSEED'] = str(seed)\n",
        "  np.random.seed(seed)\n",
        "  tf.random.set_seed(seed)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZSMgchvYiPZR"
      },
      "source": [
        "def get_logger(filename, filemode):\n",
        "  logging.basicConfig(level=logging.DEBUG,\n",
        "                    format='%(asctime)s %(name)-12s %(levelname)-8s %(message)s',\n",
        "                    datefmt='%m-%d %H:%M',\n",
        "                    filename=filename,\n",
        "                    filemode=filemode)\n",
        "  # define a Handler which writes INFO messages or higher to the sys.stderr\n",
        "  console = logging.StreamHandler()\n",
        "  console.setLevel(logging.INFO)\n",
        "  # set a format which is simpler for console use\n",
        "  formatter = logging.Formatter('%(message)s')\n",
        "  # tell the handler to use this format\n",
        "  console.setFormatter(formatter)\n",
        "  # add the handler to the root logger\n",
        "  logging.getLogger('').addHandler(console)\n",
        "\n",
        "  logger1 = logging.getLogger(\"Standard\")\n",
        "  \n",
        "  return logger1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "76vjyoU4hq5K"
      },
      "source": [
        "def get_file_lists(tfr_dir):\n",
        "\n",
        "  #list all files with training data\n",
        "  train_shards = glob.glob(tfr_dir+\"*_train.tfrecords\", recursive=False)\n",
        "  sorted(train_shards) #for same same input order over all runs \n",
        "\n",
        "  #list all files with test data\n",
        "  test_shards = glob.glob(tfr_dir+\"*_test.tfrecords\", recursive=False)\n",
        "  sorted(test_shards) #for same same input order over all runs \n",
        "\n",
        "  #list all files with validation data\n",
        "  valid_shards = glob.glob(tfr_dir+\"*_valid.tfrecords\", recursive=False)\n",
        "  sorted(valid_shards) #for same same input order over all runs \n",
        "\n",
        "  logger1.info(f\"Number of training, test, and validation shards: {len(train_shards), len(test_shards), len(valid_shards)}\")\n",
        "\n",
        "  return train_shards, test_shards, valid_shards"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LdzO7pM9iA6J"
      },
      "source": [
        "def parse_tfr_elem(element):\n",
        "  '''\n",
        "   extract the features from a TFrecord element\n",
        "  '''\n",
        "  \n",
        "  parse_dict = {\n",
        "      'x': tf.io.FixedLenFeature([], tf.int64),\n",
        "      'y':tf.io.FixedLenFeature([], tf.int64),\n",
        "      'label':tf.io.FixedLenFeature([], tf.int64),\n",
        "      'feature' : tf.io.FixedLenFeature([], tf.string)}\n",
        "\n",
        "  example_message = tf.io.parse_single_example(element, parse_dict)\n",
        "\n",
        "  x = example_message['x']\n",
        "  y = example_message['y']\n",
        "  feature = example_message['feature']\n",
        "  label = example_message['label']\n",
        "\n",
        "  feature = tf.io.parse_tensor(feature, out_type=tf.float32)\n",
        "  feature = tf.reshape(feature, shape=[x,y])\n",
        "\n",
        "  return (feature, label)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uNzabBcNieG_"
      },
      "source": [
        "def load_dataset(tfrecords, cache_path=None, batch_size=8, seed=1337):\n",
        "  \n",
        "  AUTOTUNE = tf.data.experimental.AUTOTUNE #automatic optimizer for caching, prefetching (later on)\n",
        "  \n",
        "  options = tf.data.Options()\n",
        "  options.experimental_deterministic = False #speed up file processing by not waiting for files to be in order\n",
        "\n",
        "  dataset = tf.data.TFRecordDataset(tfrecords, buffer_size=100000000)\n",
        "  dataset.with_options(options)\n",
        "\n",
        "  dataset = dataset.map(map_func=parse_tfr_elem, num_parallel_calls=AUTOTUNE)\n",
        "  dataset = dataset.prefetch(AUTOTUNE)\n",
        "  if cache_path is not None:\n",
        "    dataset = dataset.cache(cache_path)\n",
        "  else:\n",
        "    dataset = dataset.cache() #hold cache in RAM\n",
        "\n",
        "  dataset = dataset.shuffle(10, seed=seed, reshuffle_each_iteration=True)\n",
        "  dataset = dataset.batch(batch_size=batch_size)\n",
        "  dataset = dataset.repeat() #always repeat to run multiple epochs and multiple test/val runs\n",
        "\n",
        "  return dataset"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JOqRI5MRo3Z6"
      },
      "source": [
        "def get_strategy(xla=0, fp16=0, no_cuda=0):\n",
        "  '''\n",
        "  Determines the strategy under which the network is trained.\n",
        "  \n",
        "  From https://github.com/huggingface/transformers/blob/8eb7f26d5d9ce42eb88be6f0150b22a41d76a93d/src/transformers/training_args_tf.py\n",
        "  \n",
        "  returns the strategy object\n",
        "  \n",
        "  '''\n",
        "  logger1.info(\"TensorFlow: setting up strategy\")\n",
        "\n",
        "  if xla:\n",
        "    tf.config.optimizer.set_jit(True)\n",
        "\n",
        "  gpus = tf.config.list_physical_devices(\"GPU\")\n",
        "    # Set to float16 at first\n",
        "  if fp16:\n",
        "    policy = tf.keras.mixed_precision.experimental.Policy(\"mixed_float16\")\n",
        "    tf.keras.mixed_precision.experimental.set_policy(policy)\n",
        "\n",
        "  if no_cuda:\n",
        "    strategy = tf.distribute.OneDeviceStrategy(device=\"/cpu:0\")\n",
        "  else:\n",
        "    try:\n",
        "      tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n",
        "    except ValueError:\n",
        "      tpu = None\n",
        "  \n",
        "    if tpu:\n",
        "    # Set to bfloat16 in case of TPU\n",
        "      if fp16:\n",
        "        policy = tf.keras.mixed_precision.experimental.Policy(\"mixed_bfloat16\")\n",
        "        tf.keras.mixed_precision.experimental.set_policy(policy)\n",
        "      tf.config.experimental_connect_to_cluster(tpu)\n",
        "      tf.tpu.experimental.initialize_tpu_system(tpu)\n",
        "    \n",
        "      strategy = tf.distribute.experimental.TPUStrategy(tpu)\n",
        "    \n",
        "    elif len(gpus) == 0:\n",
        "        strategy = tf.distribute.OneDeviceStrategy(device=\"/cpu:0\")\n",
        "    elif len(gpus) == 1:\n",
        "      strategy = tf.distribute.OneDeviceStrategy(device=\"/gpu:0\")\n",
        "    elif len(gpus) > 1:\n",
        "      # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`\n",
        "      strategy = tf.distribute.MirroredStrategy()\n",
        "    else:\n",
        "      raise ValueError(\"Cannot find the proper strategy! Please check your environment properties.\")\n",
        "\n",
        "  logger1.info(f\"Using strategy: {strategy}\")\n",
        "  return strategy\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WoEStXk7jGgL"
      },
      "source": [
        "# Callbacks"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tvBwfkLlkH19"
      },
      "source": [
        "class ClassificationReportCallback(Callback):\n",
        "    __slots__=[\"x\", \"y\", \"file_writer\", \"class_names\", \"name\",\"batch_size\", \"labels\", \"frequency\"]\n",
        "    def __init__(self, x, y, file_writer, class_names, name, batch_size, labels, frequency=5, use_wandb_tracking:bool=False):\n",
        "      super(ClassificationReportCallback, self).__init__()\n",
        "      self.x = x\n",
        "      self.y = y\n",
        "      self.class_names = class_names\n",
        "      self.file_writer = file_writer\n",
        "      self.batch_size = batch_size\n",
        "      self.labels = labels\n",
        "      self.name = name\n",
        "      self.frequency = frequency #create report every k epochs\n",
        "      self.wandb = use_wandb_tracking\n",
        "\n",
        "    def on_epoch_end(self, epoch, logs=None):\n",
        "      if (epoch % self.frequency)==0:\n",
        "        self.generate_classification_report(epoch)\n",
        "    \n",
        "    def generate_classification_report(self, epoch):\n",
        "      test_pred_raw = self.model.predict(self.x, batch_size=self.batch_size)\n",
        "      test_pred = np.argmax(test_pred_raw, axis=1)\n",
        "\n",
        "      classification_report_text = sklearn.metrics.classification_report(self.y, test_pred, labels=self.labels, target_names=self.class_names, output_dict=False)\n",
        "      classification_report_text = \"\\n\".join(classification_report_text.splitlines()[:7])\n",
        "\n",
        "      with self.file_writer.as_default():\n",
        "        tf.summary.text(self.name, classification_report_text, step=epoch)\n",
        "\n",
        "      if self.wandb:\n",
        "        classification_report_dict = sklearn.metrics.classification_report(self.y, test_pred, labels=self.labels, target_names=self.class_names, output_dict=True)\n",
        "        wandb.log(data=classification_report_dict)\n",
        "      \n",
        "      gc.collect()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aBXmPlm5kpmt"
      },
      "source": [
        "#after https://www.tensorflow.org/tensorboard/image_summaries\n",
        "\n",
        "class ConfusionMatrixCallback(Callback):\n",
        "    __slots__=['x', 'y', 'class_names', 'image_name', 'file_writer', 'batch_size', 'frequency']\n",
        "    def __init__(self, x, y, file_writer, class_names, image_name, batch_size, frequency=5):\n",
        "        super(ConfusionMatrixCallback, self).__init__()\n",
        "        self.x = x\n",
        "        self.y = y\n",
        "        self.class_names = class_names\n",
        "        self.image_name = image_name\n",
        "        self.batch_size=batch_size\n",
        "        self.file_writer = file_writer\n",
        "        self.frequency = frequency\n",
        "\n",
        "\n",
        "    def on_epoch_end(self, epoch, logs=None):\n",
        "        if (epoch%self.frequency) == 0:\n",
        "          self.log_confusion_matrix(epoch)\n",
        "\n",
        "    \n",
        "    def log_confusion_matrix(self, epoch):\n",
        "        # Use the model to predict the values from the validation dataset.\n",
        "        test_pred_raw = self.model.predict(self.x, batch_size=self.batch_size)\n",
        "        test_pred = np.argmax(test_pred_raw, axis=-1)\n",
        "\n",
        "        # Calculate the confusion matrix.\n",
        "        cm = sklearn.metrics.confusion_matrix(self.y, test_pred)\n",
        "        # Log the confusion matrix as an image summary.\n",
        "        figure = self.plot_confusion_matrix(cm, class_names=self.class_names)\n",
        "        cm_image = self.plot_to_image(figure)\n",
        "\n",
        "        # Log the confusion matrix as an image summary.\n",
        "        with self.file_writer.as_default():\n",
        "          tf.summary.image(self.image_name, cm_image, step=epoch)\n",
        "        \n",
        "        gc.collect()\n",
        "\n",
        "\n",
        "    def plot_confusion_matrix(self, cm, class_names):\n",
        "        \"\"\"\n",
        "        Returns a matplotlib figure containing the plotted confusion matrix.\n",
        "\n",
        "        Args:\n",
        "        cm (array, shape = [n, n]): a confusion matrix of integer classes\n",
        "        class_names (array, shape = [n]): String names of the integer classes\n",
        "        \"\"\"\n",
        "        figure = plt.figure(figsize=(8, 8))\n",
        "        plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)\n",
        "        plt.title(\"Confusion matrix\")\n",
        "        plt.colorbar()\n",
        "        tick_marks = np.arange(len(class_names))\n",
        "        plt.xticks(tick_marks, class_names, rotation=45)\n",
        "        plt.yticks(tick_marks, class_names)\n",
        "\n",
        "        # Compute the labels from the normalized confusion matrix.\n",
        "        labels = np.around(cm.astype('float') / cm.sum(axis=1)[:, np.newaxis], decimals=2)\n",
        "\n",
        "        # Use white text if squares are dark; otherwise black.\n",
        "        threshold = cm.max() / 2.\n",
        "        for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):\n",
        "          color = \"white\" if cm[i, j] > threshold else \"black\"\n",
        "          plt.text(j, i, labels[i, j], horizontalalignment=\"center\", color=color)\n",
        "\n",
        "        plt.tight_layout()\n",
        "        plt.ylabel('True label')\n",
        "        plt.xlabel('Predicted label')\n",
        "\n",
        "        return figure\n",
        "\n",
        "    def plot_to_image(self, figure):\n",
        "        \"\"\"Converts the matplotlib plot specified by 'figure' to a PNG image and\n",
        "        returns it. The supplied figure is closed and inaccessible after this call.\"\"\"\n",
        "        # Save the plot to a PNG in memory.\n",
        "        buf = io.BytesIO()\n",
        "        plt.savefig(buf, format='png')\n",
        "        # Closing the figure prevents it from being displayed directly inside\n",
        "        # the notebook.\n",
        "        plt.close(figure)\n",
        "        buf.seek(0)\n",
        "        # Convert PNG buffer to TF image\n",
        "        image = tf.image.decode_png(buf.getvalue(), channels=4)\n",
        "        # Add the batch dimension\n",
        "        image = tf.expand_dims(image, 0)\n",
        "        return image\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p9kojV15jIO3"
      },
      "source": [
        "def get_callbacks(validation_dataset, steps_valid, train_monitor_x, train_monitor_y, valid_monitor_x, valid_monitor_y, args):\n",
        "\n",
        "  logger1.info(\"\\nSetting up callbacks.\")\n",
        "  \n",
        "\n",
        "  tensorboard_dir = args['out_dir']+\"tensorboard/\"\n",
        "  os.makedirs(tensorboard_dir, exist_ok=True)\n",
        "\n",
        "  file_writer = tf.summary.create_file_writer(tensorboard_dir)\n",
        "\n",
        "  #tensorboard callback\n",
        "  logger1.info(\"\\nLogging TensorBoard to {}.\".format(tensorboard_dir))\n",
        "  \n",
        "  #tensorboard callback\n",
        "  tb_callback = keras.callbacks.TensorBoard(log_dir=tensorboard_dir, histogram_freq=1, write_graph=True, write_images=True)\n",
        "\n",
        "  #confusion matrix callback (applied after each epoch)\n",
        "  cm_callback_train = ConfusionMatrixCallback(train_monitor_x, train_monitor_y, file_writer, [\"male\", \"female\"], 'Confusion Matrix Train', args['batch_size'], 3)\n",
        "  cm_callback_valid = ConfusionMatrixCallback(valid_monitor_x, valid_monitor_y, file_writer, [\"male\", \"female\"], 'Confusion Matrix Valid', args['batch_size'], 3)\n",
        "\n",
        "  #Classification report callback (applied after each epoch)\n",
        "  classification_report_cb = ClassificationReportCallback(valid_monitor_x, valid_monitor_y, file_writer, [\"male\", \"female\"], \"Classification Report Validation\", args['batch_size'], [0,1], 3)\n",
        "\n",
        "  #early stopping\n",
        "  es_callback = keras.callbacks.EarlyStopping(monitor='val_sparse_categorical_accuracy', patience=args['es_patience'], verbose=1, restore_best_weights=True)\n",
        "\n",
        "\n",
        "  callbacks = [tb_callback, es_callback, cm_callback_train, cm_callback_valid, classification_report_cb]\n",
        "  \n",
        "  if args['use_wandb_tracking']:\n",
        "    wandb_callback = WandbCallback(labels=[\"male\", \"female\"], log_weights=True, monitor=\"val_sparse_categorical_accuracy\", generator=validation_dataset, validation_steps=steps_valid, log_evaluation=True)\n",
        "    callbacks.append(wandb_callback)\n",
        "\n",
        "  logger1.info(\"\\nCallbacks are all set up.\")\n",
        "\n",
        "  return callbacks"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XBAdyjYKiuGI"
      },
      "source": [
        "# The model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SOT1zWEEitce"
      },
      "source": [
        "def get_model(config):\n",
        "  input_ = layers.Input(shape=(1323000, 1))\n",
        "\n",
        "  mel_layer = get_melspectrogram_layer(n_mels=config['n_mels'], sample_rate=22050, win_length=config['win_length'], n_fft=config['n_fft'], return_decibel=True, output_data_format='channels_last')(input_)\n",
        "\n",
        "  delta_layer = Delta()(mel_layer)\n",
        "  concat = layers.Concatenate()([mel_layer, delta_layer])\n",
        "  frame_layer = Frame(frame_length=config['frame_length'], hop_length=config['hop_length'])(concat)\n",
        "\n",
        "  x = layers.Conv2D(config['num_conv1'], config['kernel_size'], padding=\"same\")(frame_layer)\n",
        "  x = layers.BatchNormalization()(x)\n",
        "  x = layers.Activation(\"relu\")(x)\n",
        "  x = layers.Dropout(config['dropout1'])(x)\n",
        "  x = layers.MaxPooling3D(pool_size=config['pool_size'])(x)\n",
        "\n",
        "  x = layers.Conv2D(config['num_conv2'], config['kernel_size'], padding=\"same\")(x)\n",
        "  x = layers.BatchNormalization()(x)\n",
        "  x = layers.Activation(\"relu\")(x)\n",
        "  x = layers.Dropout(config['dropout2'])(x)\n",
        "  x = layers.MaxPooling3D(pool_size=config['pool_size'])(x)\n",
        "\n",
        "  x = layers.Conv2D(config['num_conv3'], config['kernel_size'], padding=\"same\")(x)\n",
        "  x = layers.BatchNormalization()(x)\n",
        "  x = layers.Activation(\"relu\")(x)\n",
        "  x = layers.Dropout(config['dropout3'])(x)\n",
        "  x = layers.MaxPooling3D(pool_size=config['pool_size'])(x)\n",
        "\n",
        "  x = layers.Conv2D(config['num_conv4'], config['kernel_size'], padding=\"same\")(x)\n",
        "  x = layers.BatchNormalization()(x)\n",
        "  x = layers.Activation(\"relu\")(x)\n",
        "  x = layers.Dropout(config['dropout4'])(x)\n",
        "  x = layers.GlobalMaxPooling3D()(x)\n",
        "\n",
        "  x = layers.Dropout(config['dropout5'])(x)\n",
        "  x = layers.Dense(config['num_dense'], activation=\"relu\", name=\"pre_out_dense\")(x)\n",
        "  x = layers.Dense(2, activation=\"softmax\")(x)\n",
        "\n",
        "  return keras.Model(inputs=[input_], outputs=[x], name=\"KapreClassifier\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I-cFo4o0mANi"
      },
      "source": [
        "def getHyperparameterDefaults(args):\n",
        "  hyperparameter_defaults = {\n",
        "     'num_conv1':24, #conv should be multiple of 8\n",
        "     'num_conv2':32,\n",
        "     'num_conv3':64,\n",
        "     'num_conv4':128,\n",
        "     'dropout1':0.1,\n",
        "     'dropout2':0.1,\n",
        "     'dropout3':0.1,\n",
        "     'dropout4':0.1,\n",
        "     'pool_size': (2,2,2),\n",
        "     'kernel_size': (3,3),\n",
        "     'win_length': 512,\n",
        "     'n_fft':2048,\n",
        "     'n_mels':60,\n",
        "     'frame_length':41,\n",
        "     'hop_length':64,\n",
        "     'dropout5':0.5,\n",
        "     'num_dense':512, #dense should be multiple of 8\n",
        "     'epochs':args['epochs']} #writing args here to log it to W&B (just in case)\n",
        "  \n",
        "  return hyperparameter_defaults"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "npaw-LGzxCFE"
      },
      "source": [
        "# Setup for training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KaZ4gb-LpenJ"
      },
      "source": [
        "def setup_model(strategy, model_config, args):\n",
        "  #setting the training strategy\n",
        "  with strategy.scope():\n",
        "    model = get_model(model_config)\n",
        "    model.compile(optimizer=keras.optimizers.Adam(1e-4),loss=keras.losses.SparseCategoricalCrossentropy(),metrics=[\"sparse_categorical_accuracy\"])\n",
        "\n",
        "  #saving a model's visualization\n",
        "  tf.keras.utils.plot_model(\n",
        "    model,\n",
        "    to_file=args[\"out_dir\"]+\"model.png\",\n",
        "    show_shapes=True,\n",
        "    show_dtype=False,\n",
        "    show_layer_names=True,\n",
        "    rankdir=\"TB\",\n",
        "    expand_nested=False,\n",
        "    dpi=150,\n",
        "  )\n",
        "\n",
        "  return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PdR6xA0ZrjqP"
      },
      "source": [
        "def save_model(model, args):\n",
        "  \n",
        "  model_directory = os.path.join(args['out_dir'], args['model_dir'])\n",
        "  \n",
        "  os.makedirs(model_directory, exist_ok=True)\n",
        "\n",
        "  model.save(model_directory)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NZETL4IOr-Zm"
      },
      "source": [
        "def evaluate_model(model):\n",
        "  \n",
        "  logger1.info(\"\\nStarting evaluation on test data.\")\n",
        "  \n",
        "  test_loss, test_acc = model.evaluate(x=test_dataset, steps=steps_test)\n",
        "  \n",
        "  logger1.info(\"\\nFinished evaluation. Loss: {:.4f}, Accuracy: {:.4f}.\".format(test_loss, test_acc))\n",
        "\n",
        "  return test_loss, test_acc"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r_L19SGPnHJO"
      },
      "source": [
        "def train(model, callbacks, args):\n",
        "\n",
        "  logger1.info(\"Starting training.\")\n",
        "  \n",
        "  model.fit(x=train_dataset, validation_data=validation_dataset, validation_steps=steps_valid, steps_per_epoch=steps_train, epochs=args['epochs'], callbacks=callbacks)\n",
        "\n",
        "  logger1.info(\"Training finished.\")\n",
        "\n",
        "  return model  \n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s7Goq6BEpO0j"
      },
      "source": [
        "def setup(args):\n",
        "\n",
        "  #seeding for reproducability\n",
        "  set_all_seeds(args['seed'])\n",
        "\n",
        "  #get model configuration\n",
        "  model_config = getHyperparameterDefaults(args)\n",
        "\n",
        "  #initialize W&B logging if requested\n",
        "  if args['use_wandb_tracking']:\n",
        "    wandb.tensorboard.patch(root_logdir=args[\"out_dir\"])\n",
        "    wandb.init(config=model_config, entity=args[\"wandb_entity\"], project=args[\"wandb_project\"], group=str(args[\"wandb_project\"]), sync_tensorboard=True)\n",
        "  \n",
        "  #determine the distribution strategy\n",
        "  distribution_strategy = get_strategy(xla=args['xla'], fp16=args['fp16'], no_cuda=args['no_cuda'])\n",
        "\n",
        "  #setup model\n",
        "  model = setup_model(strategy=distribution_strategy, model_config=model_config, args=args)\n",
        "\n",
        "  #initializing the callbacks\n",
        "  callbacks = get_callbacks(validation_dataset, steps_valid, train_monitor_x, train_monitor_y, valid_monitor_x, valid_monitor_y, args)\n",
        "  \n",
        "  return model, callbacks"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pedi7fF3rW2c"
      },
      "source": [
        "def main(args):\n",
        "  \n",
        "  #initial setup\n",
        "  model, callbacks = setup(args=args)\n",
        "\n",
        "  #summary\n",
        "  model.summary(positions=[.33, .60, .67, 1.])\n",
        "\n",
        "  #training\n",
        "  model = train(model=model, callbacks=callbacks, args=args)\n",
        "\n",
        "  #saving the model\n",
        "  save_model(model, args)  \n",
        "\n",
        "  #evaluate model:\n",
        "  test_loss, test_acc = evaluate_model(model)\n",
        "\n",
        "  if args['use_wandb_tracking']:\n",
        "    wandb.log({\"Test loss\":test_loss, \"Test accuracy\": test_acc})\n",
        "    wandb.finish()\n",
        "  \n",
        "  tf.keras.backend.clear_session() #clean up, free memory (not reliable though)\n",
        "  del model\n",
        "  gc.collect()\n",
        "\n",
        "  logger1.info(\"\\nScript finished.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZDF3Nq2E58Mo"
      },
      "source": [
        "# Commandline arguments"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PrzjNc3YtQHR"
      },
      "source": [
        "parser = argparse.ArgumentParser(description='')\n",
        "\n",
        "parser.add_argument('--out_dir', dest='out_dir', default=\"/content/custaudio/\", help=\"Directory where all the stuff is stored\")\n",
        "parser.add_argument('--tensorboard_dir', dest='tensorboard_dir', default=\"tensorboard\", help=\"Subdir to store the logs at\")\n",
        "parser.add_argument('--tfr_dir', dest='tfr_dir', default=\"/content/drive/MyDrive/custaudio/tfr_dir/\", help=\"Where the TFRecord files are stored\")\n",
        "parser.add_argument('--logfile_name', dest='logfile_name', default=\"logfile.log\", help=\"Where the logfile is stored\")\n",
        "parser.add_argument('--model_dir', dest='model_dir', default=\"model\", help=\"Name of the directory to store the model\")\n",
        "\n",
        "parser.add_argument('--epochs', dest='epochs', type=int, default=1000, help=\"Number of training epochs\")\n",
        "parser.add_argument('--seed', dest='seed', type=int, default=1337, help=\"Seed for initializing random number generators\")\n",
        "parser.add_argument('--per_device_batch_size', dest='per_device_batch_size', type=int, default=8, help=\"Batch size per device\")\n",
        "parser.add_argument('--xla', dest='xla', type=int, default=1, help=\"Use xla compiled cpu operations\")\n",
        "parser.add_argument('--fp16', dest='fp16', type=int, default=0, help=\"Use mixed precision training\")\n",
        "parser.add_argument('--no_cuda', dest='no_cuda', type=int, default=0, help=\"Run training only on CPU\")\n",
        "parser.add_argument('--es_patience', dest='es_patience', type=int, default=25, help=\"Number of epochs without change until ES terminates training and restores the best weights so far\")\n",
        "\n",
        "parser.add_argument('--wandb', dest='use_wandb_tracking', type=int, default=0, help=\"Use wandb logging\")\n",
        "parser.add_argument('--entity', dest='wandb_entity',  help=\"Entity to log the wandb stuff to\")\n",
        "parser.add_argument('--group', dest='wandb_group', type=str, help=\"Group for the wandb run\")\n",
        "parser.add_argument('--project', dest='wandb_project', default=\"speaker_classification\", help=\"Project to log the files to when using wandb\")\n",
        "\n",
        "args, unknown = parser.parse_known_args()\n",
        "args = args.__dict__"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z4GjOtwz6Aba"
      },
      "source": [
        "# Driver code"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N1nwcoc5qFAz",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "73892405-5623-4f97-9a25-fb7102e975b5"
      },
      "source": [
        "if __name__ == \"__main__\":\n",
        "  os.makedirs(args['out_dir'], exist_ok=True)\n",
        "  logger1 = get_logger(filename=args['out_dir']+args['logfile_name'], filemode=\"a\")\n",
        "\n",
        "  #defining the following variables here makes them available globally\n",
        "  num_devices = len(tf.config.list_physical_devices('GPU')) \n",
        "  multiplier = num_devices if num_devices != 0 else 1\n",
        "  batch_size = args['per_device_batch_size']*multiplier\n",
        "  args['batch_size'] = batch_size\n",
        "  logger1.info(f\"Batch size is {batch_size}\")\n",
        "  \n",
        "  #first load the shards, then create a dataset object from them\n",
        "  train_shards, test_shards, valid_shards = get_file_lists(tfr_dir=args['tfr_dir'])\n",
        "  train_dataset = load_dataset(tfrecords=train_shards, cache_path=None, batch_size=batch_size, seed=args['seed'])\n",
        "  validation_dataset = load_dataset(tfrecords=valid_shards, cache_path=None, batch_size=batch_size, seed=args['seed'])\n",
        "  test_dataset = load_dataset(tfrecords=test_shards, cache_path=None, batch_size=batch_size, seed=args['seed'])\n",
        "  \n",
        "  steps_train = 50//batch_size  \n",
        "  steps_test = 50//batch_size \n",
        "  steps_valid = 50//batch_size\n",
        "\n",
        "  logger1.info(\"Loading some batches to visualize learning\")\n",
        "  train_monitor_x = np.load(args['tfr_dir']+\"train_x_monitor.npy\") #for gradient visualization, cm plotting\n",
        "  train_monitor_y = np.load(args['tfr_dir']+\"train_y_monitor.npy\")\n",
        "  valid_monitor_x = np.load(args['tfr_dir']+\"valid_x_monitor.npy\") #for gradient visualization, cm plotting\n",
        "  valid_monitor_y = np.load(args['tfr_dir']+\"valid_y_monitor.npy\")\n",
        "\n",
        "  main(args)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Batch size is 8\n",
            "Number of training, test, and validation shards: (1, 1, 1)\n",
            "Loading some batches to visualize learning\n",
            "TensorFlow: setting up strategy\n",
            "Using strategy: <tensorflow.python.distribute.one_device_strategy.OneDeviceStrategy object at 0x7ff8105fba50>\n",
            "\n",
            "Setting up callbacks.\n",
            "\n",
            "Logging TensorBoard to /content/custaudio/tensorboard/.\n",
            "\n",
            "Callbacks are all set up.\n",
            "Starting training.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Model: \"KapreClassifier\"\n",
            "__________________________________________________________________________________________________\n",
            "Layer (type)                    Output Shape              Param  Connected to                     \n",
            "==================================================================================================\n",
            "input_1 (InputLayer)            [(None, 1323000, 1)]      0                                       \n",
            "__________________________________________________________________________________________________\n",
            "melspectrogram (Sequential)     (None, 10332, 60, 1)      0      input_1[0][0]                    \n",
            "__________________________________________________________________________________________________\n",
            "delta (Delta)                   (None, 10332, 60, 1)      0      melspectrogram[0][0]             \n",
            "__________________________________________________________________________________________________\n",
            "concatenate (Concatenate)       (None, 10332, 60, 2)      0      melspectrogram[0][0]             \n",
            "                                                                 delta[0][0]                      \n",
            "__________________________________________________________________________________________________\n",
            "frame (Frame)                   (None, 161, 41, 60, 2)    0      concatenate[0][0]                \n",
            "__________________________________________________________________________________________________\n",
            "conv2d (Conv2D)                 (None, 161, 41, 60, 24)   456    frame[0][0]                      \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization (BatchNorma (None, 161, 41, 60, 24)   96     conv2d[0][0]                     \n",
            "__________________________________________________________________________________________________\n",
            "activation (Activation)         (None, 161, 41, 60, 24)   0      batch_normalization[0][0]        \n",
            "__________________________________________________________________________________________________\n",
            "dropout (Dropout)               (None, 161, 41, 60, 24)   0      activation[0][0]                 \n",
            "__________________________________________________________________________________________________\n",
            "max_pooling3d (MaxPooling3D)    (None, 80, 20, 30, 24)    0      dropout[0][0]                    \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_1 (Conv2D)               (None, 80, 20, 30, 32)    6944   max_pooling3d[0][0]              \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_1 (BatchNor (None, 80, 20, 30, 32)    128    conv2d_1[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "activation_1 (Activation)       (None, 80, 20, 30, 32)    0      batch_normalization_1[0][0]      \n",
            "__________________________________________________________________________________________________\n",
            "dropout_1 (Dropout)             (None, 80, 20, 30, 32)    0      activation_1[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "max_pooling3d_1 (MaxPooling3D)  (None, 40, 10, 15, 32)    0      dropout_1[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_2 (Conv2D)               (None, 40, 10, 15, 64)    18496  max_pooling3d_1[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_2 (BatchNor (None, 40, 10, 15, 64)    256    conv2d_2[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "activation_2 (Activation)       (None, 40, 10, 15, 64)    0      batch_normalization_2[0][0]      \n",
            "__________________________________________________________________________________________________\n",
            "dropout_2 (Dropout)             (None, 40, 10, 15, 64)    0      activation_2[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "max_pooling3d_2 (MaxPooling3D)  (None, 20, 5, 7, 64)      0      dropout_2[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "conv2d_3 (Conv2D)               (None, 20, 5, 7, 128)     73856  max_pooling3d_2[0][0]            \n",
            "__________________________________________________________________________________________________\n",
            "batch_normalization_3 (BatchNor (None, 20, 5, 7, 128)     512    conv2d_3[0][0]                   \n",
            "__________________________________________________________________________________________________\n",
            "activation_3 (Activation)       (None, 20, 5, 7, 128)     0      batch_normalization_3[0][0]      \n",
            "__________________________________________________________________________________________________\n",
            "dropout_3 (Dropout)             (None, 20, 5, 7, 128)     0      activation_3[0][0]               \n",
            "__________________________________________________________________________________________________\n",
            "global_max_pooling3d (GlobalMax (None, 128)               0      dropout_3[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "dropout_4 (Dropout)             (None, 128)               0      global_max_pooling3d[0][0]       \n",
            "__________________________________________________________________________________________________\n",
            "pre_out_dense (Dense)           (None, 512)               66048  dropout_4[0][0]                  \n",
            "__________________________________________________________________________________________________\n",
            "dense (Dense)                   (None, 2)                 1026   pre_out_dense[0][0]              \n",
            "==================================================================================================\n",
            "Total params: 167,818\n",
            "Trainable params: 167,322\n",
            "Non-trainable params: 496\n",
            "__________________________________________________________________________________________________\n",
            "Epoch 1/1000\n",
            "6/6 [==============================] - 45s 1s/step - loss: 1.5700 - sparse_categorical_accuracy: 0.5283 - val_loss: 0.6999 - val_sparse_categorical_accuracy: 0.4792\n",
            "Epoch 2/1000\n",
            "6/6 [==============================] - 1s 147ms/step - loss: 1.9440 - sparse_categorical_accuracy: 0.5274 - val_loss: 0.9161 - val_sparse_categorical_accuracy: 0.4167\n",
            "Epoch 3/1000\n",
            "6/6 [==============================] - 4s 688ms/step - loss: 1.2711 - sparse_categorical_accuracy: 0.6675 - val_loss: 0.7743 - val_sparse_categorical_accuracy: 0.4167\n",
            "Epoch 4/1000\n",
            "6/6 [==============================] - 1s 137ms/step - loss: 1.5656 - sparse_categorical_accuracy: 0.4005 - val_loss: 0.6861 - val_sparse_categorical_accuracy: 0.5000\n",
            "Epoch 5/1000\n",
            "6/6 [==============================] - 1s 139ms/step - loss: 1.5475 - sparse_categorical_accuracy: 0.4090 - val_loss: 0.6875 - val_sparse_categorical_accuracy: 0.5417\n",
            "Epoch 6/1000\n",
            "6/6 [==============================] - 1s 138ms/step - loss: 1.1580 - sparse_categorical_accuracy: 0.5922 - val_loss: 0.7152 - val_sparse_categorical_accuracy: 0.4375\n",
            "Epoch 7/1000\n",
            "6/6 [==============================] - 1s 142ms/step - loss: 1.4923 - sparse_categorical_accuracy: 0.5891 - val_loss: 0.7086 - val_sparse_categorical_accuracy: 0.3958\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/sklearn/metrics/_classification.py:1272: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
            "  _warn_prf(average, modifier, msg_start, len(result))\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 8/1000\n",
            "6/6 [==============================] - 1s 147ms/step - loss: 1.1569 - sparse_categorical_accuracy: 0.5443 - val_loss: 0.6801 - val_sparse_categorical_accuracy: 0.7083\n",
            "Epoch 9/1000\n",
            "6/6 [==============================] - 1s 149ms/step - loss: 0.5235 - sparse_categorical_accuracy: 0.7904 - val_loss: 0.6948 - val_sparse_categorical_accuracy: 0.4792\n",
            "Epoch 10/1000\n",
            "6/6 [==============================] - 1s 134ms/step - loss: 1.0891 - sparse_categorical_accuracy: 0.5658 - val_loss: 0.7144 - val_sparse_categorical_accuracy: 0.4375\n",
            "Epoch 11/1000\n",
            "6/6 [==============================] - 1s 142ms/step - loss: 0.7812 - sparse_categorical_accuracy: 0.7069 - val_loss: 0.7203 - val_sparse_categorical_accuracy: 0.4167\n",
            "Epoch 12/1000\n",
            "6/6 [==============================] - 1s 138ms/step - loss: 0.9517 - sparse_categorical_accuracy: 0.4504 - val_loss: 0.6902 - val_sparse_categorical_accuracy: 0.5208\n",
            "Epoch 13/1000\n",
            "6/6 [==============================] - 1s 137ms/step - loss: 0.8291 - sparse_categorical_accuracy: 0.6327 - val_loss: 0.6971 - val_sparse_categorical_accuracy: 0.4375\n",
            "Epoch 14/1000\n",
            "6/6 [==============================] - 1s 138ms/step - loss: 1.0629 - sparse_categorical_accuracy: 0.5588 - val_loss: 0.7032 - val_sparse_categorical_accuracy: 0.4167\n",
            "Epoch 15/1000\n",
            "6/6 [==============================] - 1s 146ms/step - loss: 0.7953 - sparse_categorical_accuracy: 0.6827 - val_loss: 0.7731 - val_sparse_categorical_accuracy: 0.3958\n",
            "Epoch 16/1000\n",
            "6/6 [==============================] - 1s 147ms/step - loss: 0.7421 - sparse_categorical_accuracy: 0.7421 - val_loss: 0.8017 - val_sparse_categorical_accuracy: 0.4375\n",
            "Epoch 17/1000\n",
            "6/6 [==============================] - 1s 137ms/step - loss: 1.4435 - sparse_categorical_accuracy: 0.4335 - val_loss: 0.7095 - val_sparse_categorical_accuracy: 0.4375\n",
            "Epoch 18/1000\n",
            "6/6 [==============================] - 1s 139ms/step - loss: 0.8671 - sparse_categorical_accuracy: 0.6543 - val_loss: 0.7043 - val_sparse_categorical_accuracy: 0.4167\n",
            "Epoch 19/1000\n",
            "6/6 [==============================] - 1s 140ms/step - loss: 0.7069 - sparse_categorical_accuracy: 0.6502 - val_loss: 0.7592 - val_sparse_categorical_accuracy: 0.4375\n",
            "Epoch 20/1000\n",
            "6/6 [==============================] - 1s 138ms/step - loss: 0.7271 - sparse_categorical_accuracy: 0.7124 - val_loss: 0.7867 - val_sparse_categorical_accuracy: 0.4167\n",
            "Epoch 21/1000\n",
            "6/6 [==============================] - 1s 135ms/step - loss: 0.6236 - sparse_categorical_accuracy: 0.7119 - val_loss: 0.7643 - val_sparse_categorical_accuracy: 0.4167\n",
            "Epoch 22/1000\n",
            "6/6 [==============================] - 1s 151ms/step - loss: 1.2887 - sparse_categorical_accuracy: 0.4479 - val_loss: 0.6865 - val_sparse_categorical_accuracy: 0.4375\n",
            "Epoch 23/1000\n",
            "6/6 [==============================] - 1s 147ms/step - loss: 0.9681 - sparse_categorical_accuracy: 0.4645 - val_loss: 0.6422 - val_sparse_categorical_accuracy: 0.6250\n",
            "Epoch 24/1000\n",
            "6/6 [==============================] - 1s 143ms/step - loss: 0.7037 - sparse_categorical_accuracy: 0.6270 - val_loss: 0.6385 - val_sparse_categorical_accuracy: 0.6250\n",
            "Epoch 25/1000\n",
            "6/6 [==============================] - 1s 134ms/step - loss: 0.5639 - sparse_categorical_accuracy: 0.6967 - val_loss: 0.6440 - val_sparse_categorical_accuracy: 0.5417\n",
            "Epoch 26/1000\n",
            "6/6 [==============================] - 1s 136ms/step - loss: 0.4501 - sparse_categorical_accuracy: 0.8096 - val_loss: 0.7002 - val_sparse_categorical_accuracy: 0.4167\n",
            "Epoch 27/1000\n",
            "6/6 [==============================] - 1s 137ms/step - loss: 0.5937 - sparse_categorical_accuracy: 0.6954 - val_loss: 0.6789 - val_sparse_categorical_accuracy: 0.4167\n",
            "Epoch 28/1000\n",
            "6/6 [==============================] - 1s 138ms/step - loss: 0.4874 - sparse_categorical_accuracy: 0.7526 - val_loss: 0.6187 - val_sparse_categorical_accuracy: 0.6042\n",
            "Restoring model weights from the end of the best epoch.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "Training finished.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 00028: early stopping\n",
            "INFO:tensorflow:Assets written to: /content/custaudio/model/assets\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "Assets written to: /content/custaudio/model/assets\n",
            "\n",
            "Starting evaluation on test data.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "6/6 [==============================] - 4s 297ms/step - loss: 0.6901 - sparse_categorical_accuracy: 0.5625\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Finished evaluation. Loss: 0.6901, Accuracy: 0.5625.\n",
            "\n",
            "Script finished.\n"
          ],
          "name": "stderr"
        }
      ]
    }
  ]
}