Lightning_BCI-autoencoder.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 185,
   "id": "f0e55b6a",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.io, scipy.interpolate\n",
    "import pathlib\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_model_summary import summary\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "import scipy.ndimage\n",
    "import plotly.tools as tls\n",
    "from pytorch_lightning.callbacks import Callback\n",
    "from pytorch_lightning import Trainer\n",
    "import wandb\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "\n",
    "L_FREQ, H_FREQ = 40, 300 # Lower and upper filtration bounds\n",
    "CHANNELS_NUM = 64        # Number of channels in ECoG data\n",
    "WAVELET_NUM = 40         # Number of wavelets in the indicated frequency range, with which the convolution is performed\n",
    "DOWNSAMPLE_FS = 100      # Desired sampling rate\n",
    "time_delay_secs = 0.2    # Time delay hyperparameter\n",
    "\n",
    "\n",
    "current_fs = DOWNSAMPLE_FS\n",
    "\n",
    "TYPE = \"test\"  # Script modes: \"train\" and \"test\"\n",
    "model_to_test = f\"{pathlib.Path().resolve()}/checkpoints/subj3_best-corr_mean_val=0.7361.ckpt\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "id": "429740ff",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class EcogFingerflexDataset(Dataset):\n",
    "    \"\"\"\n",
    "    The class that defines the sampling unit\n",
    "    \"\"\"\n",
    "    def __init__(self, path_to_ecog_data: str,\n",
    "                 path_to_fingerflex_data: str, sample_len: int, train = False):\n",
    "        \"\"\"\n",
    "        paths should point to .npy files\n",
    "        \"\"\"\n",
    "        self.ecog_data, self.fingerflex_data = np.load(path_to_ecog_data).astype('float32'),\\\n",
    "                                            np.load(path_to_fingerflex_data).astype('float32')\n",
    "        \n",
    "        self.duration = self.ecog_data.shape[2]\n",
    "        self.sample_len = sample_len                                 # sample size\n",
    "        self.stride = 1                                              # stride between samples\n",
    "        self.ds_len = (self.duration-self.sample_len) // self.stride\n",
    "        self.train = train\n",
    "        \n",
    "        print(\"Duration: \", self.duration, \"Ds_len:\", self.ds_len)\n",
    "    def __len__(self):\n",
    "        return self.ds_len\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "\n",
    "        sample_start = index*self.stride\n",
    "        sample_end = sample_start+self.sample_len\n",
    "\n",
    "        ecog_sample = self.ecog_data[...,sample_start:sample_end] # x\n",
    "        \n",
    "        fingerflex_sample = self.fingerflex_data[...,sample_start:sample_end] # y\n",
    "        \n",
    "        return ecog_sample, fingerflex_sample\n",
    "\n",
    "\n",
    "class EcogFingerflexDatamodule(pl.LightningDataModule):\n",
    "    \"\"\"\n",
    "    A class that encapsulates different datasets (for training and validation) and their dataloaders\n",
    "    \"\"\"\n",
    "    def __init__(self, sample_len: int, data_dir = \"/home/lomtev/brain_train/data\",\n",
    "                    batch_size=128, add_name=\"\"):\n",
    "        super().__init__()\n",
    "        self.data_dir = data_dir     # Path to data folder\n",
    "        self.sample_len = sample_len # Sample size\n",
    "        self.batch_size = batch_size # Dataloader batch size\n",
    "        self.add_name = add_name     #  dataset name\n",
    "        \n",
    "    def setup(self, stage = None):\n",
    "        if stage is None or stage == \"fit\":\n",
    "            self.train = EcogFingerflexDataset(f\"{self.data_dir}/train/ecog_data{self.add_name}.npy\",\n",
    "                                              f\"{self.data_dir}/train/fingerflex_data{self.add_name}.npy\",\n",
    "                                              self.sample_len, train = True)\n",
    "            \n",
    "            self.val = EcogFingerflexDataset(f\"{self.data_dir}/val/ecog_data{self.add_name}.npy\",\n",
    "                                              f\"{self.data_dir}/val/fingerflex_data{self.add_name}.npy\",\n",
    "                                              self.sample_len)\n",
    "        \n",
    "        if stage is None or stage == \"test\":\n",
    "            self.test = EcogFingerflexDataset(f\"{self.data_dir}/test/ecog_data{self.add_name}.npy\",\n",
    "                                              f\"{self.data_dir}/test/fingerflex_data{self.add_name}.npy\",\n",
    "                                              self.sample_len)\n",
    "    \n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(self.train, batch_size=self.batch_size, num_workers=4, shuffle=True)\n",
    "    \n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(self.val, batch_size=self.batch_size)\n",
    "    \n",
    "    def test_dataloader(self):\n",
    "        return DataLoader(self.test, batch_size=self.batch_size)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "id": "e34f9309",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def correlation_metric(x, y):\n",
    "    \"\"\"\n",
    "     Cosine distance calculation metric\n",
    "    \"\"\"\n",
    "    cos_metric = nn.CosineSimilarity(dim=-1, eps=1e-08)\n",
    "\n",
    "    cos_sim = torch.mean(cos_metric(x, y))\n",
    "\n",
    "    return cos_sim\n",
    "\n",
    "def corr_metric(x, y):\n",
    "    \"\"\"\n",
    "    Pearson correlation calculation metric between univariate vectors\n",
    "    \"\"\"\n",
    "    assert x.shape == y.shape  \n",
    "    r = np.corrcoef(x, y)[0, 1]\n",
    "    return r\n",
    "\n",
    "class BaseEcogFingerflexModel(pl.LightningModule):\n",
    "    \"\"\"\n",
    "        The class which encapsulates the model, its optimizer and the training process at different stages, including logging\n",
    "    \"\"\"\n",
    "    def __init__(self, model):\n",
    "        super().__init__()\n",
    "        self.model = model # Pytorch model\n",
    "        self.lr = 8.42e-5\n",
    "        \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        \n",
    "        y_hat = self.model(x)\n",
    "        \n",
    "        loss = F.mse_loss(y_hat, y)\n",
    "        corr = correlation_metric(y_hat, y)\n",
    "\n",
    "        self.log(\"train_loss\", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n",
    "        self.log(f\"cosine_dst_train\", corr, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n",
    "\n",
    "        return 0.5*loss + 0.5*(1. - corr) # возврат значения функции потерь\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self.model(x)\n",
    "        loss = F.mse_loss(y_hat, y)\n",
    "        \n",
    "        \n",
    "        corr = correlation_metric(y_hat, y)\n",
    "\n",
    "        self.log(\"val_loss\", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n",
    "        self.log(\"cosine_dst_val\", corr, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n",
    "        \n",
    "        return y_hat # Return the result for the validation callback\n",
    "    \n",
    "    def test_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        \n",
    "        y_hat = self.model(x)\n",
    "        \n",
    "        loss = F.mse_loss(y_hat, y)\n",
    "        self.log(\"test_loss\", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr, weight_decay=1e-6) # set optimizer, lr and L2 regularization coeff\n",
    "        return optimizer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "id": "cc8b97ca",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Here are the blocks that make up the final model + the model itself\n",
    "\"\"\"\n",
    "\n",
    "class ConvBlock(nn.Module):\n",
    "    \"\"\"\n",
    "    Convolution block:\n",
    "        - 1d conv\n",
    "        - layer norm by embedding axis\n",
    "        - activation\n",
    "        - dropout\n",
    "        - Max pooling\n",
    "    \"\"\"\n",
    "    def __init__(self, in_channels, out_channels, kernel_size, \n",
    "                 stride=1, dilation=1, p_conv_drop=0.1):\n",
    "        super(ConvBlock, self).__init__()\n",
    "        \n",
    "        # use it instead stride. \n",
    "        \n",
    "        self.conv1d = nn.Conv1d(in_channels, out_channels, \n",
    "                                kernel_size=kernel_size, \n",
    "                                bias=False, \n",
    "                                padding='same')\n",
    "        \n",
    "        \n",
    "        self.norm = nn.LayerNorm(out_channels)\n",
    "        self.activation = nn.GELU()\n",
    "        self.drop = nn.Dropout(p=p_conv_drop)\n",
    "\n",
    "        self.downsample = nn.MaxPool1d(kernel_size=stride, stride=stride)\n",
    "\n",
    "        self.stride = stride\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        \n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        x = self.conv1d(x)\n",
    "        \n",
    "        # norm by last axis.\n",
    "        x = torch.transpose(x, -2, -1) \n",
    "        x = self.norm(x)\n",
    "        x = torch.transpose(x, -2, -1)\n",
    "        \n",
    "        x = self.activation(x)\n",
    "        \n",
    "        x = self.drop(x)\n",
    "        \n",
    "        x = self.downsample(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "class UpConvBlock(nn.Module):\n",
    "    \"\"\"\n",
    "    Decoder convolution block\n",
    "    \"\"\"\n",
    "    def __init__(self, scale, **args):\n",
    "        super(UpConvBlock, self).__init__()\n",
    "        self.conv_block = ConvBlock(**args)\n",
    "        self.upsample = nn.Upsample(scale_factor=scale, mode='linear', align_corners=False)\n",
    "\n",
    "            \n",
    "    def forward(self, x):\n",
    "        \n",
    "        x = self.conv_block(x)\n",
    "        x = self.upsample(x)\n",
    "        return x    \n",
    "    \n",
    "    \n",
    "\n",
    "\n",
    "    \n",
    "    \n",
    "class AutoEncoder1D(nn.Module):\n",
    "    \"\"\"\n",
    "    This is the final Encoder-Decoder model with skip connections\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 n_electrodes=30,   # Number of channels\n",
    "                 n_freqs = 16,      # Number of wavelets\n",
    "                 n_channels_out=21, # Number of fingers\n",
    "                 channels = [8, 16, 32, 32],  # Number of features on each encoder layer\n",
    "                 kernel_sizes=[3, 3, 3],\n",
    "                 strides=[4, 4, 4],\n",
    "                 dilation=[1, 1, 1]\n",
    "                 ):\n",
    "        \n",
    "        super(AutoEncoder1D, self).__init__()\n",
    "        \n",
    "\n",
    "        self.n_electrodes = n_electrodes\n",
    "        self.n_freqs = n_freqs\n",
    "        self.n_inp_features = n_freqs*n_electrodes\n",
    "        self.n_channels_out = n_channels_out\n",
    "        \n",
    "        self.model_depth = len(channels)-1\n",
    "        self.spatial_reduce = ConvBlock(self.n_inp_features, channels[0], kernel_size=3) # Dimensionality reduction\n",
    "        \n",
    "        # Encoder part\n",
    "        self.downsample_blocks = nn.ModuleList([ConvBlock(channels[i], \n",
    "                                                        channels[i+1], \n",
    "                                                        kernel_sizes[i],\n",
    "                                                        stride=strides[i], \n",
    "                                                        dilation=dilation[i]) for i in range(self.model_depth)])\n",
    "        \n",
    "\n",
    "        channels = [ch for ch in channels[:-1]] + channels[-1:] # channels\n",
    "\n",
    "        # Decoder part\n",
    "        self.upsample_blocks = nn.ModuleList([UpConvBlock(scale=strides[i],\n",
    "                                                          in_channels=channels[i+1] if i == self.model_depth-1 else channels[i+1]*2 ,\n",
    "                                                          out_channels=channels[i],\n",
    "                                                          kernel_size=kernel_sizes[i]) for i in range(self.model_depth-1, -1, -1)])\n",
    "        \n",
    "        \n",
    "        self.conv1x1_one = nn.Conv1d(channels[0]*2, self.n_channels_out, kernel_size=1, padding='same') # final 1x1 conv\n",
    "      \n",
    "    def forward(self, x):\n",
    "\n",
    "        batch, elec, n_freq, time = x.shape\n",
    "        x = x.reshape(batch, -1, time)  # flatten the input\n",
    "        x = self.spatial_reduce(x)\n",
    "        \n",
    "        skip_connection = []\n",
    "        \n",
    "        for i in range(self.model_depth):\n",
    "            skip_connection.append(x)\n",
    "            x = self.downsample_blocks[i](x)\n",
    "\n",
    "        \n",
    "        for i in range(self.model_depth):\n",
    "            x = self.upsample_blocks[i](x)\n",
    "            x = torch.cat((x, skip_connection[-1 - i]), # skip connections\n",
    "                         dim=1)\n",
    "        \n",
    "        x = self.conv1x1_one(x)\n",
    "\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "outputs": [],
   "source": [
    "class ValidationCallback(Callback):\n",
    "    \"\"\"\n",
    "    Callback calculating the correlation at the end of each validation epoch on the whole dataset\n",
    "     and its logging (with visualization) in wandb. In addition, it performs prediction smoothing with Gaussian function\n",
    "    \"\"\"\n",
    "    def __init__(self, val_x, val_y, fg_num):\n",
    "        super().__init__()\n",
    "        self.val_x = val_x.T # Набор сигналов для валидации\n",
    "        self.val_y = val_y.T # Набор движений для валидации\n",
    "        self.fg_num = fg_num # Число пальцев\n",
    "\n",
    "    def on_validation_epoch_end(self, trainer, pl_module):\n",
    "        with torch.no_grad():\n",
    "            SIZE = 64\n",
    "            #SIZE = 256\n",
    "            bound = self.val_x.shape[0]//SIZE *SIZE\n",
    "\n",
    "            X_test = self.val_x[:bound]\n",
    "            y_test = self.val_y[:bound]\n",
    "            x_batch = torch.from_numpy(X_test).float().to(\"cuda:3\")\n",
    "\n",
    "            x_batch = x_batch.T\n",
    "\n",
    "            x_batch = torch.unsqueeze(x_batch, 0)\n",
    "\n",
    "            y_hat = pl_module.model(x_batch)[0] # Running data through the model\n",
    "            y_hat = y_hat.cpu().detach().numpy()\n",
    "            STRIDE = 1 # It is possible to validate with stride\n",
    "            y_prediction = y_hat.T[::int(STRIDE*(DOWNSAMPLE_FS/100)), :]\n",
    "            y_prediction = scipy.ndimage.gaussian_filter1d(y_prediction.T,sigma=6).T # Prediction smoothing with Gaussian function\n",
    "\n",
    "            y_test = y_test[::int(STRIDE*(DOWNSAMPLE_FS/100)), :]\n",
    "\n",
    "\n",
    "            h, w = self.fg_num//2, self.fg_num - self.fg_num//2\n",
    "            fig, ax = plt.subplots(h, w, figsize = (h*5, w*6), sharex=True, sharey=True) # Making pair plots of true motion and prediction\n",
    "            corrs = []\n",
    "\n",
    "            for roi in range(self.fg_num):\n",
    "                y_hat = y_prediction[:, roi]\n",
    "                y_test_roi = y_test[:, roi]\n",
    "                corr_tmp = corr_metric(y_hat, y_test_roi) # Correlation сalculation\n",
    "                corrs.append(corr_tmp)\n",
    "                axi = ax.flat[roi]\n",
    "                axi.plot(y_hat, label= 'prediction')\n",
    "                axi.plot(y_test_roi, label = 'true')\n",
    "\n",
    "                axi.set_title(\"RoI {}_corr {:.2f}\".format(roi, corr_tmp))\n",
    "\n",
    "            corr_mean = np.mean(corrs)\n",
    "            pl_module.log(\"corr_mean_val\", corr_mean, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n",
    "            #wandb.log({\"corr_mean_val\" : corr_mean })\n",
    "            wandb.log({f\"plots\": fig}) # Logging charts\n",
    "\n",
    "\n",
    "class TestCallback:\n",
    "    \"\"\"\n",
    "    Callback, which calculates the correlation on the whole dataset and visualizes it in case of testing.\n",
    "    In addition, it also produces exactly the same prediction smoothing with the Gaussian function\n",
    "    \"\"\"\n",
    "    def __init__(self, val_x, val_y, fg_num):\n",
    "        super().__init__()\n",
    "        self.val_x = val_x.T\n",
    "        self.val_y = val_y.T\n",
    "        self.fg_num = fg_num\n",
    "\n",
    "    def test(self, pl_module):\n",
    "        with torch.no_grad():\n",
    "            SIZE = 64\n",
    "            bound = self.val_x.shape[0]//SIZE *SIZE\n",
    "\n",
    "            X_test = self.val_x[:bound]\n",
    "            y_test = self.val_y[:bound]\n",
    "            x_batch = torch.from_numpy(X_test).float()#.to(\"cuda:3\")\n",
    "\n",
    "            x_batch = x_batch.T\n",
    "\n",
    "            x_batch = torch.unsqueeze(x_batch, 0)\n",
    "\n",
    "            y_hat = pl_module.model(x_batch)[0]\n",
    "            y_hat = y_hat.cpu().detach().numpy()\n",
    "            STRIDE = 1\n",
    "            y_prediction = y_hat.T[::int(STRIDE*(DOWNSAMPLE_FS/100)), :]\n",
    "            y_prediction = scipy.ndimage.gaussian_filter1d(y_prediction.T,sigma=1).T\n",
    "\n",
    "            y_test = y_test[::int(STRIDE*(DOWNSAMPLE_FS/100)), :]\n",
    "\n",
    "            np.save(f\"{pathlib.Path().resolve()}/res_npy/prediction2.npy\", y_prediction)\n",
    "            np.save(f\"{pathlib.Path().resolve()}/res_npy/true2.npy\", y_test)\n",
    "\n",
    "            h, w = self.fg_num//2, self.fg_num - self.fg_num//2\n",
    "            fig, ax = plt.subplots(h, w, figsize = (h*35, w*6), sharex=True, sharey=True)\n",
    "            corrs = []\n",
    "\n",
    "            for roi in range(self.fg_num):\n",
    "                y_hat = y_prediction[:, roi]\n",
    "                y_test_roi = y_test[:, roi]\n",
    "                corr_tmp = corr_metric(y_hat, y_test_roi)\n",
    "                corrs.append(corr_tmp)\n",
    "                axi = ax.flat[roi]\n",
    "                axi.plot(y_hat, label= 'prediction')\n",
    "                axi.plot(y_test_roi, label = 'true')\n",
    "\n",
    "                axi.set_title(\"RoI {}_corr {:.2f}\".format(roi, corr_tmp))\n",
    "\n",
    "            corr_mean = np.mean(corrs)\n",
    "\n",
    "            plotly_fig = tls.mpl_to_plotly(fig) # Converting matplotlib image to plotly\n",
    "            print(corr_mean)\n",
    "            plotly_fig.write_html(\"res.html\") # Writing the interactive visualization to an html file\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "id": "1d1af6f1",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": "'-----------------------------------------------------------------------\\n      Layer (type)        Output Shape         Param #     Tr. Param #\\n=======================================================================\\n       ConvBlock-1        [4, 32, 256]         245,824         245,824\\n       ConvBlock-2        [4, 32, 128]           7,232           7,232\\n       ConvBlock-3         [4, 64, 64]          14,464          14,464\\n       ConvBlock-4         [4, 64, 32]          20,608          20,608\\n       ConvBlock-5        [4, 128, 16]          41,216          41,216\\n       ConvBlock-6         [4, 128, 8]          82,176          82,176\\n     UpConvBlock-7        [4, 128, 16]          82,176          82,176\\n     UpConvBlock-8         [4, 64, 32]          82,048          82,048\\n     UpConvBlock-9         [4, 64, 64]          41,088          41,088\\n    UpConvBlock-10        [4, 32, 128]          28,736          28,736\\n    UpConvBlock-11        [4, 32, 256]          14,400          14,400\\n         Conv1d-12         [4, 5, 256]             325             325\\n=======================================================================\\nTotal params: 660,293\\nTrainable params: 660,293\\nNon-trainable params: 0\\n-----------------------------------------------------------------------'"
     },
     "execution_count": 190,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "\n",
    "#256 -> 256/4 * (32) -> 64/4 * (64) -> 16 * (64)\n",
    "\n",
    "SAMPLE_LEN = 256 # Window size\n",
    "finger_num = 5   # Number of fingers\n",
    "\n",
    "hp_autoencoder = dict(channels = [32, 32, 64, 64, 128, 128], \n",
    "                        kernel_sizes=[7, 7, 5, 5, 5],\n",
    "                        strides=[2, 2, 2, 2, 2],\n",
    "                        dilation=[1, 1, 1, 1, 1],\n",
    "                        n_electrodes = CHANNELS_NUM,\n",
    "                        n_freqs = WAVELET_NUM,\n",
    "                        n_channels_out = finger_num) # A set of features for the model\n",
    "\n",
    "model = AutoEncoder1D(**hp_autoencoder).to(\"cuda:0\")\n",
    "\n",
    "\n",
    "lighning_wrapper = BaseEcogFingerflexModel(model) # Wrapping in pytorch-lightning class\n",
    "\n",
    "\n",
    "\n",
    "dm = EcogFingerflexDatamodule(sample_len=SAMPLE_LEN, add_name=\"\")\n",
    "summary(model, torch.zeros(4, CHANNELS_NUM,WAVELET_NUM, SAMPLE_LEN).to(\"cuda:0\"),\n",
    "       show_input=False) # Model structure output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "id": "b70e6d2d",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "SAVE_PATH = f\"{pathlib.Path().resolve()}/data/\"\n",
    "\n",
    "def load_data(ecog_data_path, fingerflex_data_path):\n",
    "    ecog_data = np.load(ecog_data_path)\n",
    "    fingerflex_data = np.load(fingerflex_data_path)\n",
    "    return ecog_data, fingerflex_data\n",
    "\n",
    "ecog_data_val, fingerflex_data_val = load_data(f\"{SAVE_PATH}/val/ecog_data.npy\", f\"{SAVE_PATH}/val/fingerflex_data.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.7110533417608561\n"
     ]
    }
   ],
   "source": [
    "### TO TRAIN ###\n",
    "\n",
    "if TYPE == \"train\":\n",
    "    wandb.init(project=\"BCI_comp\") # Logger initialization\n",
    "    wandb_logger = WandbLogger()\n",
    "\n",
    "\n",
    "    checkpoint_callback = ModelCheckpoint( # Initializing a callback to save model checkpoints\n",
    "        save_top_k=2,\n",
    "        monitor=\"corr_mean_val\",\n",
    "        mode=\"max\",\n",
    "        dirpath=\"checkpoints\",\n",
    "        filename=\"model-{epoch:02d}-{corr_mean_val}\",\n",
    "    )\n",
    "\n",
    "    # The Trainer class encapsulates the interaction of model, data and logger\n",
    "    trainer = Trainer(gpus=[3], max_epochs=20, logger=wandb_logger, callbacks=[ValidationCallback(ecog_data_val,\n",
    "                                                                                                  fingerflex_data_val,\n",
    "                                                                                                  finger_num),\n",
    "                                                                               checkpoint_callback])\n",
    "    trainer.fit(lighning_wrapper, dm) # Model training process\n",
    "    wandb.finish()                    # Signal to end the logging\n",
    "\n",
    "elif TYPE == \"test\":\n",
    "    ### TO TEST ###\n",
    "\n",
    "    trained_model = BaseEcogFingerflexModel.load_from_checkpoint(\n",
    "        checkpoint_path=model_to_test,\n",
    "        model=AutoEncoder1D(**hp_autoencoder))\n",
    "    test_callback = TestCallback(ecog_data_val, fingerflex_data_val, finger_num)\n",
    "    test_callback.test(trained_model) # Testing"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "language": "python",
   "display_name": "Python 3 (ipykernel)"
  },
  "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}