train_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "train_colab.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "RF2uyPfxgi8H"
      },
      "source": [
        "# TO USE: \n",
        "#    1. Upload your recorded tone sample .wav file from the SmartAmpPro plugin.\n",
        "#        Note: Optionally, upload two separate .wav files and set flags for two inputs\n",
        "#        Note: Currently supports 32 Bit FLoating Point Wav, or 16 Bit PCM;\n",
        "#                Using 16 Bit PCM (CD quality) allows for smaller .wav files\n",
        "#    2. Select a runtime (GPU recommended)\n",
        "#    3. Change the \"in_file\" variable to match your uploaded file. \n",
        "#        Note: If two .wav files used, set \"out_file\" to match uploaded file\n",
        "#    4 Change the \"name\" file to the desired model name\n",
        "#    5. Click \"Runtime\" then \"Run All\"\n",
        "#    6. To run another training session, reset the runtime to free up RAM\n",
        "#\n",
        "#     Note: Tested on CPU and GPU runtimes.\n",
        "\n",
        "# EDIT THIS SECTION FOR USER INPUTS ###########################################\n",
        "#\n",
        "name = 'my_model'\n",
        "in_file = 'my_input.wav'\n",
        "\n",
        "# If using two mono .wav files, enter uploaded out_file here\n",
        "# If using a single stereo .wav file recorded from SmartAmpPro,\n",
        "#    ensure out_file=''\n",
        "out_file = ''\n",
        "\n",
        "## Training Parameters\n",
        "epochs = 3  # Increase epochs to try extended training to improve model results\n",
        "split_data = 6        # Increase this to reduce RAM usage \n",
        "learning_rate = 0.01  # Agressiveness of training\n",
        "hidden_units = 24     # Number of units in LSTM layer to train\n",
        "\n",
        "input_size = 120 # Using input_size greater than 140 will not work in SmartAmpPro\n",
        "\n",
        "batch_size = 4096 # Number of samples to train at a time\n",
        "test_size = 0.2 # Portion of input data to use in testing (.2 = 20%)\n",
        "###############################################################################\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow.keras import Sequential\n",
        "from tensorflow.keras.layers import LSTM, Conv1D, Dense\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.backend import clear_session\n",
        "from tensorflow.keras.activations import tanh, elu, relu\n",
        "from tensorflow.keras.models import load_model\n",
        "import tensorflow.keras.backend as K\n",
        "from tensorflow.keras.utils import Sequence\n",
        "\n",
        "import os\n",
        "from scipy import signal\n",
        "from scipy.io import wavfile\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import math\n",
        "import h5py\n",
        "import json\n",
        "\n"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U22mDBe4jaf2"
      },
      "source": [
        "if not os.path.exists('models/'+name):\n",
        "    os.makedirs('models/'+name)\n",
        "else:\n",
        "    print(\"A model with the same name already exists. Please choose a new name.\")\n",
        "    exit\n"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WqI-cGt1jaG2"
      },
      "source": [
        "\n",
        "def pre_emphasis_filter(x, coeff=0.95):\n",
        "    return tf.concat([x, x - coeff * x], 1)\n",
        "    \n",
        "def error_to_signal(y_true, y_pred): \n",
        "    \"\"\"\n",
        "    Error to signal ratio with pre-emphasis filter:\n",
        "    \"\"\"\n",
        "    y_true, y_pred = pre_emphasis_filter(y_true), pre_emphasis_filter(y_pred)\n",
        "    return K.sum(tf.pow(y_true - y_pred, 2), axis=0) / K.sum(tf.pow(y_true, 2), axis=0) + 1e-10\n",
        "    \n",
        "def save_wav(name, data):\n",
        "    wavfile.write(name, 44100, data.flatten().astype(np.float32))\n",
        "\n",
        "def normalize(data):\n",
        "    data_max = max(data)\n",
        "    data_min = min(data)\n",
        "    data_norm = max(data_max,abs(data_min))\n",
        "    return data / data_norm\n",
        "\n",
        "\n",
        "'''This is a similar Tensorflow/Keras implementation of the LSTM model from the paper:\n",
        "    \"Real-Time Guitar Amplifier Emulation with Deep Learning\"\n",
        "    https://www.mdpi.com/2076-3417/10/3/766/htm\n",
        "\n",
        "    Uses a stack of two 1-D Convolutional layers, followed by LSTM, followed by \n",
        "    a Dense (fully connected) layer. Three preset training modes are available, \n",
        "    with further customization by editing the code. A Sequential tf.keras model \n",
        "    is implemented here.\n",
        "\n",
        "    Note: RAM may be a limiting factor for the parameter \"input_size\". The wav data\n",
        "      is preprocessed and stored in RAM, which improves training speed but quickly runs out\n",
        "      if using a large number for \"input_size\".  Reduce this if you are experiencing\n",
        "      RAM issues. \n",
        "'''\n",
        "\n",
        "conv1d_strides = 12\n",
        "conv1d_1_strides = 12   \n",
        "conv1d_filters = 4\n",
        "\n",
        "conv1d_KS = conv1d_strides \n",
        "conv1d_1_KS = conv1d_1_strides\n",
        "\n",
        "# Create Sequential Model ###########################################\n",
        "clear_session()\n",
        "model = Sequential()\n",
        "model.add(Conv1D(conv1d_filters, conv1d_KS, strides=conv1d_strides, activation=None, padding='same',input_shape=(input_size,1)))\n",
        "model.add(Conv1D(conv1d_filters, conv1d_1_KS, strides=conv1d_1_strides, activation=None, padding='same'))\n",
        "model.add(LSTM(hidden_units))\n",
        "model.add(Dense(1, activation=None))\n",
        "model.compile(optimizer=Adam(learning_rate=learning_rate), loss=error_to_signal, metrics=[error_to_signal])\n",
        "print(model.summary())\n",
        "\n",
        "# Load and Preprocess Data ###########################################\n",
        "in_rate, stereo_data = wavfile.read(in_file)\n",
        "\n",
        "# If a second wav file is provided, assume each file is mono and load data\n",
        "if out_file != \"\":\n",
        "    out_rate, out_data = wavfile.read(out_file)\n",
        "    in_data = stereo_data\n",
        "# Else, use the stereo wav file channel 1 and channel 2 for training data\n",
        "else:\n",
        "    in_data = stereo_data.T[0]\n",
        "    out_data = stereo_data.T[1]\n",
        "    \n",
        "\n",
        "X_all = in_data.astype(np.float32).flatten()  \n",
        "X_all = normalize(X_all).reshape(len(X_all),1)   \n",
        "y_all = out_data.astype(np.float32).flatten() \n",
        "y_all = normalize(y_all).reshape(len(y_all),1)   \n",
        "\n",
        "# If splitting the data for training, do this part\n",
        "if split_data > 1:\n",
        "    num_split = len(X_all) // split_data\n",
        "    X = X_all[0:num_split*split_data]\n",
        "    y = y_all[0:num_split*split_data]\n",
        "    X_data = np.split(X, split_data)\n",
        "    y_data = np.split(y, split_data)\n",
        "\n",
        "    # Perform training on each split dataset\n",
        "    for i in range(len(X_data)):\n",
        "        print(\"\\nTraining on split data \" + str(i+1) + \"/\" +str(len(X_data)))\n",
        "        X_split = X_data[i]\n",
        "        y_split = y_data[i]\n",
        "\n",
        "        y_ordered = y_split[input_size-1:] \n",
        "\n",
        "        indices = np.arange(input_size) + np.arange(len(X_split)-input_size+1)[:,np.newaxis] \n",
        "        X_ordered = tf.gather(X_split,indices) \n",
        "\n",
        "        shuffled_indices = np.random.permutation(len(X_ordered)) \n",
        "        X_random = tf.gather(X_ordered,shuffled_indices)\n",
        "        y_random = tf.gather(y_ordered, shuffled_indices)\n",
        "\n",
        "        # Train Model ###################################################\n",
        "        model.fit(X_random,y_random, epochs=epochs, batch_size=batch_size, validation_split=0.2)  \n",
        "\n",
        "\n",
        "    model.save('models/'+name+'/'+name+'.h5')\n",
        "\n",
        "# If training on the full set of input data in one run, do this part\n",
        "else:\n",
        "    y_ordered = y_all[input_size-1:] \n",
        "\n",
        "    indices = np.arange(input_size) + np.arange(len(X_all)-input_size+1)[:,np.newaxis] \n",
        "    X_ordered = tf.gather(X_all,indices) \n",
        "\n",
        "    shuffled_indices = np.random.permutation(len(X_ordered)) \n",
        "    X_random = tf.gather(X_ordered,shuffled_indices)\n",
        "    y_random = tf.gather(y_ordered, shuffled_indices)\n",
        "\n",
        "    # Train Model ###################################################\n",
        "    model.fit(X_random,y_random, epochs=epochs, batch_size=batch_size, validation_split=test_size)    \n",
        "\n",
        "    model.save('models/'+name+'/'+name+'.h5')\n",
        "\n",
        "# Run Prediction #################################################\n",
        "#print(\"Running prediction..\")\n",
        "\n",
        "# Get the last 20% of the wav data to run prediction and plot results\n",
        "y_the_rest, y_last_part = np.split(y_all, [int(len(y_all)*.8)])\n",
        "x_the_rest, x_last_part = np.split(X_all, [int(len(X_all)*.8)])\n",
        "y_test = y_last_part[input_size-1:] \n",
        "indices = np.arange(input_size) + np.arange(len(x_last_part)-input_size+1)[:,np.newaxis] \n",
        "X_test = tf.gather(x_last_part,indices) \n",
        "\n",
        "prediction = model.predict(X_test, batch_size=batch_size)\n",
        "\n",
        "save_wav('models/'+name+'/y_pred.wav', prediction)\n",
        "save_wav('models/'+name+'/x_test.wav', x_last_part)\n",
        "save_wav('models/'+name+'/y_test.wav', y_test)\n",
        "\n",
        "# Add additional data to the saved model (like input_size)\n",
        "filename = 'models/'+name+'/'+name+'.h5'\n",
        "f = h5py.File(filename, 'a')\n",
        "grp = f.create_group(\"info\")\n",
        "dset = grp.create_dataset(\"input_size\", (1,), dtype='int16')\n",
        "dset[0] = input_size\n",
        "dset2 = grp.create_dataset(\"conv1d_stride\", (1,), dtype='int16')\n",
        "dset3 = grp.create_dataset(\"conv1d_1_stride\", (1,), dtype='int16')\n",
        "dset[0] = input_size\n",
        "dset2[0] = conv1d_strides\n",
        "dset3[0] = conv1d_1_strides\n",
        "\n",
        "f.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OV7hybs5eXBS"
      },
      "source": [
        "# Generate json model ################################\r\n",
        "filename = 'models/'+name+'/'+ name +'.h5'\r\n",
        "json_filename = 'models/'+name+'/'+ name\r\n",
        "f = h5py.File(filename, 'r')\r\n",
        "\r\n",
        "# Load the model data\r\n",
        "data = {}\r\n",
        "for layer in f[\"model_weights\"].keys():\r\n",
        "    if layer not in data.keys():\r\n",
        "        data[layer] = {}\r\n",
        "    for item in f[\"model_weights\"][layer][layer]:\r\n",
        "        if item not in data[layer].keys():\r\n",
        "            data[layer][item] = {}\r\n",
        "            try:\r\n",
        "                data[layer][item] = f[\"model_weights\"][layer][layer][item][:].tolist()\r\n",
        "            except:\r\n",
        "                data[layer][\"kernel:0\"] = f[\"model_weights\"][layer][layer][item][\"kernel:0\"][:].tolist()\r\n",
        "                data[layer][\"bias:0\"] = f[\"model_weights\"][layer][layer][item][\"bias:0\"][:].tolist()\r\n",
        "        else:\r\n",
        "            try:\r\n",
        "                data[layer][item] = f[\"model_weights\"][layer][layer][item][:].tolist()\r\n",
        "            except:\r\n",
        "                data[layer][\"kernel:0\"] = f[\"model_weights\"][layer][layer][item][\"kernel:0\"][:].tolist()\r\n",
        "                data[layer][\"bias:0\"] = f[\"model_weights\"][layer][layer][item][\"bias:0\"][:].tolist()\r\n",
        "\r\n",
        "input_size = f[\"info\"][\"input_size\"][0]\r\n",
        "conv1d_strides = f[\"info\"][\"conv1d_stride\"][0]\r\n",
        "conv1d_1_strides = f[\"info\"][\"conv1d_1_stride\"][0]\r\n",
        "data['input_size'] = int(input_size)\r\n",
        "print(\"input_size: \",input_size)\r\n",
        "data['conv1d_stride'] = int(conv1d_strides)\r\n",
        "data['conv1d_1_stride'] = int(conv1d_1_strides)\r\n",
        "f.close()\r\n",
        "\r\n",
        "with open(json_filename + \".json\", 'w') as outfile:\r\n",
        "    json.dump(data, outfile)\r\n",
        "print(\"SmartAmpPro model generated: \", json_filename + \".json\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xJFGh1hDimp-"
      },
      "source": [
        "# Model Evaluation and Plotting\r\n",
        "\r\n",
        "def print_accuracy(loss):\r\n",
        "    print(\"Model Accuracy: \", str(round((1 - round(loss,2))*100)),\"%\")\r\n",
        "\r\n",
        "def error_to_signal(y, y_pred, use_filter=1):\r\n",
        "    \"\"\"\r\n",
        "    Error to signal ratio with pre-emphasis filter:\r\n",
        "    https://www.mdpi.com/2076-3417/10/3/766/htm\r\n",
        "    \"\"\"\r\n",
        "    if use_filter == 1:\r\n",
        "        y, y_pred = pre_emphasis_filter(y), pre_emphasis_filter(y_pred)\r\n",
        "    return np.sum(np.power(y - y_pred, 2)) / (np.sum(np.power(y, 2) + 1e-10))\r\n",
        "\r\n",
        "\r\n",
        "def pre_emphasis_filter(x, coeff=0.95):\r\n",
        "    return np.concatenate([x, np.subtract(x, np.multiply(x, coeff))])\r\n",
        "\r\n",
        "\r\n",
        "def read_wave(wav_file):\r\n",
        "    # Extract Audio and framerate from Wav File\r\n",
        "    fs, signal = wavfile.read(wav_file)\r\n",
        "    return signal, fs\r\n",
        "\r\n",
        "\r\n",
        "\r\n",
        "def analyze_pred_vs_actual(output_wav, pred_wav, input_wav, model_name, show_plots, path):\r\n",
        "    \"\"\"Generate plots to analyze the predicted signal vs the actual\r\n",
        "    signal.\r\n",
        "\r\n",
        "    Inputs:\r\n",
        "        output_wav : The actual signal, by default will use y_test.wav from the test.py output\r\n",
        "        pred_wav : The predicted signal, by default will use y_pred.wav from the test.py output\r\n",
        "        input_wav : The pre effect signal, by default will use x_test.wav from the test.py output\r\n",
        "        model_name : Used to add the model name to the plot .png filename\r\n",
        "        path   :   The save path for generated .png figures\r\n",
        "        show_plots : Default is 1 to show plots, 0 to only generate .png files and suppress plots\r\n",
        "\r\n",
        "    1. Plots the two signals\r\n",
        "    2. Calculates Error to signal ratio the same way Pedalnet evauluates the model for training\r\n",
        "    3. Plots the absolute value of pred_signal - actual_signal  (to visualize abs error over time)\r\n",
        "\r\n",
        "    \"\"\"\r\n",
        "\r\n",
        "    # Read the input wav file\r\n",
        "    signal3, fs3 = read_wave(input_wav)\r\n",
        "\r\n",
        "    # Read the output wav file\r\n",
        "    signal1, fs = read_wave(output_wav)\r\n",
        "\r\n",
        "    Time = np.linspace(0, len(signal1) / fs, num=len(signal1))\r\n",
        "    fig, (ax3, ax1, ax2) = plt.subplots(3, sharex=True, figsize=(13, 8))\r\n",
        "    fig.suptitle(\"Predicted vs Actual Signal\")\r\n",
        "    ax1.plot(Time, signal1, label=output_wav, color=\"red\")\r\n",
        "\r\n",
        "    # Read the predicted wav file\r\n",
        "    signal2, fs2 = read_wave(pred_wav)\r\n",
        "\r\n",
        "    Time2 = np.linspace(0, len(signal2) / fs2, num=len(signal2))\r\n",
        "    ax1.plot(Time2, signal2, label=pred_wav, color=\"green\")\r\n",
        "    ax1.legend(loc=\"upper right\")\r\n",
        "    ax1.set_xlabel(\"Time (s)\")\r\n",
        "    ax1.set_ylabel(\"Amplitude\")\r\n",
        "    ax1.set_title(\"Wav File Comparison\")\r\n",
        "    ax1.grid(\"on\")\r\n",
        "\r\n",
        "    error_list = []\r\n",
        "    for s1, s2 in zip(signal1, signal2):\r\n",
        "        error_list.append(abs(s2 - s1))\r\n",
        "\r\n",
        "    # Calculate error to signal ratio with pre-emphasis filter as\r\n",
        "    #    used to train the model\r\n",
        "    e2s = error_to_signal(signal1, signal2)\r\n",
        "    print_accuracy(e2s)\r\n",
        "    \r\n",
        "    e2s_no_filter = error_to_signal(signal1, signal2, use_filter=0)\r\n",
        "    print(\"Error to signal (with pre-emphasis filter): \", e2s)\r\n",
        "    print(\"Error to signal (no pre-emphasis filter): \", e2s_no_filter)\r\n",
        "    fig.suptitle(\"Predicted vs Actual Signal (error to signal: \" + str(round(e2s, 4)) + \")\")\r\n",
        "    # Plot signal difference\r\n",
        "    signal_diff = signal2 - signal1\r\n",
        "    ax2.plot(Time2, error_list, label=\"signal diff\", color=\"blue\")\r\n",
        "    ax2.set_xlabel(\"Time (s)\")\r\n",
        "    ax2.set_ylabel(\"Amplitude\")\r\n",
        "    ax2.set_title(\"abs(pred_signal-actual_signal)\")\r\n",
        "    ax2.grid(\"on\")\r\n",
        "\r\n",
        "    # Plot the original signal\r\n",
        "    Time3 = np.linspace(0, len(signal3) / fs3, num=len(signal3))\r\n",
        "    ax3.plot(Time3, signal3, label=input_wav, color=\"purple\")\r\n",
        "    ax3.legend(loc=\"upper right\")\r\n",
        "    ax3.set_xlabel(\"Time (s)\")\r\n",
        "    ax3.set_ylabel(\"Amplitude\")\r\n",
        "    ax3.set_title(\"Original Input\")\r\n",
        "    ax3.grid(\"on\")\r\n",
        "\r\n",
        "    # Save the plot\r\n",
        "    plt.savefig(path+'/'+model_name + \"_signal_comparison_e2s_\" + str(round(e2s, 4)) + \".png\", bbox_inches=\"tight\")\r\n",
        "    # Create a zoomed in plot of 0.01 seconds centered at the max input signal value\r\n",
        "    sig_temp = signal1.tolist()\r\n",
        "    plt.axis(\r\n",
        "        [\r\n",
        "            Time3[sig_temp.index((max(sig_temp)))] - 0.005,\r\n",
        "            Time3[sig_temp.index((max(sig_temp)))] + 0.005,\r\n",
        "            min(signal2),\r\n",
        "            max(signal2),\r\n",
        "        ]\r\n",
        "    )\r\n",
        "   \r\n",
        "    plt.savefig(path+'/'+model_name + \"_Detail_signal_comparison_e2s_\" + str(round(e2s, 4)) + \".png\", bbox_inches=\"tight\")\r\n",
        "    #plt.show() # Uncomment this line to show Detail plot\r\n",
        "    # Reset the axis\r\n",
        "    plt.axis([0, Time3[-1], min(signal2), max(signal2)])\r\n",
        "\r\n",
        "\r\n",
        "print(\"Plotting results..\")\r\n",
        "\r\n",
        "\r\n",
        "analyze_pred_vs_actual('models/'+name+'/y_test.wav',\r\n",
        "                       'models/'+name+'/y_pred.wav', \r\n",
        "                       'models/'+name+'/x_test.wav',\r\n",
        "                       name,\r\n",
        "                       1,\r\n",
        "                       'models/'+name)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}