s3_transfer_learning.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n",
        "# Transfer learning for EEG\n",
        "\n",
        "This example shows how to train a neural network with supervision on TUH [2]\n",
        "EEG data and transfer the model to NMT [3] EEG dataset. We follow the approach of [1]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Authors: MJ Bayazi <mj.darvishi92@gmail.com>\n",
        "#\n",
        "# License: BSD (3-clause)\n",
        "\n",
        "from braindecode.datautil.serialization import  load_concat_dataset\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "\n",
        "random_state = 2024\n",
        "n_jobs = 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Loading and preprocessing the dataset\n",
        "\n",
        "### Load  and save the raw recordings\n",
        "\n",
        "Here we assume you already load and preprocess the raw recordings for both TUAB and NMT datasetsand saved the file in `TUAB_path' and 'NMT_path' respectively. To read more see this notebook [here](https://braindecode.org/stable/auto_examples/applied_examples/plot_tuh_eeg_corpus.html) \n",
        "\n",
        "### Load the preprocessed data\n",
        "\n",
        "Now, we load a few recordings from the TUAB dataset. Running\n",
        "this example with more recordings should yield better representations and\n",
        "downstream classification performance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {},
      "outputs": [],
      "source": [
        "Exp_Path = 'YOUR_PATH'  # specify the path to the experiment folder\n",
        "TUAB_pp_path = Exp_Path + '/tuab/tuab_pp'\n",
        "NMT_pp_path = Exp_Path + '/NMT/nmt_pp'\n",
        "RESULTS_path = Exp_Path + '/results/'\n",
        "N_JOBS = 40  # specify the number of jobs for loading and windowing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Load NMT dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {},
      "outputs": [],
      "source": [
        "import mne\n",
        "mne.set_log_level('ERROR')\n",
        "\n",
        "NMT_ds = load_concat_dataset(NMT_pp_path, preload=False,\n",
        "                            target_name=['pathological','age','gender'] ,#)\n",
        "                            # ids_to_load=range(200)\n",
        "                            )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {},
      "outputs": [],
      "source": [
        "# split based on train split from dataset\n",
        "train_set = NMT_ds.split('train')['True']\n",
        "test_set = NMT_ds.split('train')['False']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "target is being set to pathological clf\n"
          ]
        }
      ],
      "source": [
        "import pandas as pd\n",
        "print(\"target is being set to pathological clf\")\n",
        "target = NMT_ds.description['pathological'].astype(int)\n",
        "for d, y in zip(NMT_ds.datasets, target):\n",
        "    d.description['pathological'] = y\n",
        "    d.target_name = 'pathological'\n",
        "    d.target = d.description[d.target_name]\n",
        "NMT_ds.set_description(pd.DataFrame([d.description for d in NMT_ds.datasets]), overwrite=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating the model\n",
        "\n",
        "We can now create the deep learning model. In this tutorial, we use DeepNet introduced in [4]."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from braindecode.models import Deep4Net\n",
        "\n",
        "n_chans = 21\n",
        "n_classes = 2\n",
        "input_window_samples = 6000\n",
        "drop_prob = 0.5\n",
        "cuda = True # Set to False if you don't have a GPU\n",
        "n_start_chans = 25\n",
        "final_conv_length = 1\n",
        "n_chan_factor = 2\n",
        "stride_before_pool = True\n",
        "# input_window_samples =6000\n",
        "model = Deep4Net(\n",
        "            n_chans, n_classes,\n",
        "            n_filters_time=n_start_chans,\n",
        "            n_filters_spat=n_start_chans,\n",
        "            input_window_samples=input_window_samples,\n",
        "            n_filters_2=int(n_start_chans * n_chan_factor),\n",
        "            n_filters_3=int(n_start_chans * (n_chan_factor ** 2.0)),\n",
        "            n_filters_4=int(n_start_chans * (n_chan_factor ** 3.0)),\n",
        "            final_conv_length=final_conv_length,\n",
        "            stride_before_pool=stride_before_pool,\n",
        "            drop_prob=drop_prob)\n",
        "            # Send model to GPU\n",
        "if cuda:\n",
        "    model.cuda()\n",
        "\n",
        "from braindecode.models.util import to_dense_prediction_model, get_output_shape\n",
        "\n",
        "to_dense_prediction_model(model)\n",
        "\n",
        "n_preds_per_input = get_output_shape(model, n_chans, input_window_samples)[2]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Extracting windows\n",
        "\n",
        "We extract 60-s windows to be used in both datasets. We use a window size of 6000 samples.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {},
      "outputs": [],
      "source": [
        "from IPython.utils import io\n",
        "from braindecode.datautil.windowers import create_fixed_length_windows\n",
        "\n",
        "with io.capture_output() as captured:\n",
        "        window_train_set = create_fixed_length_windows(train_set, \n",
        "                                                    start_offset_samples=0,\n",
        "                                                    stop_offset_samples=None,\n",
        "                                                    preload=True,\n",
        "                                                    window_size_samples=input_window_samples,\n",
        "                                                    window_stride_samples=n_preds_per_input,\n",
        "                                                    drop_last_window=True,)\n",
        "        \n",
        "with io.capture_output() as captured:\n",
        "        window_test_set = create_fixed_length_windows(test_set,\n",
        "                                                    start_offset_samples=0,\n",
        "                                                    stop_offset_samples=None,preload=False,\n",
        "                                                    window_size_samples=input_window_samples,\n",
        "                                                    window_stride_samples=n_preds_per_input,\n",
        "                                                    drop_last_window=False,)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## defining the classifier\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "classifier initialized\n",
            "Number of parameters =  277052\n"
          ]
        }
      ],
      "source": [
        "from braindecode import EEGClassifier\n",
        "from braindecode.training.losses import CroppedLoss\n",
        "import torch\n",
        "from skorch.callbacks import LRScheduler\n",
        "from skorch.helper import predefined_split\n",
        "\n",
        "weight_decay = 0.5 * 0.001\n",
        "lr = 0.0625 * 0.01\n",
        "n_epochs = 10\n",
        "batch_size = 64\n",
        "\n",
        "clf = EEGClassifier(\n",
        "                model,\n",
        "                cropped=True,\n",
        "                classes=[0, 1],\n",
        "                criterion=CroppedLoss,\n",
        "                criterion__loss_function=torch.nn.functional.nll_loss,\n",
        "                optimizer=torch.optim.AdamW,\n",
        "                train_split=predefined_split(window_test_set), \n",
        "                optimizer__lr=lr,\n",
        "                optimizer__weight_decay=weight_decay,\n",
        "                iterator_train__shuffle=True,\n",
        "                batch_size=batch_size,\n",
        "                callbacks=[\n",
        "                    \"accuracy\", \"balanced_accuracy\",\"f1\",(\"lr_scheduler\", LRScheduler('CosineAnnealingLR', T_max=n_epochs - 1)),\n",
        "                    ],\n",
        "                device='cuda' if cuda else 'cpu',\n",
        "                )\n",
        "\n",
        "clf.initialize() # This is important!\n",
        "print('classifier initialized')\n",
        "print(\"Number of parameters = \", sum(p.numel() for p in model.parameters() if p.requires_grad))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Training\n",
        "\n",
        "We trained our model in S2 on the TUAB. We use similar\n",
        "hyperparameters as in [1]_, but reduce the number of epochs and\n",
        "increase the learning rate to account for the smaller setting of\n",
        "this example.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## loading the pre trained model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "pre-trained model loaded using pytorch\n"
          ]
        }
      ],
      "source": [
        "load_path = RESULTS_path + 'state_dict_2024.pt'\n",
        "state_dicts = torch.load(load_path) \n",
        "model.load_state_dict(state_dicts, strict= False)\n",
        "print('pre-trained model loaded using pytorch')\n",
        "\n",
        "## freeze layers ##\n",
        "freez = False\n",
        "if freez:\n",
        "    for ii, (name, param) in enumerate(model.named_parameters()):\n",
        "        # if 'temporal_block_0' in name or 'temporal_block_1' in name or 'temporal_block_2' in name or 'temporal_block_3' in name: # or 'temporal_block_5' in name or 'conv_classifier' in name:\n",
        "        if not 'conv_classifier' in name:\n",
        "            param.requires_grad = False\n",
        "            print('param:', name, param.requires_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## fine tuning the model on NMT dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "classifier initialized\n",
            "Number of parameters =  277052\n"
          ]
        }
      ],
      "source": [
        "from braindecode import EEGClassifier\n",
        "from braindecode.training.losses import CroppedLoss\n",
        "import torch\n",
        "from skorch.callbacks import LRScheduler\n",
        "from skorch.helper import predefined_split\n",
        "\n",
        "weight_decay = 0.5 * 0.001\n",
        "lr = 0.0625 * 0.01\n",
        "n_epochs = 10\n",
        "batch_size = 64\n",
        "\n",
        "clf = EEGClassifier(\n",
        "                model,\n",
        "                cropped=True,\n",
        "                classes=[0, 1],\n",
        "                criterion=CroppedLoss,\n",
        "                criterion__loss_function=torch.nn.functional.nll_loss,\n",
        "                optimizer=torch.optim.AdamW,\n",
        "                train_split=predefined_split(window_test_set), \n",
        "                optimizer__lr=lr,\n",
        "                optimizer__weight_decay=weight_decay,\n",
        "                iterator_train__shuffle=True,\n",
        "                batch_size=batch_size,\n",
        "                callbacks=[\n",
        "                    \"accuracy\", \"balanced_accuracy\",\"f1\",(\"lr_scheduler\", LRScheduler('CosineAnnealingLR', T_max=n_epochs - 1)),\n",
        "                    ],\n",
        "                device='cuda' if cuda else 'cpu',\n",
        "                )\n",
        "\n",
        "clf.initialize() # This is important!\n",
        "print('classifier initialized')\n",
        "print(\"Number of parameters = \", sum(p.numel() for p in model.parameters() if p.requires_grad))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Re-initializing module.\n",
            "Re-initializing criterion because the following parameters were re-set: loss_function.\n",
            "Re-initializing optimizer.\n",
            "  epoch    train_accuracy    train_balanced_accuracy    train_f1    valid_accuracy    valid_balanced_accuracy    valid_f1    valid_loss      lr     dur\n",
            "-------  ----------------  -------------------------  ----------  ----------------  -------------------------  ----------  ------------  ------  ------\n",
            "      1            \u001b[36m0.3333\u001b[0m                     \u001b[32m0.5000\u001b[0m      \u001b[35m0.5000\u001b[0m            \u001b[31m0.5000\u001b[0m                     \u001b[94m0.5000\u001b[0m      \u001b[36m0.6667\u001b[0m       \u001b[32m27.6027\u001b[0m  0.0006  0.0805\n",
            "      2            0.3333                     0.5000      0.5000            0.5000                     0.5000      0.6667       27.6027  0.0006  0.0710\n",
            "      3            0.3333                     0.5000      0.5000            0.5000                     0.5000      0.6667       27.6027  0.0006  0.0879\n",
            "      4            0.3333                     0.5000      0.5000            0.5000                     0.5000      0.6667       27.6027  0.0005  0.0714\n",
            "      5            0.3333                     0.5000      0.5000            0.5000                     0.5000      0.6667       27.6027  0.0004  0.0697\n",
            "      6            0.3333                     0.5000      0.5000            0.5000                     0.5000      0.6667       27.6027  0.0003  0.0735\n",
            "      7            0.3333                     0.5000      0.5000            0.5000                     0.5000      0.6667       27.6027  0.0002  0.0694\n",
            "      8            0.3333                     0.5000      0.5000            0.5000                     0.5000      0.6667       27.6027  0.0001  0.0704\n",
            "      9            0.3333                     0.5000      0.5000            0.5000                     0.5000      0.6667       27.6027  0.0000  0.0724\n",
            "     10            0.3333                     0.5000      0.5000            0.5000                     0.5000      0.6667       27.6027  0.0000  0.0715\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "<class 'braindecode.classifier.EEGClassifier'>[initialized](\n",
              "  module_=============================================================================================================================================\n",
              "  Layer (type (var_name):depth-idx)        Input Shape               Output Shape              Param #                   Kernel Shape\n",
              "  ============================================================================================================================================\n",
              "  Deep4Net (Deep4Net)                      [1, 21, 6000]             [1, 2, 5794]              --                        --\n",
              "  ├─Ensure4d (ensuredims): 1-1             [1, 21, 6000]             [1, 21, 6000, 1]          --                        --\n",
              "  ├─Rearrange (dimshuffle): 1-2            [1, 21, 6000, 1]          [1, 1, 6000, 21]          --                        --\n",
              "  ├─CombinedConv (conv_time_spat): 1-3     [1, 1, 6000, 21]          [1, 25, 5991, 1]          13,400                    --\n",
              "  ├─BatchNorm2d (bnorm): 1-4               [1, 25, 5991, 1]          [1, 25, 5991, 1]          50                        --\n",
              "  ├─Expression (conv_nonlin): 1-5          [1, 25, 5991, 1]          [1, 25, 5991, 1]          --                        --\n",
              "  ├─MaxPool2d (pool): 1-6                  [1, 25, 5991, 1]          [1, 25, 5989, 1]          --                        [3, 1]\n",
              "  ├─Expression (pool_nonlin): 1-7          [1, 25, 5989, 1]          [1, 25, 5989, 1]          --                        --\n",
              "  ├─Dropout (drop_2): 1-8                  [1, 25, 5989, 1]          [1, 25, 5989, 1]          --                        --\n",
              "  ├─Conv2d (conv_2): 1-9                   [1, 25, 5989, 1]          [1, 50, 5980, 1]          12,500                    [10, 1]\n",
              "  ├─BatchNorm2d (bnorm_2): 1-10            [1, 50, 5980, 1]          [1, 50, 5980, 1]          100                       --\n",
              "  ├─Expression (nonlin_2): 1-11            [1, 50, 5980, 1]          [1, 50, 5980, 1]          --                        --\n",
              "  ├─MaxPool2d (pool_2): 1-12               [1, 50, 5980, 1]          [1, 50, 5974, 1]          --                        [3, 1]\n",
              "  ├─Expression (pool_nonlin_2): 1-13       [1, 50, 5974, 1]          [1, 50, 5974, 1]          --                        --\n",
              "  ├─Dropout (drop_3): 1-14                 [1, 50, 5974, 1]          [1, 50, 5974, 1]          --                        --\n",
              "  ├─Conv2d (conv_3): 1-15                  [1, 50, 5974, 1]          [1, 100, 5947, 1]         50,000                    [10, 1]\n",
              "  ├─BatchNorm2d (bnorm_3): 1-16            [1, 100, 5947, 1]         [1, 100, 5947, 1]         200                       --\n",
              "  ├─Expression (nonlin_3): 1-17            [1, 100, 5947, 1]         [1, 100, 5947, 1]         --                        --\n",
              "  ├─MaxPool2d (pool_3): 1-18               [1, 100, 5947, 1]         [1, 100, 5929, 1]         --                        [3, 1]\n",
              "  ├─Expression (pool_nonlin_3): 1-19       [1, 100, 5929, 1]         [1, 100, 5929, 1]         --                        --\n",
              "  ├─Dropout (drop_4): 1-20                 [1, 100, 5929, 1]         [1, 100, 5929, 1]         --                        --\n",
              "  ├─Conv2d (conv_4): 1-21                  [1, 100, 5929, 1]         [1, 200, 5848, 1]         200,000                   [10, 1]\n",
              "  ├─BatchNorm2d (bnorm_4): 1-22            [1, 200, 5848, 1]         [1, 200, 5848, 1]         400                       --\n",
              "  ├─Expression (nonlin_4): 1-23            [1, 200, 5848, 1]         [1, 200, 5848, 1]         --                        --\n",
              "  ├─MaxPool2d (pool_4): 1-24               [1, 200, 5848, 1]         [1, 200, 5794, 1]         --                        [3, 1]\n",
              "  ├─Expression (pool_nonlin_4): 1-25       [1, 200, 5794, 1]         [1, 200, 5794, 1]         --                        --\n",
              "  ├─Sequential (final_layer): 1-26         [1, 200, 5794, 1]         [1, 2, 5794]              --                        --\n",
              "  │    └─Conv2d (conv_classifier): 2-1     [1, 200, 5794, 1]         [1, 2, 5794, 1]           402                       [1, 1]\n",
              "  │    └─LogSoftmax (logsoftmax): 2-2      [1, 2, 5794, 1]           [1, 2, 5794, 1]           --                        --\n",
              "  │    └─Expression (squeeze): 2-3         [1, 2, 5794, 1]           [1, 2, 5794]              --                        --\n",
              "  ============================================================================================================================================\n",
              "  Total params: 277,052\n",
              "  Trainable params: 277,052\n",
              "  Non-trainable params: 0\n",
              "  Total mult-adds (G): 1.54\n",
              "  ============================================================================================================================================\n",
              "  Input size (MB): 0.50\n",
              "  Forward/backward pass size (MB): 34.30\n",
              "  Params size (MB): 1.05\n",
              "  Estimated Total Size (MB): 35.86\n",
              "  ============================================================================================================================================,\n",
              ")"
            ]
          },
          "execution_count": 36,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "clf.fit(window_train_set, y=None, epochs=n_epochs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Conclusion\n",
        "\n",
        "In this example, we used transfer learning (TL) as a way to learn\n",
        "representations from a large EEG data and transfer to a smaller dataset. \n",
        "\n",
        "\n",
        "## References\n",
        "\n",
        ".. [1] Darvishi-Bayazi, M. J., Ghaemi, M. S., Lesort, T., Arefin, M. R., Faubert, J., & Rish, I. (2024). Amplifying pathological detection in EEG signaling pathways through cross-dataset transfer learning. Computers in Biology and Medicine, 169, 107893.\n",
        "\n",
        ".. [2] Shawki, N., Shadin, M. G., Elseify, T., Jakielaszek, L., Farkas, T., Persidsky, Y., ... & Picone, J. (2022). Correction to: The temple university hospital digital pathology corpus. In Signal Processing in Medicine and Biology: Emerging Trends in Research and Applications (pp. C1-C1). Cham: Springer International Publishing..\n",
        "\n",
        ".. [3] Khan, H. A., Ul Ain, R., Kamboh, A. M., & Butt, H. T. (2022). The NMT scalp EEG dataset: an open-source annotated dataset of healthy and pathological EEG recordings for predictive modeling. Frontiers in neuroscience, 15, 755817.\n",
        "\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "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.10.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}