Using_mne_and_braindecode.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "colab": {
      "name": "Using mne and braindecode.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qK8qZ72TiHAi"
      },
      "source": [
        "# install the mne and braindecode librarie "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "25mEbCwgN2yI",
        "outputId": "ba9a833c-70c2-44c5-ea94-afa359e69c21",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "source": [
        "# for Open-source Python software for exploring, visualizing, \n",
        "# and analyzing human neurophysiological data: MEG, EEG, sEEG, ECoG, and more.\n",
        "!pip install mne \n",
        "# A deep learning toolbox to decode raw time-domain EEG.\n",
        "!pip install https://github.com/TNTLFreiburg/braindecode/archive/master.zip \n"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: mne in /usr/local/lib/python3.7/dist-packages (0.19.2)\n",
            "Requirement already satisfied: scipy>=0.17.1 in /usr/local/lib/python3.7/dist-packages (from mne) (1.4.1)\n",
            "Requirement already satisfied: numpy>=1.11.3 in /usr/local/lib/python3.7/dist-packages (from mne) (1.21.6)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting https://github.com/TNTLFreiburg/braindecode/archive/master.zip\n",
            "  Using cached https://github.com/TNTLFreiburg/braindecode/archive/master.zip\n",
            "Requirement already satisfied: mne==0.19.2 in /usr/local/lib/python3.7/dist-packages (from Braindecode==0.4.85) (0.19.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from Braindecode==0.4.85) (1.21.6)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from Braindecode==0.4.85) (1.3.5)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from Braindecode==0.4.85) (1.4.1)\n",
            "Requirement already satisfied: resampy in /usr/local/lib/python3.7/dist-packages (from Braindecode==0.4.85) (0.2.2)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.7/dist-packages (from Braindecode==0.4.85) (3.2.2)\n",
            "Requirement already satisfied: h5py in /usr/local/lib/python3.7/dist-packages (from Braindecode==0.4.85) (3.1.0)\n",
            "Requirement already satisfied: cached-property in /usr/local/lib/python3.7/dist-packages (from h5py->Braindecode==0.4.85) (1.5.2)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->Braindecode==0.4.85) (2.8.2)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib->Braindecode==0.4.85) (0.11.0)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->Braindecode==0.4.85) (1.4.3)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib->Braindecode==0.4.85) (3.0.9)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from kiwisolver>=1.0.1->matplotlib->Braindecode==0.4.85) (4.1.1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.1->matplotlib->Braindecode==0.4.85) (1.15.0)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas->Braindecode==0.4.85) (2022.1)\n",
            "Requirement already satisfied: numba>=0.32 in /usr/local/lib/python3.7/dist-packages (from resampy->Braindecode==0.4.85) (0.51.2)\n",
            "Requirement already satisfied: llvmlite<0.35,>=0.34.0.dev0 in /usr/local/lib/python3.7/dist-packages (from numba>=0.32->resampy->Braindecode==0.4.85) (0.34.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from numba>=0.32->resampy->Braindecode==0.4.85) (57.4.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1k5PHDErNuyW"
      },
      "source": [
        "import mne\n",
        "from mne.io import concatenate_raws\n",
        "import numpy as np\n",
        "import torch as th\n",
        "from torch import nn\n",
        "import torch.nn.functional as F\n",
        "from braindecode.torch_ext.util import set_random_seeds\n",
        "from braindecode.models.util import to_dense_prediction_model\n",
        "from braindecode.datautil.iterators import CropsFromTrialsIterator\n",
        "from braindecode.torch_ext.util import np_to_var, var_to_np\n",
        "from braindecode.experiments.monitors import compute_preds_per_trial_from_crops\n",
        "from braindecode.torch_ext.optimizers import AdamW\n",
        "from braindecode.torch_ext.schedulers import ScheduledOptimizer, CosineAnnealing\n",
        "from braindecode.datautil.iterators import get_balanced_batches\n",
        "from braindecode.datautil.signal_target import SignalAndTarget\n",
        "from braindecode.models.shallow_fbcsp import ShallowFBCSPNet\n",
        "from numpy.random import RandomState\n"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HHRvl2Uwk0l0"
      },
      "source": [
        "# (open and close the left or right fist)               ---> 3,7,11 \n",
        "\n",
        "# (imagine opening and closing the left or right fist)  ---> 4,8,12 <--- I chose this\n",
        "\n",
        "# (open and close both fists or both feet)              ---> 5,9,13\n",
        "\n",
        "# (imagine opening and closing both fists or both feet) ---> 6,10,14\n",
        "\n",
        "# subject_id = [] # carefully cherry-picked to give nice results on such limited data :)\n",
        "# This will download the files if you don't have them yet,and then return the paths to the files.\n",
        "physionet_paths = [mne.datasets.eegbci.load_data(sub_id,[4,8,12]) for sub_id in range(1,80)]\n",
        "physionet_paths = np.concatenate(physionet_paths)\n",
        "\n",
        "# Load each of the files\n",
        "parts = [mne.io.read_raw_edf(path, preload=True,stim_channel='auto',verbose='WARNING')\n",
        "         for path in physionet_paths]\n",
        "\n",
        "# Concatenate them\n",
        "raw = concatenate_raws(parts)\n",
        "\n",
        "# Find the events in this dataset\n",
        "events, _ = mne.events_from_annotations(raw) #<--- no use the event_id for that is the \"_\"\n",
        "\n",
        "# Use only EEG channels\n",
        "eeg_channel_inds = mne.pick_types(raw.info, meg=False, eeg=True, stim=False, eog=False,\n",
        "                   exclude='bads')\n",
        "\n",
        "# Each annotation includes one of three codes (T0=1, T1=2, or T2=3):\n",
        "\n",
        "# T0 corresponds to rest\n",
        "# T1 corresponds to onset of motion (real or imagined) of\n",
        "# the left fist (in runs 3, 4, 7, 8, 11, and 12)\n",
        "# both fists (in runs 5, 6, 9, 10, 13, and 14)\n",
        "# T2 corresponds to onset of motion (real or imagined) of\n",
        "# the right fist (in runs 3, 4, 7, 8, 11, and 12)\n",
        "# both feet (in runs 5, 6, 9, 10, 13, and 14)\n",
        "\n",
        "# Extract trials, only using EEG channels\n",
        "epoched = mne.Epochs(raw, events, dict(left=2, right=3), tmin=1, tmax=4.1, proj=False, picks=eeg_channel_inds,\n",
        "                baseline=None, preload=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-IacSJpwNuyf"
      },
      "source": [
        "# Convert data from volt to millivolt\n",
        "# Pytorch expects float32 for input and int64 for labels.\n",
        "X = (epoched.get_data() * 1e6).astype(np.float32)\n",
        "y = (epoched.events[:,2] - 2).astype(np.int64) #2,3 -> 0,1 "
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XyvL6GV7Nuyk"
      },
      "source": [
        "epoched"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tFQqzhaVNuyp"
      },
      "source": [
        "epoched.events[:,2]-2 #---> the 0 represent the left and the 1 represent the right "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yG-9hEfaNuyt"
      },
      "source": [
        "train_set = SignalAndTarget(X[:2500], y=y[:2500]) \n",
        "valid_set = SignalAndTarget(X[2700:3000], y=y[2700:3000]) "
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ht-2g6yDNuyx"
      },
      "source": [
        "train_set.X.shape[2]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qYm4lPmSNuy1"
      },
      "source": [
        "train_set.X.shape[1]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fflfaBKKNuy6"
      },
      "source": [
        "# Set if you want to use GPU\n",
        "# You can also use torch.cuda.is_available() to determine if cuda is available on your machine.\n",
        "\n",
        "cuda = th.cuda.is_available()\n",
        "set_random_seeds(seed=20170629, cuda=cuda)\n",
        "\n",
        "# This will determine how many crops are processed in parallel\n",
        "input_time_length = train_set.X.shape[2]\n",
        "n_classes = 2\n",
        "in_chans = train_set.X.shape[1]\n",
        "\n",
        "# final_conv_length determines the size of the receptive field of the ConvNet\n",
        "model = ShallowFBCSPNet(in_chans=in_chans, n_classes=n_classes, \n",
        "                        input_time_length=input_time_length,final_conv_length = \"auto\").create_network()\n",
        "to_dense_prediction_model(model)\n",
        "\n",
        "if cuda:\n",
        "    model.cuda()"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4VITc7e7Yt4p"
      },
      "source": [
        "cuda"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kIQgRE5QjJTv"
      },
      "source": [
        "# determine output size\n",
        "test_input = np_to_var(np.ones((2, in_chans, input_time_length, 1), dtype=np.float32))\n",
        "if cuda:\n",
        "    test_input = test_input.cuda()\n",
        "out = model(test_input)\n",
        "n_preds_per_input = out.cpu().data.numpy().shape[2]\n",
        "print(\"{:d} predictions per input/trial\".format(n_preds_per_input))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jORSF-M0NuzB"
      },
      "source": [
        "# The iterator has the method get_batches, which can be used to get randomly shuffled training batches \n",
        "# with shuffle=True or ordered batches (i.e. first from trial 1, then from trial 2, etc.) with shuffle=False\n",
        "iterator = CropsFromTrialsIterator(batch_size=64,input_time_length=input_time_length,\n",
        "                                  n_preds_per_input=n_preds_per_input)"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RrWcMLl_NuzH"
      },
      "source": [
        "# Diferents values for the optimizer, uncoment and try ;) \n",
        "#optimizer = AdamW(model.parameters(), lr=2*0.01, weight_decay=0.01) # these are good values for the deep model\n",
        "#optimizer = AdamW(model.parameters(), lr=1*0.001, weight_decay=0.0001)\n",
        "optimizer = AdamW(model.parameters(), lr=0.06255 * 0.015, weight_decay=0)\n",
        "#optimizer = AdamW(model.parameters(), lr=1*0.01, weight_decay=0.5*0.001) # these are good values for the deep model\n",
        "#optimizer = AdamW(model.parameters(), lr=0.0625 * 0.01, weight_decay=0)\n",
        "\n",
        "# Need to determine number of batch passes per epoch for cosine annealing\n",
        "\n",
        "n_updates_per_epoch = len([None for b in iterator.get_batches(train_set, True)])\n",
        "n_epochs = 50\n",
        "scheduler = CosineAnnealing(n_epochs * n_updates_per_epoch)\n",
        "\n",
        "# schedule_weight_decay must be True for AdamW\n",
        "optimizer = ScheduledOptimizer(scheduler, optimizer, schedule_weight_decay=True)"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IWBe7A63NuzQ"
      },
      "source": [
        "#lets go to train !!! \n",
        "for i_epoch in range(50):\n",
        "    # Set model to training mode\n",
        "    model.train()\n",
        "    for batch_X, batch_y in iterator.get_batches(train_set, shuffle=True):\n",
        "        net_in = np_to_var(batch_X)\n",
        "        if cuda:\n",
        "            net_in = net_in.cuda()\n",
        "        net_target = np_to_var(batch_y)\n",
        "        if cuda:\n",
        "            net_target = net_target.cuda()\n",
        "        # Remove gradients of last backward pass from all parameters\n",
        "        optimizer.zero_grad()\n",
        "        outputs = model(net_in)\n",
        "        # Mean predictions across trial\n",
        "        # Note that this will give identical gradients to computing\n",
        "        # a per-prediction loss (at least for the combination of log softmax activation\n",
        "        # and negative log likelihood loss which we are using here)\n",
        "        outputs = th.mean(outputs, dim=2, keepdim=False)\n",
        "        loss = F.nll_loss(outputs, net_target)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "    # Print some statistics each epoch\n",
        "    model.eval()\n",
        "    print(\"Epoch {:d}\".format(i_epoch))\n",
        "    for setname, dataset in (('Train', train_set),('Valid', valid_set)):\n",
        "        # Collect all predictions and losses\n",
        "        all_preds = []\n",
        "        all_losses = []\n",
        "        batch_sizes = []\n",
        "        for batch_X, batch_y in iterator.get_batches(dataset, shuffle=False):\n",
        "            net_in = np_to_var(batch_X)\n",
        "            if cuda:\n",
        "                net_in = net_in.cuda()\n",
        "            net_target = np_to_var(batch_y)\n",
        "            if cuda:\n",
        "                net_target = net_target.cuda()\n",
        "            outputs = model(net_in)\n",
        "            all_preds.append(var_to_np(outputs))\n",
        "            outputs = th.mean(outputs, dim=2, keepdim=False)\n",
        "            loss = F.nll_loss(outputs, net_target)\n",
        "            loss = float(var_to_np(loss))\n",
        "            all_losses.append(loss)\n",
        "            batch_sizes.append(len(batch_X))\n",
        "        # Compute mean per-input loss\n",
        "        loss = np.mean(np.array(all_losses) * np.array(batch_sizes) /\n",
        "                       np.mean(batch_sizes))\n",
        "        print(\"{:6s} Loss: {:.5f}\".format(setname, loss))\n",
        "        # Assign the predictions to the trials\n",
        "        preds_per_trial = compute_preds_per_trial_from_crops(all_preds,\n",
        "                                                          input_time_length,\n",
        "                                                          dataset.X)\n",
        "        # preds per trial are now trials x classes x timesteps/predictions\n",
        "        # Now mean across timesteps for each trial to get per-trial predictions\n",
        "        meaned_preds_per_trial = np.array([np.mean(p, axis=1) for p in preds_per_trial])\n",
        "        predicted_labels = np.argmax(meaned_preds_per_trial, axis=1)\n",
        "        accuracy = np.mean(predicted_labels == dataset.y)\n",
        "        print(\"{:6s} Accuracy: {:.1f}%\".format(\n",
        "            setname, accuracy * 100))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "av1-vxmBNuzU"
      },
      "source": [
        "# now we will test the model\n",
        "test_set = SignalAndTarget(X[2500:2700], y=y[2500:2700])\n",
        "\n",
        "model.eval()\n",
        "# Collect all predictions and losses\n",
        "all_preds = []\n",
        "all_losses = []\n",
        "batch_sizes = []\n",
        "for batch_X, batch_y in iterator.get_batches(test_set, shuffle=False):\n",
        "    net_in = np_to_var(batch_X)\n",
        "    if cuda:\n",
        "        net_in = net_in.cuda()\n",
        "    net_target = np_to_var(batch_y)\n",
        "    if cuda:\n",
        "        net_target = net_target.cuda()\n",
        "    outputs = model(net_in)\n",
        "    all_preds.append(var_to_np(outputs))\n",
        "    outputs = th.mean(outputs, dim=2, keepdim=False)\n",
        "    loss = F.nll_loss(outputs, net_target)\n",
        "    loss = float(var_to_np(loss))\n",
        "    all_losses.append(loss)\n",
        "    batch_sizes.append(len(batch_X))\n",
        "# Compute mean per-input loss\n",
        "loss = np.mean(np.array(all_losses) * np.array(batch_sizes) /\n",
        "               np.mean(batch_sizes))\n",
        "print(\"Test Loss: {:.5f}\".format(loss))\n",
        "# Assign the predictions to the trials\n",
        "preds_per_trial = compute_preds_per_trial_from_crops(all_preds,\n",
        "                                                  input_time_length,\n",
        "                                                  test_set.X)\n",
        "# preds per trial are now trials x classes x timesteps/predictions\n",
        "# Now mean across timesteps for each trial to get per-trial predictions\n",
        "meaned_preds_per_trial = np.array([np.mean(p, axis=1) for p in preds_per_trial])\n",
        "predicted_labels = np.argmax(meaned_preds_per_trial, axis=1)\n",
        "accuracy = np.mean(predicted_labels == test_set.y)\n",
        "print(\"Test Accuracy: {:.1f}%\".format(accuracy * 100))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "50B5iPd_NuzY"
      },
      "source": [
        "# print the model arquitecture \n",
        "model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TR3a4QWkNuzv"
      },
      "source": [
        "# re training?\n",
        "### lets go!!! \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VvwXwq45Nuzw"
      },
      "source": [
        "# load the data\n",
        "physionet_paths_v2 = [mne.datasets.eegbci.load_data(sub_id,[4,8,12]) for sub_id in range(81,108)]\n",
        "physionet_paths_v2 = np.concatenate(physionet_paths_v2)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Yz_m3zMQNuz1"
      },
      "source": [
        "# Load each of the files\n",
        "parts_v2 = [mne.io.read_raw_edf(path, preload=True,stim_channel='auto',verbose='WARNING')\n",
        "         for path in physionet_paths_v2]\n",
        "\n",
        "# Concatenate them\n",
        "raw_v2 = concatenate_raws(parts)\n",
        "\n",
        "# Find the events in this dataset\n",
        "events_v2, _ = mne.events_from_annotations(raw_v2) #<--- no use the event_id for that is the \"_\"\n",
        "\n",
        "# Use only EEG channels\n",
        "eeg_channel_inds_v2 = mne.pick_types(raw.info, meg=False, eeg=True, stim=False, eog=False,\n",
        "                   exclude='bads')\n",
        "\n",
        "#Each annotation includes one of three codes (T0=1, T1=2, or T2=3):\n",
        "\n",
        "#T0 corresponds to rest\n",
        "#T1 corresponds to onset of motion (real or imagined) of\n",
        "#the left fist (in runs 3, 4, 7, 8, 11, and 12)\n",
        "#both fists (in runs 5, 6, 9, 10, 13, and 14)\n",
        "#T2 corresponds to onset of motion (real or imagined) of\n",
        "#the right fist (in runs 3, 4, 7, 8, 11, and 12)\n",
        "#both feet (in runs 5, 6, 9, 10, 13, and 14)\n",
        "\n",
        "# Extract trials, only using EEG channels\n",
        "epoched_2 = mne.Epochs(raw_v2, events_v2, dict(left=2, right=3), tmin=1, tmax=4.1, proj=False, picks=eeg_channel_inds_v2,\n",
        "                baseline=None, preload=True)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pTBtr0TJNuz6"
      },
      "source": [
        "# Convert data from volt to millivolt\n",
        "# Pytorch expects float32 for input and int64 for labels.\n",
        "X = (epoched_2.get_data() * 1e6).astype(np.float32)\n",
        "y = (epoched_2.events[:,2] - 2).astype(np.int64) #2,3 -> 0,1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mc-Lh1DyNuz9"
      },
      "source": [
        "epoched_2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hKWcuXbhNu0B"
      },
      "source": [
        "epoched_2.events[:,2]-2 #---> the 0 represent the left and the 1 represent the right "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B92kUwPdNu0F"
      },
      "source": [
        "train_set = SignalAndTarget(X[:2535], y=y[:2535]) # we have 3377 events and we will use 1636 for the training\n",
        "valid_set = SignalAndTarget(X[2535:3100], y=y[2535:3100]) #and for valid we will use 2450-1636 = 814 "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lO8epgpnNu0J"
      },
      "source": [
        "# Set if you want to use GPU\n",
        "# You can also use torch.cuda.is_available() to determine if cuda is available on your machine.\n",
        "cuda = torch.cuda.is_available()\n",
        "set_random_seeds(seed=20170629, cuda=cuda)\n",
        "\n",
        "# This will determine how many crops are processed in parallel\n",
        "input_time_length = train_set.X.shape[2]\n",
        "n_classes = 2\n",
        "in_chans = train_set.X.shape[1]\n",
        "# final_conv_length determines the size of the receptive field of the ConvNet\n",
        "model = ShallowFBCSPNet(in_chans=in_chans, n_classes=n_classes, input_time_length=input_time_length,\n",
        "                        final_conv_length=8).create_network()\n",
        "to_dense_prediction_model(model)\n",
        "\n",
        "if cuda:\n",
        "    model.cuda()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yW5uEvfkouqR"
      },
      "source": [
        "# determine output size\n",
        "test_input = np_to_var(np.ones((2, in_chans, input_time_length, 1), dtype=np.float32))\n",
        "if cuda:\n",
        "    test_input = test_input.cuda()\n",
        "out = model(test_input)\n",
        "n_preds_per_input = out.cpu().data.numpy().shape[2]\n",
        "print(\"{:d} predictions per input/trial\".format(n_preds_per_input))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dCtfll25pgpJ"
      },
      "source": [
        "iterator = CropsFromTrialsIterator(batch_size=64,input_time_length=input_time_length,\n",
        "                                  n_preds_per_input=n_preds_per_input)\n",
        "\n",
        "optimizer = AdamW(model.parameters(), lr=1*0.001, weight_decay=0.0001)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XxZckHlip4g1"
      },
      "source": [
        "# Need to determine number of batch passes per epoch for cosine annealing\n",
        "\n",
        "n_updates_per_epoch = len([None for b in iterator.get_batches(train_set, True)])\n",
        "n_epochs = 50\n",
        "scheduler = CosineAnnealing(n_epochs * n_updates_per_epoch)\n",
        "\n",
        "# schedule_weight_decay must be True for AdamW\n",
        "optimizer = ScheduledOptimizer(scheduler, optimizer, schedule_weight_decay=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VdHMZK9PqMZT"
      },
      "source": [
        "# lets go to train again xD!!!\n",
        "# but no with the same data  \n",
        "for i_epoch in range(50):\n",
        "    # Set model to training mode\n",
        "    model.train()\n",
        "    for batch_X, batch_y in iterator.get_batches(train_set, shuffle=True):\n",
        "        net_in = np_to_var(batch_X)\n",
        "        if cuda:\n",
        "            net_in = net_in.cuda()\n",
        "        net_target = np_to_var(batch_y)\n",
        "        if cuda:\n",
        "            net_target = net_target.cuda()\n",
        "        # Remove gradients of last backward pass from all parameters\n",
        "        optimizer.zero_grad()\n",
        "        outputs = model(net_in)\n",
        "        # Mean predictions across trial\n",
        "        # Note that this will give identical gradients to computing\n",
        "        # a per-prediction loss (at least for the combination of log softmax activation\n",
        "        # and negative log likelihood loss which we are using here)\n",
        "        outputs = th.mean(outputs, dim=2, keepdim=False)\n",
        "        loss = F.nll_loss(outputs, net_target)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "    # Print some statistics each epoch\n",
        "    model.eval()\n",
        "    print(\"Epoch {:d}\".format(i_epoch))\n",
        "    for setname, dataset in (('Train', train_set),('Valid', valid_set)):\n",
        "        # Collect all predictions and losses\n",
        "        all_preds = []\n",
        "        all_losses = []\n",
        "        batch_sizes = []\n",
        "        for batch_X, batch_y in iterator.get_batches(dataset, shuffle=False):\n",
        "            net_in = np_to_var(batch_X)\n",
        "            if cuda:\n",
        "                net_in = net_in.cuda()\n",
        "            net_target = np_to_var(batch_y)\n",
        "            if cuda:\n",
        "                net_target = net_target.cuda()\n",
        "            outputs = model(net_in)\n",
        "            all_preds.append(var_to_np(outputs))\n",
        "            outputs = th.mean(outputs, dim=2, keepdim=False)\n",
        "            loss = F.nll_loss(outputs, net_target)\n",
        "            loss = float(var_to_np(loss))\n",
        "            all_losses.append(loss)\n",
        "            batch_sizes.append(len(batch_X))\n",
        "        # Compute mean per-input loss\n",
        "        loss = np.mean(np.array(all_losses) * np.array(batch_sizes) /\n",
        "                       np.mean(batch_sizes))\n",
        "        print(\"{:6s} Loss: {:.5f}\".format(setname, loss))\n",
        "        # Assign the predictions to the trials\n",
        "        preds_per_trial = compute_preds_per_trial_from_crops(all_preds,\n",
        "                                                          input_time_length,\n",
        "                                                          dataset.X)\n",
        "        # preds per trial are now trials x classes x timesteps/predictions\n",
        "        # Now mean across timesteps for each trial to get per-trial predictions\n",
        "        meaned_preds_per_trial = np.array([np.mean(p, axis=1) for p in preds_per_trial])\n",
        "        predicted_labels = np.argmax(meaned_preds_per_trial, axis=1)\n",
        "        accuracy = np.mean(predicted_labels == dataset.y)\n",
        "        print(\"{:6s} Accuracy: {:.1f}%\".format(\n",
        "            setname, accuracy * 100))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nu_B8XZfqXU3"
      },
      "source": [
        "# :( no improve much"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "dxJKU0Wh2aZZ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}