RippleNet_training_unidirectional.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ubB9AcSmaT5z"
   },
   "source": [
    "# RippleNet_training_unidirectional\n",
    "Training of simple unidirectional recurrent neural network (RNN) implementation in `tensorflow.keras` using LSTM (long short-term memory) layers to identify time of occurence of sharp wave ripple (SPW-R) events in temporal LFP data.\n",
    "\n",
    "Author: Espen Hagen (<https://github.com/espenhgn>)\n",
    "\n",
    "LICENSE: <https://github.com/CINPLA/RippleNet/blob/master/LICENSE>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 139
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 36615,
     "status": "ok",
     "timestamp": 1587994505948,
     "user": {
      "displayName": "Espen Hagen",
      "photoUrl": "",
      "userId": "16098989268258426650"
     },
     "user_tz": -120
    },
    "id": "NiELVop4aT54",
    "outputId": "ec286503-2ab7-4057-d235-110beec4ee6a"
   },
   "outputs": [],
   "source": [
    "# allow running on Google Colab for training using Google Drive for file access\n",
    "try:\n",
    "    from google.colab import drive\n",
    "    drive.mount('/content/gdrive')\n",
    "    %cd gdrive/My\\ Drive/Colab\\ Notebooks/RippleNet\n",
    "    %tensorflow_version 2.x\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PMwOuBEXaT6H"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mUcZ1pzZaT6M"
   },
   "outputs": [],
   "source": [
    "# import modules\n",
    "import os\n",
    "import numpy as np\n",
    "import scipy.signal as ss\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.gridspec import GridSpec\n",
    "from matplotlib import colors\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras.utils import plot_model\n",
    "import ripplenet.models\n",
    "import h5py\n",
    "import pickle\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 544
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 48882,
     "status": "ok",
     "timestamp": 1587994518236,
     "user": {
      "displayName": "Espen Hagen",
      "photoUrl": "",
      "userId": "16098989268258426650"
     },
     "user_tz": -120
    },
    "id": "wOvpLagbaT6S",
    "outputId": "4db75a6c-8b1c-43c9-d169-eb9fb9a34745"
   },
   "outputs": [],
   "source": [
    "from tensorflow.python.client import device_lib\n",
    "print(device_lib.list_local_devices())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 48872,
     "status": "ok",
     "timestamp": 1587994518237,
     "user": {
      "displayName": "Espen Hagen",
      "photoUrl": "",
      "userId": "16098989268258426650"
     },
     "user_tz": -120
    },
    "id": "DwjTugQJaT6Y",
    "outputId": "0dc76e2d-8f7b-45dd-e817-228958eeb398"
   },
   "outputs": [],
   "source": [
    "print(tf.__version__)\n",
    "print(tf.test.gpu_device_name())\n",
    "print(\"Num GPUs Available: \", len(tf.config.experimental.list_physical_devices('GPU')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tz9QovhfaT6d"
   },
   "outputs": [],
   "source": [
    "# set random seeds with some additional environment variables to ensure deterministic output\n",
    "random_seed = 123\n",
    "os.environ['TF_DETERMINISTIC_OPS'] = '1'\n",
    "os.environ['PYTHONHASHSEED']=str(random_seed)\n",
    "random.seed(random_seed)\n",
    "np.random.seed(random_seed)\n",
    "tf.random.set_seed(random_seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "z1uNRJlXHC3-"
   },
   "outputs": [],
   "source": [
    "# select dataset (may have generated different sets.)\n",
    "dataset_index = 0 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OzZCG0OdaT6j"
   },
   "source": [
    "# Load training/validation data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Lty8L4jbkKSM"
   },
   "outputs": [],
   "source": [
    "# select species for training/validation data (mouse, rat or both)\n",
    "mouse = True\n",
    "rat = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "S4Izedhc3s_g"
   },
   "outputs": [],
   "source": [
    "# output destination\n",
    "output_folder = 'trained_networks'\n",
    "if not os.path.isdir(output_folder):\n",
    "    os.mkdir(output_folder)\n",
    "    \n",
    "# prefix for trained network files (training loss/MSE, weights, `best' weights)\n",
    "rnn_prefix = 'ripplenet_unidirectional'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "e2PaOdJzC0vb"
   },
   "outputs": [],
   "source": [
    "if mouse:\n",
    "    # training and validation files\n",
    "    f_name_train = 'train_{:02}.h5'\n",
    "    f_name_val = 'validation_{:02}.h5'\n",
    "\n",
    "    # training data\n",
    "    f = h5py.File(os.path.join('data', \n",
    "                               f_name_train.format(dataset_index)), \n",
    "                  'r')\n",
    "    X_train = np.expand_dims(f['X0'][:], -1)\n",
    "    Y_train = f['Y'][:]\n",
    "    f.close()\n",
    "\n",
    "    # validation data\n",
    "    f = h5py.File(os.path.join('data', \n",
    "                               f_name_val.format(dataset_index)), \n",
    "                  'r')\n",
    "    X_val = np.expand_dims(f['X0'][:], -1)\n",
    "    Y_val = f['Y'][:]\n",
    "    f.close()\n",
    "\n",
    "    # load some data for plotting\n",
    "    f = h5py.File(os.path.join('data', \n",
    "                               f_name_val.format(dataset_index)), 'r')\n",
    "    X0 = f['X0'][:]\n",
    "    X1 = f['X1'][:]\n",
    "    S = f['S'][:]\n",
    "    Y = f['Y'][:]\n",
    "    S_freqs = f['S_freqs'][:]\n",
    "    f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uwZ4kvu11Va0"
   },
   "outputs": [],
   "source": [
    "# Add rat training/validation data to sets\n",
    "if rat and mouse:\n",
    "    # rat \n",
    "    f_name_train = 'train_tingley_{:02}.h5'\n",
    "    f_name_val = 'validation_tingley_{:02}.h5'\n",
    "\n",
    "    # training data\n",
    "    f = h5py.File(os.path.join('data', \n",
    "                            f_name_train.format(dataset_index)), \n",
    "                'r')\n",
    "    X_train = np.concatenate((X_train, np.expand_dims(f['X0'][:], -1)))\n",
    "    Y_train = np.concatenate((Y_train, f['Y'][:]))\n",
    "    f.close()\n",
    "\n",
    "    # validation data\n",
    "    f = h5py.File(os.path.join('data', \n",
    "                            f_name_val.format(dataset_index)), \n",
    "                'r')\n",
    "    X_val = np.concatenate((X_val, np.expand_dims(f['X0'][:], -1)))\n",
    "    Y_val = np.concatenate((Y_val, f['Y'][:]))\n",
    "    f.close()\n",
    "\n",
    "    # load some data for plotting\n",
    "    f = h5py.File(os.path.join('data', \n",
    "                            f_name_val.format(dataset_index)), 'r')\n",
    "    X0 = np.concatenate((X0, f['X0'][:]))\n",
    "    X1 = np.concatenate((X1, f['X1'][:]))\n",
    "    S = np.concatenate((S, f['S'][:]))\n",
    "    Y = np.concatenate((Y, f['Y'][:]))\n",
    "    f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "q10NkE38kKSa"
   },
   "outputs": [],
   "source": [
    "if rat and not mouse:\n",
    "    # rat \n",
    "    f_name_train = 'train_tingley_{:02}.h5'\n",
    "    f_name_val = 'validation_tingley_{:02}.h5'\n",
    "\n",
    "    # training data\n",
    "    f = h5py.File(os.path.join('..', 'data', \n",
    "                               f_name_train.format(dataset_index)), \n",
    "                  'r')\n",
    "    X_train = np.expand_dims(f['X0'][:], -1)\n",
    "    Y_train = f['Y'][:]\n",
    "    f.close()\n",
    "\n",
    "    # validation data\n",
    "    f = h5py.File(os.path.join('data', \n",
    "                               f_name_val.format(dataset_index)), \n",
    "                  'r')\n",
    "    X_val = np.expand_dims(f['X0'][:], -1)\n",
    "    Y_val = f['Y'][:]\n",
    "    f.close()\n",
    "\n",
    "    # load some data for plotting\n",
    "    f = h5py.File(os.path.join('data', \n",
    "                               f_name_val.format(dataset_index)), 'r')\n",
    "    X0 = f['X0'][:]\n",
    "    X1 = f['X1'][:]\n",
    "    S = f['S'][:]\n",
    "    Y = f['Y'][:]\n",
    "    S_freqs = f['S_freqs'][:]\n",
    "    f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Dl-v81WHkKSe"
   },
   "outputs": [],
   "source": [
    "# needed parameters\n",
    "Fs = 1250 # Hz, sampling freq\n",
    "time = np.arange(X0.shape[1]) / Fs\n",
    "\n",
    "# center raw data\n",
    "X0 = (X0.T - X0.mean(axis=-1)).T\n",
    "\n",
    "# total number of samples\n",
    "n_samples = X0.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 730
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 68251,
     "status": "ok",
     "timestamp": 1587994537657,
     "user": {
      "displayName": "Espen Hagen",
      "photoUrl": "",
      "userId": "16098989268258426650"
     },
     "user_tz": -120
    },
    "id": "VszoOq-WaT60",
    "outputId": "342d026d-a0c6-467d-8e2e-9d740273b8db"
   },
   "outputs": [],
   "source": [
    "# plot all labels and raw data matrices\n",
    "fig, axes = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(12, 12))\n",
    "axes[0].pcolormesh(time, np.arange(n_samples), Y[:, :, 0], shading='auto')\n",
    "axes[0].set_ylabel('#')\n",
    "axes[0].set_title('labels (y)')\n",
    "axes[1].pcolormesh(time, np.arange(n_samples), X0, vmin=-X0.std()*3, vmax=X0.std()*3, shading='auto')\n",
    "axes[1].set_ylabel('#')\n",
    "axes[1].set_xlabel('t (s)')\n",
    "axes[1].set_title('raw data (X)')\n",
    "for ax in axes:\n",
    "    ax.axis(ax.axis('tight'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 71156,
     "status": "ok",
     "timestamp": 1587994540569,
     "user": {
      "displayName": "Espen Hagen",
      "photoUrl": "",
      "userId": "16098989268258426650"
     },
     "user_tz": -120
    },
    "id": "xqGT2ICzaT64",
    "outputId": "7cc20c9d-23af-4035-cda3-a577e6e20d90"
   },
   "outputs": [],
   "source": [
    "# plot wavelet spectrograms vs. labels and raw data for some samples\n",
    "for i in range(5):\n",
    "    gs = GridSpec(2, 1)\n",
    "    fig = plt.figure(figsize=(12, 6))\n",
    "    ax0 = fig.add_subplot(gs[0, 0])\n",
    "    ax0.plot(time, X0[i, ], label='$X(t)$')\n",
    "    ax0.plot(time, X1[i, ], label=r'$\\phi_\\mathrm{bp}(t)$')\n",
    "    ax0.plot(time, Y[i, :, 0], label='label ($y$)' )\n",
    "    ax0.legend(ncol=2)\n",
    "    ax0.axis(ax0.axis('tight'))\n",
    "    ax0.set_title('label, raw data and spectrograms')\n",
    "    plt.setp(ax0.get_xticklabels(), visible=False)\n",
    "    \n",
    "    ax1 = fig.add_subplot(gs[1:, 0], sharex=ax0)\n",
    "    vmin, vmax = np.exp(np.percentile(np.log(S), [1, 99]))\n",
    "    im = ax1.pcolormesh(time, S_freqs, S[i, ].T, norm=colors.LogNorm(vmin=vmin, vmax=vmax),\n",
    "                        cmap='inferno', \n",
    "                        shading='auto')\n",
    "    ax1.axis(ax1.axis('tight'))\n",
    "    ax1.set_ylabel('$f$ (Hz)')\n",
    "    ax1.set_xlabel('$t$ (s)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GF-Om-cnaT6-"
   },
   "source": [
    "# Set up recurrent neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3VeJA7jpkKSt"
   },
   "outputs": [],
   "source": [
    "model = ripplenet.models.get_unidirectional_LSTM_model(input_shape=(None, X_train.shape[2]), \n",
    "                                                       layer_sizes=[20, 10, 20, 20],\n",
    "                                                       seed=random_seed+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 697
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 71144,
     "status": "ok",
     "timestamp": 1587994540570,
     "user": {
      "displayName": "Espen Hagen",
      "photoUrl": "",
      "userId": "16098989268258426650"
     },
     "user_tz": -120
    },
    "id": "nqH3VcApaT7J",
    "outputId": "584f845c-17d3-49f7-e025-96e98596da43"
   },
   "outputs": [],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4IE_8BYTRz-4"
   },
   "outputs": [],
   "source": [
    "# plot_model(model, show_shapes=True, expand_nested=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Wu5rxBaaaT7N"
   },
   "outputs": [],
   "source": [
    "# callback for model checkpoints when validation MSE improves\n",
    "filepath = os.path.join(output_folder, '{}_best_random_seed{}.h5'.format(rnn_prefix, random_seed))\n",
    "checkpoint_best = keras.callbacks.ModelCheckpoint(filepath, monitor='val_mse', \n",
    "                                             verbose=1, save_best_only=True, \n",
    "                                             mode='min')\n",
    "# MSE/loss callback to .csv\n",
    "callback_hist = keras.callbacks.CSVLogger(os.path.join(output_folder, \n",
    "                                                       '{}_history_random_seed{}.csv'.format(rnn_prefix, \n",
    "                                                                                             random_seed)))\n",
    "callbacks_list = [checkpoint_best, callback_hist]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1533519,
     "status": "ok",
     "timestamp": 1587996002963,
     "user": {
      "displayName": "Espen Hagen",
      "photoUrl": "",
      "userId": "16098989268258426650"
     },
     "user_tz": -120
    },
    "id": "ktXzlvIeaT7S",
    "outputId": "a8c1bd23-db96-4e61-ccaa-f590430a18e5"
   },
   "outputs": [],
   "source": [
    "# train model\n",
    "history = model.fit(X_train, Y_train, \n",
    "                    batch_size=20, \n",
    "                    epochs=50, \n",
    "                    callbacks=callbacks_list,\n",
    "                    validation_data=(X_val, Y_val))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IOaPE2duWVjl"
   },
   "outputs": [],
   "source": [
    "# save history to a pickle so we can load it later\n",
    "with open(os.path.join(output_folder, '{}_history_random_seed{}.pkl'.format(rnn_prefix, random_seed)\n",
    "                      ), 'wb') as f:\n",
    "    pickle.dump(history.history, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 747
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1535187,
     "status": "ok",
     "timestamp": 1587996004646,
     "user": {
      "displayName": "Espen Hagen",
      "photoUrl": "",
      "userId": "16098989268258426650"
     },
     "user_tz": -120
    },
    "id": "jijEvsoEaT7X",
    "outputId": "d032fabd-5d6e-4fd9-85ca-0279b9cae6be"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12, 12))\n",
    "plt.semilogy(history.history['loss'], '-o', label='loss')\n",
    "plt.semilogy(history.history['val_loss'], '-o', label='val_loss')\n",
    "plt.semilogy(history.history['mse'], '-o', label='mse')\n",
    "plt.semilogy(history.history['val_mse'], '-o', label='val_mse')\n",
    "plt.legend()\n",
    "plt.xlabel('epochs')\n",
    "plt.ylabel('MSE')\n",
    "plt.title('training/validation MSE')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xLF11w-saT7a"
   },
   "outputs": [],
   "source": [
    "# Save the trained model\n",
    "model.save(os.path.join(output_folder, '{}_random_seed{}.h5'.format(rnn_prefix, random_seed)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PcjPv13i2PEk"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "RippleNet_training_unidirectional.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}