VoiceEraser Training Pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Keras with MENext backend training routine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:50:14.107218Z",
     "start_time": "2019-06-25T19:50:10.351613Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import subprocess\n",
    "from pdb import set_trace\n",
    "import gc\n",
    "import time\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from IPython import display\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "#!pip install pyfftw\n",
    "import pyfftw\n",
    "from pyfftw.interfaces.scipy_fftpack import rfft, irfft\n",
    "\n",
    "#!pip install librosa\n",
    "import librosa\n",
    "\n",
    "#!pip install soundfile\n",
    "import soundfile as sf\n",
    "\n",
    "import pywt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:56:14.767250Z",
     "start_time": "2019-06-25T19:56:14.763133Z"
    }
   },
   "outputs": [],
   "source": [
    "# Data loading\n",
    "# Contains full paths to files, \n",
    "# with 'simPath' column for merged audio (X)\n",
    "# and 'crowdSegPath' for target (y)\n",
    "train_df_path = './train.csv' \n",
    "test_df_path = './test.csv'\n",
    "val_df_path = './val.csv'\n",
    "audio_len = 65536 # length of audio clips\n",
    "\n",
    "# Model\n",
    "previous_weights = None#'./train1/model_5_weights-improvement-02-0.26.hdf5' # continue from previous training\n",
    "\n",
    "# Training\n",
    "epochs = 2000 # number of epochs to train\n",
    "batch_size=16 # batch size for training\n",
    "train_num = 1 # training results will be saved in the trainNum folder\n",
    "model_num = 1 # model results will be saved as model_num file\n",
    "\n",
    "# number of GPUs used for training. \n",
    "# For multiple GPUs: make sure ((num_train) % batch_size) / num_gpu == 0\n",
    "num_gpus = 1\n",
    "\n",
    "# Data transformation used for training\n",
    "datatype='wavelet' #[wavelet, time (no transformation)]\n",
    "\n",
    "# Saving results\n",
    "train_dir = './train'\n",
    "save_dir = '/data/output_data/'\n",
    "wavelet_dir = '/data/simulated_wavelet10' # if training using wavelet, make sure this directory exists"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:50:23.210641Z",
     "start_time": "2019-06-25T19:50:20.244850Z"
    }
   },
   "outputs": [],
   "source": [
    "#import tensorflow as tf\n",
    "import keras\n",
    "from keras import Input, metrics, losses\n",
    "from keras.models import Sequential, Model, load_model, model_from_json\n",
    "from keras.layers import Input, Dense, Activation, Dropout, Flatten, BatchNormalization, Lambda\n",
    "from keras.layers import Conv1D, Conv2DTranspose, MaxPooling1D, AveragePooling1D, UpSampling1D, UpSampling2D, K\n",
    "from keras.layers import Concatenate, LeakyReLU\n",
    "from keras.layers.convolutional import _UpSampling\n",
    "from keras.legacy import interfaces\n",
    "from keras.callbacks import ModelCheckpoint, EarlyStopping #, TensorBoard\n",
    "from keras.regularizers import l1, l2, l1_l2\n",
    "from keras.initializers import VarianceScaling\n",
    "from keras.callbacks import CSVLogger\n",
    "from keras.utils import multi_gpu_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:50:23.360287Z",
     "start_time": "2019-06-25T19:50:23.355620Z"
    },
    "code_folding": [
     1,
     7
    ]
   },
   "outputs": [],
   "source": [
    "# Second level decomposition, Haar\n",
    "def time2wavelet(x):\n",
    "        cA2, cD2, cD1 = pywt.wavedec(x, 'haar', mode='constant', level=2)\n",
    "        w = np.concatenate((cA2, cD2, cD1)) # enhancing higher frequency space\n",
    "        w = w / np.max(np.abs(w))\n",
    "        return w\n",
    "\n",
    "def wavelet2time(w, weights=[1, 1, 1]):\n",
    "    coeffs = w[:16384], w[16384:32768], w[32768:]\n",
    "    x = pywt.waverec(coeffs, 'haar')\n",
    "    x = x / np.max(np.abs(x))\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:55:30.408190Z",
     "start_time": "2019-06-25T19:55:30.345845Z"
    },
    "code_folding": [
     3,
     26,
     31,
     43,
     49,
     74,
     86
    ]
   },
   "outputs": [],
   "source": [
    "# Generating data\n",
    "class DataGenerator(keras.utils.Sequence):\n",
    "    \"\"\"Generate data for Keras\"\"\"\n",
    "    def __init__(self, df, batch_size=50, dataset='train', shuffle=True, save_transform=False, datatype='time', \n",
    "                 input_dims=[32000], data_format='channels_first'):\n",
    "        \"\"\"\n",
    "        * df: data frame that stores the meta data of the dataset. \n",
    "            The generator will iterate through this dataframe.\n",
    "        * batch_size: batch sample of each iteration\n",
    "        * dataset: use to label the dataset of the current generator\n",
    "        * shuffle: whether or not shuffle the dataframe before itering over it. Default True!\n",
    "        * datatype: type of data to load: 'time' or 'fft'\n",
    "        * savefft: save the fft transformation used to train the data as a file to speed up training over epochs\n",
    "        * input_dims: dimension of the input\n",
    "        * duration: duration of the data to get (the last index). By default, get all data. \n",
    "        \"\"\"\n",
    "        self.batch_size = batch_size\n",
    "        self.dims = input_dims\n",
    "        self.dataset = dataset\n",
    "        self.shuffle = shuffle\n",
    "        self.datatype = datatype\n",
    "        self.df = df\n",
    "        self.save_transform = save_transform\n",
    "        self.data_format = data_format\n",
    "        self.on_epoch_end()\n",
    "    \n",
    "    def __len__(self):\n",
    "        \"\"\"Denotes the number of batches per epoch\"\"\"\n",
    "        num_batches = self.df.shape[0] // self.batch_size + (self.df.shape[0] % self.batch_size > 0)\n",
    "        return num_batches\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"Generate one batch of data\"\"\"\n",
    "        start_index = index * self.batch_size\n",
    "        end_index = (index+1) * self.batch_size\n",
    "        if end_index > self.df.shape[0]:\n",
    "            end_index = None\n",
    "        rows = self.df.index[start_index:end_index]\n",
    "        \n",
    "        X, y = self.__data_generation(rows)\n",
    "\n",
    "        return X, y\n",
    "    \n",
    "    def on_epoch_end(self):\n",
    "        \"\"\"Update indexes after each epoch\"\"\"\n",
    "        start_time = time.clock()\n",
    "        if self.shuffle:\n",
    "            self.df = self.df.sample(frac=1, replace=False).reset_index(drop=True)\n",
    "    \n",
    "    def __data_generation(self, rows):\n",
    "        nrows = len(rows)\n",
    "        # crowd + vice\n",
    "        dims = [nrows] +  [1] + list(self.dims) \\\n",
    "                if self.data_format == 'channels_first' \\\n",
    "                else [nrows] +  list(self.dims) + [1] # crowd + noise\n",
    "        X = np.empty(dims)\n",
    "        # crowd and voice separately\n",
    "        dims = [nrows] +  [1] + list(self.dims) \\\n",
    "                if self.data_format == 'channels_first' \\\n",
    "                else [nrows] +  list(self.dims) + [1] # crowd + noise\n",
    "        y = np.empty(dims)\n",
    "        \n",
    "        func = self.__load_wavelet if self.datatype == 'wavelet' else self.__load_data\n",
    "        #start_time = time.clock()\n",
    "        for n, r in enumerate(rows):\n",
    "            # Load crowd + voice file\n",
    "            if self.data_format == 'channels_first':\n",
    "                X[n, 0, :] = func(self.df.loc[r, 'simPath'])\n",
    "                y[n, 0, :] = func(self.df.loc[r, 'crowdSegPath'])\n",
    "            else:\n",
    "                X[n, :, 0] = func(self.df.loc[r, 'simPath'])\n",
    "                y[n, :, 0] = func(self.df.loc[r, 'crowdSegPath'])\n",
    "        return X, y\n",
    "    \n",
    "    def __load_data(self, path=None, normalize=True, *args, **kwargs):\n",
    "        y, fs = sf.read(path)\n",
    "        # Take only the first channel\n",
    "        y = y[:self.dims[0], 0]\n",
    "        y[np.isnan(y)] = 0\n",
    "        \n",
    "        if normalize:\n",
    "            y = 1.0/np.max(abs(y), axis=0) * y\n",
    "            #y = (y - y.min(axis=0)) / (y.max(axis=0) - y.min(axis=0))\n",
    "            \n",
    "        return y\n",
    "    \n",
    "    def __load_wavelet(self, path=None, normalize=True):\n",
    "        \"\"\"\n",
    "        Load the wavelet coefficients of the times series file.\n",
    "        If not already exist, create it.\n",
    "        normalize: normalize the time series to 0 to 1 before converting fft\n",
    "        self.save_transform: [True|False] save the new fft files in a drive, given a path here.\n",
    "        \"\"\"\n",
    "        wl_path = path.replace('.wav', '_wavelet.npz').replace('simulated_wav', 'simulated_wavelet')\n",
    "        if not os.path.isfile(wl_path): # file specified but not there!\n",
    "            # load file\n",
    "            y, fs = sf.read(path)\n",
    "            # Take only the first two seconds and first channel\n",
    "            y = y[:self.dims[0], 0]\n",
    "            y[np.isnan(y)] = 0\n",
    "            #set_trace()\n",
    "            if normalize: # -1 and 1\n",
    "                y = 1.0/np.max(abs(y), axis=0) * y\n",
    "                #y = (y - y.min(axis=0)) / (y.max(axis=0) - y.min(axis=0))\n",
    "\n",
    "            w = time2wavelet(y)\n",
    "\n",
    "            if self.save_transform:\n",
    "                np.savez(wl_path, w=w)\n",
    "        else: # already exists, simply load it\n",
    "            w = np.load(wl_path)['w']\n",
    "            \n",
    "        return w\n",
    "    \n",
    "test_ = True\n",
    "save_transform=True\n",
    "start_fresh = True # recalculating any transformations\n",
    "data_format = 'channels_first' # batch x channel x time; channel_last: batch x time x channel (slow on GPU)\n",
    "\n",
    "# Start fresh\n",
    "try:\n",
    "    cmd = 'mkdir ' + wavelet_dir\n",
    "    subprocess.call(cmd, shell=True)\n",
    "except:\n",
    "    pass\n",
    "    \n",
    "if start_fresh:\n",
    "    cmd = 'rm -f ' + wavelet_dir +'/*.npz'\n",
    "    subprocess.call(cmd, shell=True)\n",
    "\n",
    "# Initialize data generator\n",
    "if test_:\n",
    "    df_train = pd.read_csv(train_df_path)\n",
    "    df_test = pd.read_csv(test_df_path)\n",
    "    df_val = pd.read_csv(val_df_path)\n",
    "    train_generator = DataGenerator(df_train.reset_index(drop=True), batch_size=batch_size, shuffle=True, \n",
    "                                    save_transform=save_transform, dataset='train', input_dims=[audio_len],\n",
    "                                    datatype=datatype, data_format=data_format)\n",
    "    test_generator = DataGenerator(df_test.reset_index(drop=True), batch_size=batch_size, shuffle=False, \n",
    "                                   save_transform=save_transform, dataset='test', input_dims=[audio_len], \n",
    "                                   datatype=datatype, data_format=data_format)\n",
    "    val_generator = DataGenerator(df_val.reset_index(drop=True), batch_size=batch_size, shuffle=False, \n",
    "                                  save_transform=save_transform, dataset='val', input_dims=[audio_len], \n",
    "                                  datatype=datatype,  data_format=data_format)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:55:31.783247Z",
     "start_time": "2019-06-25T19:55:31.689917Z"
    }
   },
   "outputs": [],
   "source": [
    "X, y = train_generator.__getitem__(0)\n",
    "print(X.shape)\n",
    "print(y.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define the denoising autoencoder with U-net structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:55:37.340295Z",
     "start_time": "2019-06-25T19:55:37.331775Z"
    },
    "code_folding": [
     1,
     5,
     19,
     24,
     29,
     38
    ]
   },
   "outputs": [],
   "source": [
    "# custom losses\n",
    "def corr_loss(y_true, y_pred):\n",
    "    r = corr_metric(y_true, y_pred)\n",
    "    return 1 - K.square(r)\n",
    "    \n",
    "def corr_metric(y_true, y_pred):\n",
    "    x = y_true\n",
    "    y = y_pred\n",
    "    mx = K.mean(x)\n",
    "    my = K.mean(y)\n",
    "    xm, ym = x-mx, y-my\n",
    "    r_num = K.sum(xm * ym) #K.sum(tf.multiply(xm, ym))\n",
    "    r_den = K.sqrt(K.sum(K.square(xm)) * K.sum(K.square(ym))) #K.sqrt(tf.multiply(K.sum(K.square(xm)), K.sum(K.square(ym))))\n",
    "    r = r_num / (r_den + 1E-12) # prevent nans\n",
    "\n",
    "    r = K.maximum(K.minimum(r, 1.0), -1.0)\n",
    "        \n",
    "    return r\n",
    "\n",
    "def corr_crowd_metric(y_true, y_pred): # channels_last, tf backend\n",
    "    y_true_crowd = K.slice(y_true, [0, 0, 0], [-1, -1, 1])\n",
    "    y_pred_crowd = K.slice(y_pred, [0, 0, 0], [-1, -1, 1]) \n",
    "    return corr_metric(y_true_crowd, y_pred_crowd)\n",
    "\n",
    "def corr_voice_metric(y_true, y_pred): # channels_last, tf backend\n",
    "    y_true_voice = K.slice(y_true, [0, 0, 1], [-1, -1, 1])\n",
    "    y_pred_voice = K.slice(y_pred, [0, 0, 1], [-1, -1, 1])\n",
    "    return corr_metric(y_true_voice, y_pred_voice)\n",
    "\n",
    "def corr_crowd_mxnet(y_true, y_pred):\n",
    "    if data_format == 'channels_last':\n",
    "        y_true_crowd = Lambda(lambda x: x[:, :, 0])(y_true)\n",
    "        y_pred_crowd = Lambda(lambda x: x[:, :, 0])(y_pred)\n",
    "    else:\n",
    "        y_true_crowd = Lambda(lambda x: x[:, 0, :])(y_true)\n",
    "        y_pred_crowd = Lambda(lambda x: x[:, 0, :])(y_pred)\n",
    "    return corr_metric(y_true_crowd, y_pred_crowd) # channels_first\n",
    "\n",
    "def corr_voice_mxnet(y_true, y_pred):\n",
    "    if data_format == 'channels_last':\n",
    "        y_true_crowd = Lambda(lambda x: x[:, :, 1])(y_true)\n",
    "        y_pred_crowd = Lambda(lambda x: x[:, :, 1])(y_pred)\n",
    "    else:\n",
    "        y_true_crowd = Lambda(lambda x: x[:, 1, :])(y_true)\n",
    "        y_pred_crowd = Lambda(lambda x: x[:, 1, :])(y_pred)\n",
    "    return corr_metric(y_true_crowd, y_pred_crowd) # channels_first\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:55:37.772946Z",
     "start_time": "2019-06-25T19:55:37.767276Z"
    },
    "code_folding": [
     0,
     25
    ]
   },
   "outputs": [],
   "source": [
    "class UpSampling1D_2(_UpSampling):\n",
    "    \"\"\"Upsampling layer for 1D inputs.\n",
    "    Repeats each temporal step `size` times along the time axis.\n",
    "    # Arguments\n",
    "        size: integer. Upsampling factor.\n",
    "    # Input shape\n",
    "        3D tensor with shape: `(batch, steps, features)`.\n",
    "    # Output shape\n",
    "        3D tensor with shape: `(batch, upsampled_steps, features)`.\n",
    "    \"\"\"\n",
    "\n",
    "    @interfaces.legacy_upsampling1d_support\n",
    "    def __init__(self, size=2, **kwargs):\n",
    "        super(UpSampling1D_2, self).__init__((int(size),), 'channels_first', **kwargs)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        output = K.repeat_elements(inputs, self.size[0], axis=2)\n",
    "        return output\n",
    "\n",
    "    def get_config(self):\n",
    "        config = super(UpSampling1D_2, self).get_config()\n",
    "        config['size'] = self.size[0]\n",
    "        config.pop('data_format')\n",
    "        return config\n",
    "\n",
    "def UpSampling1D_interpolate(input_tensor, size, interpolation='bilinear', data_format='channels_last'):\n",
    "    x = Lambda(lambda x: K.expand_dims(x, axis=1))(input_tensor)\n",
    "    x = UpSampling2D(size=(size, 1), data_format=data_format, interpolation=interpolation)(x)\n",
    "    x = Lambda(lambda x: K.squeeze(x, axis=1))(x)\n",
    "    return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:56:23.509501Z",
     "start_time": "2019-06-25T19:56:23.242728Z"
    },
    "code_folding": [
     4,
     18,
     64
    ],
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "K.clear_session()\n",
    "alpha_leakyrelu = 0.1\n",
    "axis = 1 if data_format == 'channels_first' else 2\n",
    "\n",
    "def conv_block(x, filter_size=64, nconv='double', kernel_size=3, padding='same'):\n",
    "    if nconv=='double':\n",
    "        conv = Conv1D(filters=filter_size, kernel_size=kernel_size, padding=padding, activation=None, \n",
    "                     data_format=data_format)(x)\n",
    "        #conv = BatchNormalization(axis=axis)(conv)\n",
    "        conv = LeakyReLU(alpha_leakyrelu)(conv)\n",
    "    else:\n",
    "        conv = x\n",
    "    conv = Conv1D(filters=filter_size, kernel_size=kernel_size, padding=padding, activation=None, \n",
    "                 data_format=data_format)(conv)\n",
    "    conv = BatchNormalization(axis=axis)(conv)\n",
    "    conv = LeakyReLU(alpha_leakyrelu)(conv)\n",
    "    return conv\n",
    "\n",
    "def build_generator():\n",
    "    inputs = Input(shape=(1, audio_len)) if  data_format == 'channels_first' \\\n",
    "                                         else Input(shape=(audio_len, 1))\n",
    "    # Encoder\n",
    "    conv1 = conv_block(inputs, 64, 'double') # L x 64\n",
    "    pool1 = MaxPooling1D(pool_size=2, padding='same', data_format=data_format)(conv1) # L/2 x 64\n",
    "    \n",
    "    conv2 = conv_block(pool1, 128, 'double') # L/2 x 128\n",
    "    #conv2 = Dropout(0.5)(conv2) # drop out layer regularization\n",
    "    pool2 = MaxPooling1D(pool_size=2, padding='same', data_format=data_format)(conv2) # L/4 x 128\n",
    "    \n",
    "    conv3 = conv_block(pool2, 256, 'double') # L/4 x 256\n",
    "    #conv3 = Dropout(0.5)(conv3) # drop out layer regularization\n",
    "    pool3 = MaxPooling1D(pool_size=2, padding='same', data_format=data_format)(conv3) # L/8 x 256\n",
    "    \n",
    "    conv4 = conv_block(pool3, 256, 'double') # L/8 x 256\n",
    "    #conv4 = Dropout(0.5)(conv4) # drop out layer regularization\n",
    "\n",
    "    \n",
    "    # Decoder\n",
    "    upsamp1 = UpSampling1D(size=2)(conv4) if data_format == 'channels_last' \\\n",
    "                                            else UpSampling1D_2(size=2)(conv4) # L/4 x 256\n",
    "    deconv1 = Concatenate(axis=axis)([upsamp1, conv3]) # U: L/4 x 512 (2 x 256)\n",
    "    deconv1 = conv_block(deconv1, 128, 'double') # L/4 x 128\n",
    "    \n",
    "    \n",
    "    upsamp2 = UpSampling1D(size=2)(deconv1) if data_format == 'channels_last' \\\n",
    "                                            else UpSampling1D_2(size=2)(deconv1) # L/2 x 128\n",
    "    deconv2 = Concatenate(axis=axis)([upsamp2, conv2]) # U: L/2 x 256 (2 x 128)\n",
    "    deconv2 = conv_block(deconv2, 64, 'double') # L/2 x 64\n",
    "    \n",
    "    upsamp3 = UpSampling1D(size=2)(deconv2) if data_format == 'channels_last' \\\n",
    "                                            else UpSampling1D_2(size=2)(deconv2) # L x 64\n",
    "    deconv3 = Concatenate(axis=axis)([upsamp3, conv1]) # U: L x 129 (2 x 64 + 1)\n",
    "    deconv3 = Conv1D(filters=16, kernel_size=3, padding='same', activation=None, \n",
    "                  data_format=data_format)(deconv3)  # L x 16\n",
    "    deconv3 = LeakyReLU(alpha_leakyrelu)(deconv3)\n",
    "    \n",
    "    deconv4 = Concatenate(axis=axis)([deconv3, inputs]) # U: L x 17  (16 + 1)\n",
    "    deconv4 = Conv1D(filters=1, kernel_size=3, padding='same', activation='tanh', \n",
    "                  data_format=data_format)(deconv4)  # L x 2\n",
    "    \n",
    "    model = Model(inputs, deconv4)\n",
    "    \n",
    "    return model\n",
    "\n",
    "def backend_agnostic_compile(model, loss, optimizer, metrics, num_gpus=8):\n",
    "    if keras.backend._backend == 'mxnet':\n",
    "        gpu_list = [\"gpu(%d)\" % i for i in range(num_gpus)]\n",
    "        model.compile(loss=loss,\n",
    "            optimizer=optimizer,\n",
    "            metrics=metrics, \n",
    "            context = gpu_list)\n",
    "    else:\n",
    "        model = multi_gpu_model(model, gpus=num_gpus)\n",
    "        model.compile(optimizer=optimizer, loss=loss,\n",
    "                       metrics=metrics)\n",
    "        \n",
    "    return model\n",
    "\n",
    "model = build_generator()\n",
    "opt = keras.optimizers.Adam(lr=0.0001, decay=1e-5)\n",
    "\n",
    "# load previous model\n",
    "if previous_weights:\n",
    "    model.load_weights(previous_weights)\n",
    "\n",
    "if num_gpus > 1:\n",
    "    # check if the num batches can be evenly distributed\n",
    "    if (df_train.shape[0] % batch_size) / num_gpus !=0 or \\\n",
    "        (df_test.shape[0] % batch_size) / num_gpus !=0 or \\\n",
    "        (df_val.shape[0] % batch_size) / num_gpus !=0:\n",
    "        raise(BaseException('Cannot train on multile gpus!\\nMake sure ((num_train) % batch_size) / num_gpu == 0'))\n",
    "    else:\n",
    "        model = backend_agnostic_compile(model, optimizer=opt, \n",
    "                         loss = corr_loss, #losses.mean_squared_error, \n",
    "                         metrics =[corr_crowd_mxnet, metrics.mean_squared_error], #corr_crowd_metric, corr_voice_metric,\n",
    "                         num_gpus=num_gpus)\n",
    "else:\n",
    "    model.compile(optimizer=opt,\n",
    "                  loss=losses.mean_squared_error, # with regularization added on\n",
    "                  metrics=[corr_crowd_mxnet, metrics.mean_squared_error])\n",
    "\n",
    "print(model.summary())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:59:12.289860Z",
     "start_time": "2019-06-25T19:58:50.196846Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    os.mkdir(os.path.join(train_dir,\"train\"+str(train_num)))\n",
    "except:\n",
    "    pass\n",
    "checkfilepath = os.path.join(train_dir, \"train\"+str(train_num)+\"/model_\"+str(model_num)+\\\n",
    "            \"_weights-improvement-{epoch:02d}-{val_loss:.2f}.hdf5\")\n",
    "checkpoint = ModelCheckpoint(checkfilepath, monitor='val_loss', verbose=0, save_best_only=True)\n",
    "#tensorboard = TensorBoard(log_dir=\"./train\"+str(train_num), histogram_freq=5,write_graph=True,write_images=True)\n",
    "csv_logger = CSVLogger(os.path.join(train_dir, 'train'+str(train_num)+'/model_'+str(model_num)+'_training.log'))\n",
    "early_stop = EarlyStopping(monitor='val_loss', patience=20)\n",
    "\n",
    "history = model.fit_generator(train_generator,\n",
    "                              epochs=epochs,\n",
    "                              validation_data=val_generator, \n",
    "                              callbacks = [csv_logger, checkpoint, early_stop],# tensorboard, \n",
    "                              verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:59:25.794252Z",
     "start_time": "2019-06-25T19:59:25.679553Z"
    }
   },
   "outputs": [],
   "source": [
    "# Serialize the model to json\n",
    "model.save(os.path.join(train_dir, 'train'+str(train_num)+'/CNN_model_'+str(model_num)+'_corr_loss.hdf5'))\n",
    "model_json = model.to_json()\n",
    "with open(os.path.join(train_dir, \"train\"+str(train_num)+\"/CNN_model_\"+str(model_num)+\\\n",
    "                       \"_corr_loss_specification.json\"), \"w\") as json_file:\n",
    "    json_file.write(model_json)\n",
    "#Serialize weights to HDF5\n",
    "model.save_weights(os.path.join(train_dir, \"train\"+str(train_num)+\\\n",
    "                                \"/CNN_model_\"+str(model_num)+\"_corr_loss_weights.h5\"))\n",
    "# save history\n",
    "try:\n",
    "    with open( os.path.join(train_dir, 'train'+str(train_num)+'/history_model_'+\\\n",
    "                            str(model_num)+'_corr_loss.json'), 'w') as f:\n",
    "        json.dump(history.history, f)\n",
    "except:\n",
    "    pass\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:59:34.651609Z",
     "start_time": "2019-06-25T19:59:34.248846Z"
    }
   },
   "outputs": [],
   "source": [
    "# Do some plots\n",
    "fig, axs = plt.subplots(1, 2)\n",
    "history_df = pd.read_csv(os.path.join(train_dir, 'train'+str(train_num)+\\\n",
    "                                      '/model_'+str(model_num)+'_training.log'))\n",
    "axs[0].plot(history_df['loss'])\n",
    "axs[0].plot(history_df['val_loss'])\n",
    "axs[0].set_title('loss')\n",
    "axs[1].plot(history_df['corr_crowd_mxnet'])\n",
    "axs[1].plot(history_df['val_corr_crowd_mxnet'])\n",
    "axs[1].set_title('crowd correlation')\n",
    "fig.set_size_inches(12, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save outputs on the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:59:37.011829Z",
     "start_time": "2019-06-25T19:59:37.007143Z"
    },
    "code_folding": [
     0,
     19
    ]
   },
   "outputs": [],
   "source": [
    "def load_autoencoder_model(model_file, weight_file):\n",
    "    # load json and create model\n",
    "    json_file = open(model_file, 'r')\n",
    "    loaded_model_json = json_file.read()\n",
    "    json_file.close()\n",
    "    model = model_from_json(loaded_model_json)\n",
    "    # load weights into new model\n",
    "    model.load_weights(weight_file)\n",
    "    \n",
    "    #model = load_model(model_file) if saved as hdf5 directly\n",
    "    \n",
    "    opt = keras.optimizers.Adam(lr=0.1, decay=0.01)\n",
    "\n",
    "    model.compile(optimizer=opt,\n",
    "              loss=K.binary_crossentropy,\n",
    "              metrics=[metrics.binary_accuracy, metrics.mean_squared_error])\n",
    "    \n",
    "    return model\n",
    "\n",
    "def predict_audio(model, audios):\n",
    "    predicted_audio_segments = model.predict(audios)\n",
    "    return predicted_audio_segments#.squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:59:37.811063Z",
     "start_time": "2019-06-25T19:59:37.550150Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# Write Files DWT domain\n",
    "fs = 16000\n",
    "count = 0\n",
    "\n",
    "df = pd.DataFrame({'id':np.nan, 'mse':np.nan, 'corr':np.nan}, index=[])\n",
    "from scipy.signal import butter, filtfilt\n",
    "def filter_audio(audio, freq=[3500/8000, 4300/8000]):\n",
    "    b, a = butter(5, freq, 'bandstop')\n",
    "    audio = filtfilt(b, a, audio)\n",
    "    return audio\n",
    "\n",
    "for i in range(len(test_generator)):\n",
    "    X, y = test_generator.__getitem__(i)\n",
    "    y_pred = predict_audio(model, X)\n",
    "        \n",
    "    for index in range(X.shape[0]):\n",
    "        X_merge = X[index, :, 0] if data_format == 'channels_last' else X[index, 0, :]\n",
    "        y_crowd = y[index, :, 0] if data_format == 'channels_last' else y[index, 0, :]\n",
    "        y_crowd_pred = y_pred[index, :, 0] if data_format == 'channels_last' else y_pred[index, 0, :]\n",
    "\n",
    "        arrays_dwt = [X_merge, y_crowd,  y_crowd_pred]#, y_voice, y_voice_pred]\n",
    "        wave_types = ['merge', 'crowd', 'recon']\n",
    "        arrays = [[]] * 3\n",
    "\n",
    "        for n, (k, w) in enumerate(zip(arrays_dwt, wave_types)):\n",
    "            out_wav = wavelet2time(k)\n",
    "            out_wav = filter_audio(out_wav)\n",
    "            arrays[n] = out_wav\n",
    "            librosa.output.write_wav(os.path.join(save_dir, \"{:04d}_{}.wav\".format(count, w)), out_wav, 16000)\n",
    "        \n",
    "        df.loc[count, 'id'] = \"{:04d}\".format(count)\n",
    "        df.loc[count, 'mse'] = np.mean((arrays[1] - arrays[2])**2)\n",
    "        df.loc[count, 'corr'] = np.corrcoef(arrays[1], arrays[2])[0, 1]\n",
    "        count = count + 1\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-25T19:59:45.378260Z",
     "start_time": "2019-06-25T19:59:45.189015Z"
    }
   },
   "outputs": [],
   "source": [
    "df_test_out = pd.concat([df_test, df], axis=1)\n",
    "df_test_out.to_csv(os.path.join(save_dir, 'test_results.csv'), index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optional Further Decomposition Using ICA\n",
    "Can sometimes enhance the final separation. Need to manually determine which source is the target."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2019-06-24T20:51:22.024138Z",
     "start_time": "2019-06-24T20:51:21.854439Z"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "from sklearn.decomposition import FastICA\n",
    "\n",
    "out_mixed = np.c_[arrays[0], arrays[2]]\n",
    "fica = FastICA(n_components=2)\n",
    "out_unmixed = fica.fit_transform(out_mixed)\n",
    "\n",
    "# Do some filtering before display\n",
    "out_unmixed_filtered = filter_audio(out_unmixed)\n",
    "\n",
    "print('Original')\n",
    "display.display(display.Audio(arrays[0], rate=fs))\n",
    "print('Source 1')\n",
    "display.display(display.Audio(out_unmixed_filtered[:, 0], rate=fs))\n",
    "print('Source 2')\n",
    "display.display(display.Audio(out_unmixed_filtered[:, 1], rate=fs))\n",
    "\"\"\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_mxnet_p36)",
   "language": "python",
   "name": "conda_mxnet_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}