Data_Pipeline_Part2(Autoencoder).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Going to recreate 16 ECoG channels:\n",
        "Autoencoder = PCA; Walk down from 129 to 16 and from 16 to 129.\n",
        "Overall results: the error does not converge in terms of training loss - a continuous oscillation of error, but the error is under 0.1 loss, which is ok.\n",
        "\n",
        "Encoder, Decoder used to reduce and demap target and estimate respectively."
      ],
      "metadata": {
        "id": "GgPWPGGED9FE"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YtkfalYOZ4vV",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0370a592-c723-4f64-e172-631a6befff07"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/gdrive\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install torch torchvision mat73 pymatreader matplotlib tensorboard mne joblib"
      ],
      "metadata": {
        "id": "XD_pC8O_aFeX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch.autograd import Variable\n",
        "import statistics\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import os"
      ],
      "metadata": {
        "id": "hGra3asu0_ui"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Autoencoder Class\n",
        "#class definition of model\n",
        "\n",
        "class Encoder(torch.nn.Module):\n",
        "  def __init__(self,input_size, bottleneck_size):\n",
        "    super(Encoder,self).__init__()\n",
        "    self.linear_list = torch.nn.ModuleList()\n",
        "    self.activation_function = torch.nn.ModuleList()\n",
        "    shift_len = 0\n",
        "\n",
        "    #expect the shiftlen to be the pwr of 2 higher than input_size (ex. 129 should be 256)\n",
        "    while (input_size >> shift_len) > 0:\n",
        "      shift_len += 1\n",
        "    shift_len -=2 # go from 64\n",
        "    output_size = 1 << shift_len\n",
        "    while output_size >= bottleneck_size:\n",
        "      print(f\"Encoder: Created Linear=({input_size,output_size})\")\n",
        "      self.linear_list.append(torch.nn.Linear(input_size,output_size))\n",
        "      self.activation_function.append(torch.nn.Tanh()) #restricts the output(including from last layer) to be between -1 to 1\n",
        "      input_size = output_size\n",
        "      output_size >>= 1\n",
        "    assert len(self.activation_function) == len(self.linear_list)\n",
        "\n",
        "  def forward(self,X):\n",
        "    reduced_data = X\n",
        "    for i in range(len(self.activation_function)):\n",
        "      reduced_data = self.linear_list[i](reduced_data)\n",
        "      reduced_data = self.activation_function[i](reduced_data)\n",
        "    return reduced_data\n",
        "\n",
        "class Decoder(torch.nn.Module):\n",
        "  def __init__(self,input_size, output_size):\n",
        "    super(Decoder,self).__init__()\n",
        "    self.linear_list = torch.nn.ModuleList()\n",
        "    self.activation_function = torch.nn.ModuleList()\n",
        "    shift_len = 0 \n",
        "\n",
        "    while (input_size << 1) < 128 : \n",
        "      print(f\"Decoder: Created Linear=({input_size,input_size << 1})\")\n",
        "      self.linear_list.append(torch.nn.Linear(input_size,input_size << 1))\n",
        "      self.activation_function.append(torch.nn.Tanh())\n",
        "      input_size = input_size << 1\n",
        "    \n",
        "    print(f\"Decoder: Created Linear=({input_size,output_size})\")\n",
        "    self.linear_list.append(torch.nn.Linear(input_size,output_size))\n",
        "    self.activation_function.append(torch.nn.Tanh())\n",
        "    assert len(self.activation_function) == len(self.linear_list)\n",
        "  def forward(self, X):\n",
        "    encoded_data = X\n",
        "    for i in range(len(self.activation_function)):\n",
        "      encoded_data = self.linear_list[i](encoded_data)\n",
        "      encoded_data = self.activation_function[i](encoded_data)\n",
        "    return encoded_data\n",
        "\n",
        "#expect auto-encoder to accept normalized(CRA_Sig(Y)) + dimensionally reduce\n",
        "#normalized(CRA_Sig(Y)) => normalized(CRA_Sig(Y))_hat\n",
        "class Autoencoder(torch.nn.Module):\n",
        "  def __init__(self,input_size, bottleneck):\n",
        "    super(Autoencoder,self).__init__()\n",
        "    self.encoder = Encoder(input_size,bottleneck)\n",
        "    self.decoder = Decoder(bottleneck,input_size)\n",
        "\n",
        "  def forward(self,X):\n",
        "    encoded = self.encoder(X)\n",
        "    decoded = self.decoder(encoded)\n",
        "    return decoded"
      ],
      "metadata": {
        "id": "XWuaAk0m3dPE",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#how we obtain the PCA/demapped results of the ECoG; only used after full training is done\n",
        "@torch.no_grad()\n",
        "def get_latent_space(model, X):\n",
        "  return model.encoder(X)\n",
        "\n",
        "@torch.no_grad()\n",
        "def decode_latent(model, Y):\n",
        "  return model.decoder(Y)"
      ],
      "metadata": {
        "id": "ZW8MywnE5pMQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Plotting Graphs-- General Methods\n",
        "def plot_loss(model_loss,epoch_id, model_type, error_type,img_dir, range_key):\n",
        "  assert len(model_loss) == 2\n",
        "  print(f'Plotting {model_type}_{epoch_id} {error_type} Loss in Range {range_key} ....')\n",
        "  x_vals = list(range(0,len(model_loss[0])))\n",
        "  plot_error(x_vals,model_loss,epoch_id,model_type,error_type,img_dir,range_key)\n",
        "\n",
        "def plot_error(x_vals, y_vals, epoch_i, model_type, error_type, images_dir, range_key):\n",
        "  plt.figure(figsize=(10,10))\n",
        "  counter = 0\n",
        "  dataset_type = [\"train\",\"validation\",\"test\"]\n",
        "  colors = ['r','b','y']\n",
        "  style = ['-','--','.']\n",
        "  for y_coors in y_vals:\n",
        "    assert len(y_coors) == len(x_vals)\n",
        "    plt.plot(x_vals,y_coors,colors[counter]+style[counter], label=dataset_type[counter])\n",
        "    counter += 1\n",
        "  plt.title(f'Autoencoder {model_type} Loss Curves until Epoch {epoch_i}')\n",
        "  plt.xlabel('Epoch (#)')\n",
        "  plt.ylabel(f'Error ({error_type})')\n",
        "  plt.legend(loc='best')\n",
        "  plt.savefig(f\"{images_dir}/{error_type}_{range_key}_curve_{epoch_i}_{model_type}.png\")\n",
        "  plt.show()\n",
        "\n",
        "def plot_test(model_loss,model_type,error_type,img_dir, range_key):\n",
        "  plt.figure(figsize=(10,10))\n",
        "  counts, bins = np.histogram(model_loss)\n",
        "  plt.hist(bins[:-1], bins, weights=counts)\n",
        "  plt.title(f'Frequency of RMSE Loss in Minibatches of {model_type} Set')\n",
        "  plt.savefig(f\"{img_dir}/{error_type}_{range_key}_histogram_{model_type}.png\")\n",
        "  plt.show()"
      ],
      "metadata": {
        "id": "STKZ3BOl8ADl",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Autoencoder Training/Testing Methods\n",
        "\n",
        "from enum import Enum\n",
        "class Mode(Enum):\n",
        "  TRAIN = 0\n",
        "  VALIDATION = 1\n",
        "  TEST = 2\n",
        "\n",
        "def one_epoch(data, model, cost_function, optimizer, mode, epoch,feature):\n",
        "  seq_count = 0\n",
        "  loss_list = []\n",
        "  assert mode != Mode.TEST\n",
        "  print(f'Running {feature}_Epoch {epoch}....',end='')\n",
        "  for seq in range(len(data)):\n",
        "    #seq.shape = 8x100xN\n",
        "    recreated_batch = model(data[seq])\n",
        "    assert recreated_batch.shape == data[seq].shape\n",
        "    train_loss = cost_function(recreated_batch,data[seq])\n",
        "    loss_list.append(train_loss.item())\n",
        "    if mode == Mode.TRAIN:\n",
        "        train_loss.backward()\n",
        "        optimizer.step()  \n",
        "        optimizer.zero_grad()\n",
        "    if mode == Mode.TRAIN:\n",
        "      type_str = \"Train\"\n",
        "    elif mode == Mode.VALIDATION:\n",
        "      type_str = \"Validation\"\n",
        "    print(f'\\tBatch Loss = {loss_list[-1]}')\n",
        "  total_loss = statistics.mean(loss_list)\n",
        "  print(f\"avg loss = {total_loss}\")\n",
        "  return total_loss\n",
        "\n",
        "def train_driver(epochs, inp_dataset, optimizer, cost_function, model, checkpoint_dir, feature, img_dir,range_key,bookmark=15):\n",
        "  train_loss = []\n",
        "  val_loss = []\n",
        "  inp_train_dataset = torch.split(inp_dataset[0],100,dim=1)\n",
        "  inp_val_dataset = torch.split(inp_dataset[1],100,dim=1)\n",
        "\n",
        "  for i in range(epochs + 1):\n",
        "    model.train(True)\n",
        "    train_loss.append(one_epoch(inp_train_dataset,model,cost_function,optimizer,Mode.TRAIN,i,feature))\n",
        "    model.train(False)\n",
        "    with torch.no_grad():\n",
        "      model.eval()\n",
        "      val_loss.append(one_epoch(inp_val_dataset, model, cost_function, optimizer, Mode.VALIDATION, i,feature))\n",
        "    if i > 0 and i % bookmark == 0:\n",
        "      print(f'Saving Autoencoder_{feature} Model....')\n",
        "      torch.save(model.state_dict(),f'{checkpoint_dir}/autoencoder_{range_key}_{feature}_model.pt')\n",
        "      plot_loss([train_loss,val_loss],i,feature,\"MSE\",img_dir,range_key)\n",
        "  return model\n",
        "  \n",
        "def test_driver(inp_dataset, cost_function, model, checkpoint_dir,feature,img_dir, range_key):\n",
        "  if model is None:\n",
        "    model = torch.load(f'{checkpoint_dir}/autoencoder_{range_key}_{feature}_model.pt') #load model from file\n",
        "  model.eval()\n",
        "  test_loss = []\n",
        "  inp_test_dataset = torch.split(inp_dataset,100,dim=1)\n",
        "  print(f'Testing {feature}....')\n",
        "  for batch in inp_test_dataset:\n",
        "    estimate = model(batch)\n",
        "    loss = cost_function(batch,estimate)\n",
        "    print(f'\\tBatch Loss = {loss.item()}')\n",
        "    test_loss.append(loss.item())\n",
        "  rmse_test_loss = np.sqrt(test_loss)\n",
        "  print(f'Avg RMSE for {feature}: {statistics.mean(rmse_test_loss)}')\n",
        "  plot_test(rmse_test_loss,feature,\"RMSE\",img_dir,range_key)"
      ],
      "metadata": {
        "id": "bdjz9LR67EuK",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Autoencoder Driver Methods\n",
        "def auto_subdriver(current_directory,epochs,dataset_locations,input_size,hidden_size,learning_rate, feature, range_key):\n",
        "  checkpt_dir = os.path.join(current_directory,f\"autoencoder_{range_key}\")\n",
        "  if not os.path.exists(checkpt_dir):\n",
        "    os.makedirs(checkpt_dir)\n",
        "  img_dir = os.path.join(checkpt_dir,f\"autoencoder_{range_key}_{feature}_results\")\n",
        "  if not os.path.exists(img_dir):\n",
        "    os.makedirs(img_dir)\n",
        "  \n",
        "  print(f'Driving Autoencoder_{feature}')\n",
        "  torch.autograd.set_detect_anomaly(True)\n",
        "  model = Autoencoder(input_size,hidden_size)\n",
        "  optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
        "  cost_function = torch.nn.MSELoss()\n",
        "\n",
        "  #get normalized datasets\n",
        "  print(dataset_locations[0])\n",
        "  train_dataset = torch.load(dataset_locations[0]).float()\n",
        "  print(f'Train Dim: {train_dataset.shape}')\n",
        "  validation_dataset = torch.load(dataset_locations[1]).float()\n",
        "  print(f'Validation Dim: {validation_dataset.shape}')\n",
        "  test_dataset = torch.load(dataset_locations[2]).float()\n",
        "  print(f'Test Dim: {test_dataset.shape}')\n",
        "\n",
        "  model = train_driver(epochs,[train_dataset,validation_dataset],optimizer,cost_function,model,checkpt_dir,feature,img_dir,range_key,bookmark=15)\n",
        "  with torch.no_grad():\n",
        "    test_driver(test_dataset,cost_function,model,checkpt_dir,feature,img_dir,range_key)\n",
        "\n",
        "\n",
        "def auto_driver(current_directory, epochs, data_key, range_key, specific_index = -1):\n",
        "  deeper_level = current_directory + f'/checkpoint3/normalized_ecog_{range_key}'\n",
        "  features = ['power','zero','mean','data']\n",
        "  dataset_type = ['train','validation','test']\n",
        "  locations = []\n",
        "  for feature_type in features:\n",
        "    filename_list = []\n",
        "    data_dir = f\"{deeper_level}/{feature_type}_dir/\"\n",
        "    for data_type in dataset_type:\n",
        "      filename_list.append(data_dir + f'{data_key}_{feature_type}_{data_type}.pt')\n",
        "    assert len(filename_list) == 3\n",
        "    locations.append(filename_list)\n",
        "  if specific_index >= 0:\n",
        "    assert specific_index < len(locations)\n",
        "    auto_subdriver(current_directory,epochs,locations[specific_index],129,16,1e-3,features[specific_index],range_key)\n",
        "  else:\n",
        "    for loc in range(len(locations)):\n",
        "      auto_subdriver(current_directory,epochs,locations[loc],129,16,1e-3,features[loc],range_key)"
      ],
      "metadata": {
        "id": "k4juEoSU1rPK",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "auto_driver(\"/content/gdrive/MyDrive/BCI_Project/Datasets\",150,'ECoG_Norm',\"01\",specific_index=0)"
      ],
      "metadata": {
        "id": "b9bFlPJhPDwW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Reduction Driver\n",
        "\"\"\"\n",
        "A method which stores reduced ecog data into files to reduce the potential ram usage when using get_latent_space\n",
        "\"\"\"\n",
        "import torch\n",
        "\n",
        "def generate_dataset_locations(current_directory,data_key,features,norm_key):\n",
        "  deeper_level = current_directory + f'/checkpoint3/normalized_ecog_{norm_key}'\n",
        "  locations = []\n",
        "  dataset_type = [\"train\",\"validation\",\"test\"]\n",
        "  for feature_type in features:\n",
        "    filename_list = []\n",
        "    data_dir = f\"{deeper_level}/{feature_type}_dir/\"\n",
        "    for data_type in dataset_type:\n",
        "      filename_list.append(data_dir + f'{data_key}_{norm_key}_{feature_type}_{data_type}.pt')\n",
        "    assert len(filename_list) == 3\n",
        "    locations.append(filename_list)\n",
        "  return locations\n",
        "\n",
        "def generate_model_locations(current_directory,norm_key,features):\n",
        "  model_list = []\n",
        "  deeper_level = current_directory + f'/autoencoder_{norm_key}'\n",
        "  for feature_type in features:\n",
        "    model_list.append(f'{deeper_level}/autoencoder_{norm_key}_{feature_type}_model.pt')\n",
        "  assert len(model_list)==4\n",
        "  return model_list\n",
        "\n",
        "def autoencoder_enc_dec(autoencoder_model,split_ecog_tensor,function):\n",
        "  trans_ecog_split = []\n",
        "  for item in split_ecog_tensor:\n",
        "    trans_ecog_split.append(function(autoencoder_model,item))\n",
        "  return torch.cat(trans_ecog_split,dim=1)\n",
        "\n",
        "\n",
        "def reduction_sub_driver(ecog_data, model_loc):\n",
        "  autoencode = Autoencoder(129,16)\n",
        "  autoencode.load_state_dict(torch.load(model_loc))\n",
        "  autoencode.eval()\n",
        "\n",
        "  for dataset_type in range(len(ecog_data)):\n",
        "    ecog_tensor = torch.load(ecog_data[dataset_type]).float()\n",
        "    print(f\"Shape of {ecog_data[dataset_type]}: {ecog_tensor.shape}\")\n",
        "    split_ecog_tensor = list(torch.split(ecog_tensor,100,dim=1))\n",
        "    reduced_ecog_tensor = autoencoder_enc_dec(autoencode,split_ecog_tensor,get_latent_space)\n",
        "    print(f\"\\tReduced Shape:{reduced_ecog_tensor.shape}\")\n",
        "    split_filename = ecog_data[dataset_type].split('/')\n",
        "    filename = \"reduced_\" + split_filename[-1].split('.')[0]\n",
        "    final_abs_filename = '/'.join(split_filename[:-1]) + '/'+filename+'.pt'\n",
        "    print(f\"\\tSaving tensor in {final_abs_filename}\")\n",
        "    torch.save(reduced_ecog_tensor,final_abs_filename)\n",
        "\n",
        "def reduction_caller(ecog_data_list, model_list, specific_index = -1):\n",
        "  #driver to iterate\n",
        "  assert len(ecog_data_list) == len(model_list)\n",
        "  if specific_index < 0:\n",
        "    for model_ind in range(len(model_list)):\n",
        "      reduction_sub_driver(ecog_data_list[model_ind],model_list[model_ind])\n",
        "  else:\n",
        "    assert specific_index < len(ecog_data_list) \n",
        "    reduction_sub_driver(ecog_data_list[specific_index], model_list[specific_index])\n",
        "\n",
        "def reduction_driver(current_directory,data_key,norm_key,specific_index_top = -1):\n",
        "  features = ['power','zero','mean','data']\n",
        "  locations = generate_dataset_locations(current_directory, data_key, features,norm_key)\n",
        "  print(locations)\n",
        "  model_list = generate_model_locations(current_directory,norm_key,features)\n",
        "  reduction_caller(locations,model_list,specific_index=specific_index_top)"
      ],
      "metadata": {
        "id": "tQ21sfQiWdw3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "reduction_driver(\"/content/gdrive/MyDrive/BCI_Project/Datasets\",'ECoG_Norm','01',specific_index_top=3)"
      ],
      "metadata": {
        "id": "mwvSs1Qld-oG"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}