UnendingPain.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3704e353-1ea3-4ef5-a48c-de044638b589",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import mne\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.signal as signal\n",
    "%matplotlib qt\n",
    "\n",
    "from sklearn.metrics import mean_squared_error, accuracy_score, precision_score\n",
    "from scipy.stats import pearsonr\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "from skimage.metrics import structural_similarity as ssim\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import torch\n",
    "import snntorch as snn\n",
    "from snntorch import spikegen, surrogate, utils\n",
    "import snntorch.functional as SF\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "import logging\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "logging.getLogger('mne').setLevel(logging.WARNING)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "324d66d1-1548-40a8-a078-0e6a3317e5f1",
   "metadata": {},
   "source": [
    "# 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "236a99fb-f9d8-4f3d-a70c-81235edd301b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([128, 1, 28, 28])\n",
      "torch.Size([128])\n"
     ]
    }
   ],
   "source": [
    "# dataloader arguments\n",
    "batch_size = 128\n",
    "data_path='/tmp/data/mnist'\n",
    "\n",
    "dtype = torch.float\n",
    "device = torch.device(\"cuda\")\n",
    "\n",
    "transform = transforms.Compose([\n",
    "            transforms.Resize((28, 28)),\n",
    "            transforms.Grayscale(),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0,), (1,))])\n",
    "\n",
    "mnist_train = datasets.MNIST(data_path, train=True, download=True, transform=transform)\n",
    "mnist_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)\n",
    "\n",
    "# Create DataLoaders\n",
    "train_loader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "\n",
    "for inputs, targets in train_loader:\n",
    "    print(inputs.shape)\n",
    "    print(targets.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "66bd91a1-569c-42d6-8320-050c62d7e0ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "spike_grad = surrogate.fast_sigmoid(slope=25)\n",
    "beta = 0.5\n",
    "num_steps = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "67a0fa6f-da03-4e7d-86e5-d5d5bd801515",
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward_pass(input_data, model, num_steps):\n",
    "    spk_rec = []\n",
    "    mem_rec = []\n",
    "    utils.reset(model)\n",
    "\n",
    "    for i in range(num_steps):\n",
    "        spk, mem = model(input_data)\n",
    "        spk_rec.append(spk)\n",
    "        mem_rec.append(mem)\n",
    "\n",
    "    return torch.stack(spk_rec), torch.stack(mem_rec)\n",
    "\n",
    "def batch_acc(data_loader, model, num_steps):\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        acc = 0\n",
    "        total = 0\n",
    "        for inputs, targets in iter(data_loader):\n",
    "            inputs = inputs.to(device)\n",
    "            targets = targets.to(device)\n",
    "        \n",
    "            spk_rec, mem_rec = forward_pass(inputs, model, num_steps)\n",
    "            acc += SF.accuracy_rate(spk_rec, targets) * spk_rec.size(1)\n",
    "            total += spk_rec.size(1)\n",
    "    return acc / total    \n",
    "    \n",
    "class MyEEGSNNModel(nn.Module):\n",
    "    def __init__(self, n_outputs):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.conv1 = nn.Conv2d(1, 12, 5)\n",
    "        self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        self.conv2 = nn.Conv2d(12, 64, 5)\n",
    "        self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        self.fc1 = nn.Linear(64*4*4, n_outputs)\n",
    "        self.lif3 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        # self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=(1, 5))\n",
    "        # self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        # self.conv2 = nn.Conv2d(in_channels=16, out_channels=64, kernel_size=(1, 5))\n",
    "        # self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        # self.fc1 = nn.Linear(64*14*157, n_outputs)\n",
    "        # self.lif3 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        mem1 = self.lif1.init_leaky()\n",
    "        mem2 = self.lif2.init_leaky()\n",
    "        mem3 = self.lif3.init_leaky()\n",
    "\n",
    "        cur1 = F.max_pool2d(self.conv1(x), kernel_size=(2, 2))\n",
    "        # cur1 = F.max_pool2d(self.conv1(x), kernel_size=(1, 2))\n",
    "        spk1, mem1 = self.lif1(cur1, mem1)\n",
    "         \n",
    "        cur2 = F.max_pool2d(self.conv2(spk1), kernel_size=(2, 2))\n",
    "        # cur2 = F.max_pool2d(self.conv2(spk1), kernel_size=(1, 2))\n",
    "        spk2, mem2 = self.lif2(cur2, mem2)\n",
    "        \n",
    "        cur3 = self.fc1(spk2.view(batch_size, -1))\n",
    "        spk3, mem3 = self.lif3(cur3, mem3)\n",
    "\n",
    "        return spk3, mem3   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d361f245-9ae1-41be-99cc-3b65466e273d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                          | 0/468 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([20, 128, 10])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "num_classes = 10\n",
    "model = MyEEGSNNModel(num_classes).to(\"cuda\")\n",
    "\n",
    "for inputs, targets in tqdm(train_loader):\n",
    "    inputs = inputs.to(device)\n",
    "    targets = targets.to(device)\n",
    "\n",
    "    spk_rec, _= forward_pass(inputs, model, num_steps)\n",
    "    print(spk_rec.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6f09c3b0-8e3b-46cd-be24-2b9a30894471",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 468/468 [01:14<00:00,  6.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/5], Loss: 2.3026, Accuracy: 0.0979\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 468/468 [01:13<00:00,  6.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [2/5], Loss: 2.0290, Accuracy: 0.4005\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 468/468 [01:13<00:00,  6.35it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [3/5], Loss: 1.5172, Accuracy: 0.8882\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 468/468 [01:13<00:00,  6.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [4/5], Loss: 1.5029, Accuracy: 0.9171\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 468/468 [01:13<00:00,  6.34it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [5/5], Loss: 1.4988, Accuracy: 0.9263\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "n_epochs = 5\n",
    "learning_rate = 1e-3\n",
    "loss_history = []\n",
    "acc_history = []\n",
    "\n",
    "num_classes = 10\n",
    "# num_classes = 4\n",
    "model = MyEEGSNNModel(num_classes).to(\"cuda\")\n",
    "for param in model.parameters():\n",
    "    print(param.requires_grad)\n",
    "criterion = SF.ce_rate_loss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, betas=(0.9, 0.999))\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    epoch_loss = 0.0\n",
    "    train_acc = 0.0\n",
    "    counter = 0\n",
    "    acc_co = 0\n",
    "    \n",
    "    for inputs, targets in tqdm(iter(train_loader)):\n",
    "        inputs = inputs.to(device)\n",
    "        targets = targets.to(device)\n",
    "        \n",
    "        model.train()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        spk_rec, _ = forward_pass(inputs, model, num_steps)\n",
    "        loss = criterion(spk_rec, targets)\n",
    "            \n",
    "        loss.backward()         \n",
    "        optimizer.step()  \n",
    "        epoch_loss += loss.item()\n",
    "        counter+=1\n",
    "\n",
    "        if counter % 75 == 0:\n",
    "        # if counter % 6 == 0:\n",
    "            train_acc += batch_acc(test_loader, model, num_steps)\n",
    "            acc_co+=1\n",
    "        \n",
    "    train_loss = epoch_loss / counter\n",
    "    train_acc = train_acc.item() / acc_co\n",
    "    # train_acc = train_acc.item() / acc_co\n",
    "    loss_history.append(train_loss)\n",
    "    acc_history.append(train_acc)\n",
    "    print(f'Epoch [{epoch + 1}/{n_epochs}], Loss: {train_loss:.4f}, Accuracy: {train_acc:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccee1120-4a6c-4fa1-a4a5-1eb2d808f898",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "for inputs, targets in tqdm(test_loader):\n",
    "    with torch.no_grad():\n",
    "        spk_rec, mem_rec = model(inputs, num_steps)\n",
    "        spk_count = torch.sum(spk_rec, dim=0)\n",
    "        outputs = F.softmax(spk_count)\n",
    "        \n",
    "    loss = criterion(outputs, targets)\n",
    "    epoch_loss += loss.item()\n",
    "print(f'Loss: {epoch_loss / len(test_loader):.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5d4fb93-30dd-4b6b-a6f0-7075846c2d00",
   "metadata": {},
   "source": [
    "# 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7c4ce8c9-8600-406f-8116-0a0424391d4a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 112/112 [00:02<00:00, 39.83it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(4480, 14, 640)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████| 4480/4480 [00:01<00:00, 3533.89it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x1674dad7100>]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "folder_path = \"Extras/GAMEEMO_EPOCH/5sec\"\n",
    "epochs_data = []\n",
    "\n",
    "for file in tqdm(os.listdir(folder_path)):\n",
    "    file_path = os.path.join(folder_path, file)\n",
    "    epochs = mne.read_epochs(file_path, preload=True)\n",
    "    epochs_choice_idx = np.random.choice(epochs.get_data().shape[0], size=40, replace=False)\n",
    "    epochs_choice = epochs.get_data()[epochs_choice_idx]\n",
    "    epochs_data.append(epochs_choice)\n",
    "\n",
    "epochs_data = np.stack(epochs_data, axis=0)\n",
    "epochs_data_reshaped = epochs_data.reshape(-1, epochs_data.shape[-2], epochs_data.shape[-1])\n",
    "print(epochs_data_reshaped.shape)\n",
    "\n",
    "for i in tqdm(range(epochs_data_reshaped.shape[0])):\n",
    "    for j in range(epochs_data_reshaped.shape[1]):\n",
    "        data = epochs_data_reshaped[i,j,:]\n",
    "        epochs_data_reshaped[i,j,:] = (data - np.min(data)) / (np.max(data) - np.min(data))\n",
    "\n",
    "plt.figure(1)\n",
    "plt.plot(epochs_data_reshaped[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f6528dfa-9d74-4e49-9d56-191fc965ef7e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(4480,)\n",
      "(4480, 4)\n",
      "[0 0 0 ... 3 3 3]\n",
      "[[1. 0. 0. 0.]\n",
      " [1. 0. 0. 0.]\n",
      " [1. 0. 0. 0.]\n",
      " ...\n",
      " [0. 0. 0. 1.]\n",
      " [0. 0. 0. 1.]\n",
      " [0. 0. 0. 1.]]\n"
     ]
    }
   ],
   "source": [
    "label_list = [0, 1, 2, 3]\n",
    "labels = []\n",
    "for i in range(int(112/4)):\n",
    "    labels.extend(label_list)\n",
    "\n",
    "labels = np.array(labels)\n",
    "labels = np.repeat(labels, int(epochs_data_reshaped.shape[0]/112))\n",
    "oh_labels = np.eye(4)[labels]\n",
    "print(labels.shape)\n",
    "print(oh_labels.shape)\n",
    "np.set_printoptions(threshold=20)\n",
    "print(labels)\n",
    "print(oh_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d425a921-d91f-495c-9f76-bc9b9d2ad9cc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3360, 14, 640]) torch.Size([1120, 14, 640]) torch.Size([3360]) torch.Size([1120])\n"
     ]
    }
   ],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(epochs_data_reshaped, labels, test_size=0.25, random_state=56)\n",
    "\n",
    "num_steps = 32\n",
    "device = torch.device(\"cuda\") \n",
    "X_train = torch.tensor(X_train, dtype=torch.float32)\n",
    "X_test  = torch.tensor(X_test, dtype=torch.float32)\n",
    "y_train = torch.tensor(y_train, dtype=torch.long)\n",
    "y_test  = torch.tensor(y_test, dtype=torch.long)\n",
    "print(X_train.shape, X_test.shape, y_train.shape, y_test.shape)\n",
    "\n",
    "batch_size = 80\n",
    "train_dataset = TensorDataset(X_train.unsqueeze(1), y_train)\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "\n",
    "test_dataset = TensorDataset(X_test.unsqueeze(1), y_test)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, drop_last=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "56b0b3ac-2c22-45ca-af4a-afda5ebfd3c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "spike_grad = surrogate.fast_sigmoid(slope=25)\n",
    "beta = 0.5\n",
    "num_steps = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cb7a45bc-c96f-4c91-92b2-135e99f4447a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward_pass(input_data, model, num_steps):\n",
    "    spk_rec = []\n",
    "    mem_rec = []\n",
    "    utils.reset(model)\n",
    "\n",
    "    for i in range(num_steps):\n",
    "        spk, mem = model(input_data)\n",
    "        spk_rec.append(spk)\n",
    "        mem_rec.append(mem)\n",
    "\n",
    "    return torch.stack(spk_rec), torch.stack(mem_rec)\n",
    "\n",
    "def batch_acc(data_loader, model, num_steps):\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        acc = 0\n",
    "        total = 0\n",
    "        for inputs, targets in iter(data_loader):\n",
    "            inputs = inputs.to(device)\n",
    "            targets = targets.to(device)\n",
    "        \n",
    "            spk_rec, mem_rec = forward_pass(inputs, model, num_steps)\n",
    "            acc += SF.accuracy_rate(spk_rec, targets) * spk_rec.size(1)\n",
    "            total += spk_rec.size(1)\n",
    "    return acc / total    \n",
    "    \n",
    "class MyEEGSNNModel(nn.Module):\n",
    "    def __init__(self, n_outputs):\n",
    "        super().__init__()\n",
    "        \n",
    "        # self.conv1 = nn.Conv2d(1, 12, 5)\n",
    "        # self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        # self.conv2 = nn.Conv2d(12, 64, 5)\n",
    "        # self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        # self.fc1 = nn.Linear(64*4*4, n_outputs)\n",
    "        # self.lif3 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=(1, 5))\n",
    "        self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        self.conv2 = nn.Conv2d(in_channels=16, out_channels=64, kernel_size=(1, 5))\n",
    "        self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        self.fc1 = nn.Linear(64*14*157, n_outputs)\n",
    "        self.lif3 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        mem1 = self.lif1.init_leaky()\n",
    "        mem2 = self.lif2.init_leaky()\n",
    "        mem3 = self.lif3.init_leaky()\n",
    "\n",
    "        # cur1 = F.max_pool2d(self.conv1(x), kernel_size=(2, 2))\n",
    "        cur1 = F.max_pool2d(self.conv1(x), kernel_size=(1, 2))\n",
    "        spk1, mem1 = self.lif1(cur1, mem1)\n",
    "         \n",
    "        # cur2 = F.max_pool2d(self.conv2(spk1), kernel_size=(2, 2))\n",
    "        cur2 = F.max_pool2d(self.conv2(spk1), kernel_size=(1, 2))\n",
    "        spk2, mem2 = self.lif2(cur2, mem2)\n",
    "        \n",
    "        cur3 = self.fc1(spk2.view(batch_size, -1))\n",
    "        spk3, mem3 = self.lif3(cur3, mem3)\n",
    "\n",
    "        return spk3, mem3   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "96310131-81b2-44ba-a329-153ab095c523",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                           | 0/42 [00:01<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([20, 80, 4])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "num_classes = 4\n",
    "model = MyEEGSNNModel(num_classes).to(\"cuda\")\n",
    "\n",
    "for inputs, targets in tqdm(train_loader):\n",
    "    inputs = inputs.to(device)\n",
    "    targets = targets.to(device)\n",
    "\n",
    "    spk_rec, _= forward_pass(inputs, model, num_steps)\n",
    "    print(spk_rec.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "75b1af9b-2015-4c74-80f6-a0cd6c5c6fe0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 42/42 [02:20<00:00,  3.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/5], Loss: 1.3863, Accuracy: 0.2536\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 42/42 [02:19<00:00,  3.31s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [2/5], Loss: 1.3863, Accuracy: 0.2536\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 42/42 [02:18<00:00,  3.30s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [3/5], Loss: 1.3863, Accuracy: 0.2536\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 42/42 [02:18<00:00,  3.31s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [4/5], Loss: 1.3863, Accuracy: 0.2536\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 42/42 [02:19<00:00,  3.31s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [5/5], Loss: 1.3863, Accuracy: 0.2536\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "n_epochs = 5\n",
    "learning_rate = 1e-3\n",
    "loss_history = []\n",
    "acc_history = []\n",
    "\n",
    "# num_classes = 10\n",
    "num_classes = 4\n",
    "model = MyEEGSNNModel(num_classes).to(\"cuda\")\n",
    "for param in model.parameters():\n",
    "    print(param.requires_grad)\n",
    "criterion = SF.ce_rate_loss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, betas=(0.9, 0.999))\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    epoch_loss = 0.0\n",
    "    train_acc = 0.0\n",
    "    counter = 0\n",
    "    acc_co = 0\n",
    "    \n",
    "    for inputs, targets in tqdm(iter(train_loader)):\n",
    "        inputs = inputs.to(device)\n",
    "        targets = targets.to(device)\n",
    "        \n",
    "        model.train()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        spk_rec, _ = forward_pass(inputs, model, num_steps)\n",
    "        loss = criterion(spk_rec, targets)\n",
    "            \n",
    "        loss.backward()         \n",
    "        optimizer.step()  \n",
    "        epoch_loss += loss.item()\n",
    "        counter+=1\n",
    "\n",
    "        # if counter % 75 == 0:\n",
    "        if counter % 6 == 0:\n",
    "            train_acc += batch_acc(test_loader, model, num_steps)\n",
    "            acc_co+=1\n",
    "        \n",
    "    train_loss = epoch_loss / counter\n",
    "    train_acc = train_acc.item() / acc_co\n",
    "    loss_history.append(train_loss)\n",
    "    acc_history.append(train_acc)\n",
    "    print(f'Epoch [{epoch + 1}/{n_epochs}], Loss: {train_loss:.4f}, Accuracy: {train_acc:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "381ec1ac-a15c-4492-a112-af017b283d81",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "360427bf-6faa-4430-95b1-846bb44a9d21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataloader arguments\n",
    "batch_size = 128\n",
    "data_path='/tmp/data/mnist'\n",
    "\n",
    "dtype = torch.float\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"mps\") if torch.backends.mps.is_available() else torch.device(\"cpu\")\n",
    "\n",
    "# Define a transform\n",
    "transform = transforms.Compose([\n",
    "            transforms.Resize((28, 28)),\n",
    "            transforms.Grayscale(),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0,), (1,))])\n",
    "\n",
    "mnist_train = datasets.MNIST(data_path, train=True, download=True, transform=transform)\n",
    "mnist_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)\n",
    "\n",
    "# Create DataLoaders\n",
    "train_loader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "\n",
    "# neuron and simulation parameters\n",
    "spike_grad = surrogate.fast_sigmoid(slope=25)\n",
    "beta = 0.5\n",
    "num_steps = 50\n",
    "\n",
    "# Define Network\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        # Initialize layers\n",
    "        self.conv1 = nn.Conv2d(1, 12, 5)\n",
    "        self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        self.conv2 = nn.Conv2d(12, 64, 5)\n",
    "        self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        self.fc1 = nn.Linear(64*4*4, 10)\n",
    "        self.lif3 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        # Initialize hidden states and outputs at t=0\n",
    "        mem1 = self.lif1.init_leaky()\n",
    "        mem2 = self.lif2.init_leaky()\n",
    "        mem3 = self.lif3.init_leaky()\n",
    "\n",
    "        cur1 = F.max_pool2d(self.conv1(x), 2)\n",
    "        spk1, mem1 = self.lif1(cur1, mem1)\n",
    "\n",
    "        cur2 = F.max_pool2d(self.conv2(spk1), 2)\n",
    "        spk2, mem2 = self.lif2(cur2, mem2)\n",
    "\n",
    "        cur3 = self.fc1(spk2.view(batch_size, -1))\n",
    "        spk3, mem3 = self.lif3(cur3, mem3)\n",
    "\n",
    "        return spk3, mem3\n",
    "\n",
    "def forward_pass(net, num_steps, data):\n",
    "  mem_rec = []\n",
    "  spk_rec = []\n",
    "  utils.reset(net)  # resets hidden states for all LIF neurons in net\n",
    "\n",
    "  for step in range(num_steps):\n",
    "      spk_out, mem_out = net(data)\n",
    "      spk_rec.append(spk_out)\n",
    "      mem_rec.append(mem_out)\n",
    "\n",
    "  return torch.stack(spk_rec), torch.stack(mem_rec)\n",
    "\n",
    "def batch_accuracy(train_loader, net, num_steps):\n",
    "  with torch.no_grad():\n",
    "    total = 0\n",
    "    acc = 0\n",
    "    net.eval()\n",
    "\n",
    "    train_loader = iter(train_loader)\n",
    "    for data, targets in train_loader:\n",
    "      data = data.to(device)\n",
    "      targets = targets.to(device)\n",
    "      spk_rec, _ = forward_pass(net, num_steps, data)\n",
    "\n",
    "      acc += SF.accuracy_rate(spk_rec, targets) * spk_rec.size(1)\n",
    "      total += spk_rec.size(1)\n",
    "\n",
    "  return acc/total\n",
    "\n",
    "# already imported snntorch.functional as SF\n",
    "loss_fn = SF.ce_rate_loss()\n",
    "\n",
    "net = Net().to(device)\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=1e-2, betas=(0.9, 0.999))\n",
    "num_epochs = 1\n",
    "loss_hist = []\n",
    "test_acc_hist = []\n",
    "counter = 0\n",
    "\n",
    "# Outer training loop\n",
    "for epoch in range(num_epochs):\n",
    "\n",
    "    # Training loop\n",
    "    for data, targets in tqdm(iter(train_loader)):\n",
    "        data = data.to(device)\n",
    "        targets = targets.to(device)\n",
    "\n",
    "        # forward pass\n",
    "        net.train()\n",
    "        spk_rec, _ = forward_pass(net, num_steps, data)\n",
    "\n",
    "        # initialize the loss & sum over time\n",
    "        loss_val = loss_fn(spk_rec, targets)\n",
    "\n",
    "        # Gradient calculation + weight update\n",
    "        optimizer.zero_grad()\n",
    "        loss_val.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # Store loss history for future plotting\n",
    "        loss_hist.append(loss_val.item())\n",
    "\n",
    "        # Test set\n",
    "        if counter % 50 == 0:\n",
    "            with torch.no_grad():\n",
    "                net.eval()\n",
    "    \n",
    "                # Test set forward pass\n",
    "                test_acc = batch_accuracy(test_loader, net, num_steps)\n",
    "                print(f\"Iteration {counter}, Test Acc: {test_acc * 100:.2f}%\\n\")\n",
    "                test_acc_hist.append(test_acc.item())\n",
    "\n",
    "        counter += 1\n",
    "\n",
    "print(batch_accuracy(test_loader, net, num_steps))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd7b109b-fb53-4b22-b81f-674c5e346e04",
   "metadata": {},
   "source": [
    "# 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0228af7d-d9af-46f9-bfb9-f556d5f23cf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import mne\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.signal as signal\n",
    "import cv2\n",
    "from PIL import Image\n",
    "%matplotlib qt\n",
    "\n",
    "from sklearn.metrics import mean_squared_error, accuracy_score, precision_score\n",
    "from scipy.stats import pearsonr\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "from skimage.metrics import structural_similarity as ssim\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import torch\n",
    "import snntorch as snn\n",
    "from snntorch import spikegen, surrogate, utils\n",
    "import snntorch.functional as SF\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "import logging\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "logging.getLogger('mne').setLevel(logging.WARNING)\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "# device = \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c128f599-3ff5-492a-8719-6f0bf63a3d21",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(672,)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 672/672 [00:07<00:00, 85.94it/s]\n"
     ]
    }
   ],
   "source": [
    "img_dir = \"shuvankar_spectro_snn\"\n",
    "img_files = np.array(os.listdir(img_dir))\n",
    "print(img_files.shape)\n",
    "\n",
    "data_array = []\n",
    "labels_array = []\n",
    "\n",
    "for files in tqdm(img_files):\n",
    "    data_path = os.path.join(img_dir, files)\n",
    "    read_data = cv2.imread(data_path, cv2.COLOR_BGR2GRAY)\n",
    "    read_data = (read_data - np.min(read_data))/(np.max(read_data) - np.min(read_data))\n",
    "    read_data = cv2.resize(read_data, (200, 200))\n",
    "    data_array.append(read_data)\n",
    "    labels_array.append(int(files[4]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e6ba0446-7d4f-4765-8e35-3d5538929e91",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(data_array, labels_array, test_size=0.2, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ad91d5e0-41b7-45cb-9f11-b8634385625b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_data_array = torch.tensor(X_train, device=device, dtype=torch.float32)\n",
    "train_labels_array = torch.tensor(y_train, device=device, dtype=torch.long)\n",
    "test_data_array = torch.tensor(X_test, device=device, dtype=torch.float32)\n",
    "test_labels_array = torch.tensor(y_test, device=device, dtype=torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f602f651-daaa-4d23-a244-9f2290116c73",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = TensorDataset(train_data_array, train_labels_array)\n",
    "test_dataset = TensorDataset(test_data_array, test_labels_array)\n",
    "\n",
    "batch_size = 32\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, drop_last=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f5a64510-9b9a-4345-ae2f-73ceb07c3a42",
   "metadata": {},
   "outputs": [],
   "source": [
    "spike_grad = surrogate.fast_sigmoid(slope=25)\n",
    "beta = 0.5\n",
    "num_steps = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ad38a2f6-069f-40fa-b128-42328d30f165",
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward_pass(input_data, model, num_steps):\n",
    "    spk_rec = []\n",
    "    mem_rec = []\n",
    "    utils.reset(model)\n",
    "\n",
    "    for i in range(num_steps):\n",
    "        spk, mem = model(input_data)\n",
    "        spk_rec.append(spk)\n",
    "        mem_rec.append(mem)\n",
    "\n",
    "    return torch.stack(spk_rec), torch.stack(mem_rec)\n",
    "\n",
    "def batch_acc(data_loader, model, num_steps):\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        acc = 0\n",
    "        total = 0\n",
    "        for inputs, targets in iter(data_loader):\n",
    "            inputs = inputs.to(device)\n",
    "            targets = targets.to(device)\n",
    "        \n",
    "            spk_rec, mem_rec = forward_pass(inputs, model, num_steps)\n",
    "            acc += SF.accuracy_rate(spk_rec, targets) * spk_rec.size(1)\n",
    "            total += spk_rec.size(1)\n",
    "    return acc / total    \n",
    "    \n",
    "class MyEEGSNNModel(nn.Module):\n",
    "    def __init__(self, n_outputs):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.conv1 = nn.Conv2d(3, 16, 5)\n",
    "        self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        self.conv2 = nn.Conv2d(16, 64, 5)\n",
    "        self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        self.fc1 = nn.Linear(64*47*47, n_outputs)\n",
    "        self.lif3 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        # self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=(1, 5))\n",
    "        # self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        # self.conv2 = nn.Conv2d(in_channels=16, out_channels=64, kernel_size=(1, 5))\n",
    "        # self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "        # self.fc1 = nn.Linear(64*14*157, n_outputs)\n",
    "        # self.lif3 = snn.Leaky(beta=beta, spike_grad=spike_grad)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.permute(0, 3, 1, 2)\n",
    "\n",
    "        mem1 = self.lif1.init_leaky()\n",
    "        mem2 = self.lif2.init_leaky()\n",
    "        mem3 = self.lif3.init_leaky()\n",
    "\n",
    "        cur1 = F.max_pool2d(self.conv1(x), kernel_size=(2, 2))\n",
    "        # cur1 = F.max_pool2d(self.conv1(x), kernel_size=(1, 2))\n",
    "        spk1, mem1 = self.lif1(cur1, mem1)\n",
    "         \n",
    "        cur2 = F.max_pool2d(self.conv2(spk1), kernel_size=(2, 2))\n",
    "        # cur2 = F.max_pool2d(self.conv2(spk1), kernel_size=(1, 2))\n",
    "        spk2, mem2 = self.lif2(cur2, mem2)\n",
    "        \n",
    "        cur3 = self.fc1(spk2.contiguous().view(batch_size, -1))\n",
    "        spk3, mem3 = self.lif3(cur3, mem3)\n",
    "\n",
    "        return spk3, mem3  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3106531f-2dd6-49cc-a6fc-78a204a11154",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                           | 0/16 [00:01<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([20, 32, 4])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "num_classes = 4\n",
    "model = MyEEGSNNModel(num_classes).to(device)\n",
    "\n",
    "for inputs, targets in tqdm(train_loader):\n",
    "    inputs = inputs.to(device)\n",
    "    targets = targets.to(device)\n",
    "\n",
    "    spk_rec, _= forward_pass(inputs, model, num_steps)\n",
    "    print(spk_rec.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0fe0ef1a-09eb-4f00-ba69-86fb0175202d",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA error: device-side assert triggered\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\n",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[13], line 8\u001b[0m\n\u001b[0;32m      6\u001b[0m num_classes \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m4\u001b[39m\n\u001b[0;32m      7\u001b[0m \u001b[38;5;66;03m# num_classes = 4\u001b[39;00m\n\u001b[1;32m----> 8\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[43mMyEEGSNNModel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnum_classes\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m      9\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m param \u001b[38;5;129;01min\u001b[39;00m model\u001b[38;5;241m.\u001b[39mparameters():\n\u001b[0;32m     10\u001b[0m     \u001b[38;5;28mprint\u001b[39m(param\u001b[38;5;241m.\u001b[39mrequires_grad)\n",
      "File \u001b[1;32mD:\\Anaconda\\envs\\trchenv3_10\\lib\\site-packages\\torch\\nn\\modules\\module.py:1174\u001b[0m, in \u001b[0;36mModule.to\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1171\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m   1172\u001b[0m             \u001b[38;5;28;01mraise\u001b[39;00m\n\u001b[1;32m-> 1174\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_apply\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconvert\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mD:\\Anaconda\\envs\\trchenv3_10\\lib\\site-packages\\torch\\nn\\modules\\module.py:780\u001b[0m, in \u001b[0;36mModule._apply\u001b[1;34m(self, fn, recurse)\u001b[0m\n\u001b[0;32m    778\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m recurse:\n\u001b[0;32m    779\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchildren():\n\u001b[1;32m--> 780\u001b[0m         \u001b[43mmodule\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_apply\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfn\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    782\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompute_should_use_set_data\u001b[39m(tensor, tensor_applied):\n\u001b[0;32m    783\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_has_compatible_shallow_copy_type(tensor, tensor_applied):\n\u001b[0;32m    784\u001b[0m         \u001b[38;5;66;03m# If the new tensor has compatible tensor type as the existing tensor,\u001b[39;00m\n\u001b[0;32m    785\u001b[0m         \u001b[38;5;66;03m# the current behavior is to change the tensor in-place using `.data =`,\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    790\u001b[0m         \u001b[38;5;66;03m# global flag to let the user control whether they want the future\u001b[39;00m\n\u001b[0;32m    791\u001b[0m         \u001b[38;5;66;03m# behavior of overwriting the existing tensor or not.\u001b[39;00m\n",
      "File \u001b[1;32mD:\\Anaconda\\envs\\trchenv3_10\\lib\\site-packages\\torch\\nn\\modules\\module.py:805\u001b[0m, in \u001b[0;36mModule._apply\u001b[1;34m(self, fn, recurse)\u001b[0m\n\u001b[0;32m    801\u001b[0m \u001b[38;5;66;03m# Tensors stored in modules are graph leaves, and we don't want to\u001b[39;00m\n\u001b[0;32m    802\u001b[0m \u001b[38;5;66;03m# track autograd history of `param_applied`, so we have to use\u001b[39;00m\n\u001b[0;32m    803\u001b[0m \u001b[38;5;66;03m# `with torch.no_grad():`\u001b[39;00m\n\u001b[0;32m    804\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[1;32m--> 805\u001b[0m     param_applied \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mparam\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    806\u001b[0m p_should_use_set_data \u001b[38;5;241m=\u001b[39m compute_should_use_set_data(param, param_applied)\n\u001b[0;32m    808\u001b[0m \u001b[38;5;66;03m# subclasses may have multiple child tensors so we need to use swap_tensors\u001b[39;00m\n",
      "File \u001b[1;32mD:\\Anaconda\\envs\\trchenv3_10\\lib\\site-packages\\torch\\nn\\modules\\module.py:1160\u001b[0m, in \u001b[0;36mModule.to.<locals>.convert\u001b[1;34m(t)\u001b[0m\n\u001b[0;32m   1153\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m convert_to_format \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m t\u001b[38;5;241m.\u001b[39mdim() \u001b[38;5;129;01min\u001b[39;00m (\u001b[38;5;241m4\u001b[39m, \u001b[38;5;241m5\u001b[39m):\n\u001b[0;32m   1154\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m t\u001b[38;5;241m.\u001b[39mto(\n\u001b[0;32m   1155\u001b[0m             device,\n\u001b[0;32m   1156\u001b[0m             dtype \u001b[38;5;28;01mif\u001b[39;00m t\u001b[38;5;241m.\u001b[39mis_floating_point() \u001b[38;5;129;01mor\u001b[39;00m t\u001b[38;5;241m.\u001b[39mis_complex() \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[0;32m   1157\u001b[0m             non_blocking,\n\u001b[0;32m   1158\u001b[0m             memory_format\u001b[38;5;241m=\u001b[39mconvert_to_format,\n\u001b[0;32m   1159\u001b[0m         )\n\u001b[1;32m-> 1160\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mt\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   1161\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1162\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdtype\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mis_floating_point\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mis_complex\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[0;32m   1163\u001b[0m \u001b[43m        \u001b[49m\u001b[43mnon_blocking\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1164\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1165\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[0;32m   1166\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mstr\u001b[39m(e) \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCannot copy out of meta tensor; no data!\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n",
      "\u001b[1;31mRuntimeError\u001b[0m: CUDA error: device-side assert triggered\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\n"
     ]
    }
   ],
   "source": [
    "n_epochs = 5\n",
    "learning_rate = 1e-3\n",
    "loss_history = []\n",
    "acc_history = []\n",
    "\n",
    "num_classes = 4\n",
    "# num_classes = 4\n",
    "model = MyEEGSNNModel(num_classes).to(device)\n",
    "for param in model.parameters():\n",
    "    print(param.requires_grad)\n",
    "criterion = SF.ce_rate_loss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, betas=(0.9, 0.999))\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    epoch_loss = 0.0\n",
    "    train_acc = 0.0\n",
    "    counter = 0\n",
    "    acc_co = 0\n",
    "    \n",
    "    for inputs, targets in tqdm(iter(train_loader)):\n",
    "        inputs = inputs.to(device)\n",
    "        targets = targets.to(device)\n",
    "        \n",
    "        model.train()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        spk_rec, _ = forward_pass(inputs, model, num_steps)\n",
    "        loss = criterion(spk_rec, targets)\n",
    "            \n",
    "        loss.backward()         \n",
    "        optimizer.step()  \n",
    "        epoch_loss += loss.item()\n",
    "        counter+=1\n",
    "\n",
    "        if counter % 75 == 0:\n",
    "        # if counter % 6 == 0:\n",
    "            train_acc += batch_acc(test_loader, model, num_steps)\n",
    "            acc_co+=1\n",
    "        \n",
    "    train_loss = epoch_loss / counter\n",
    "    train_acc = train_acc.item() / acc_co\n",
    "    # train_acc = train_acc.item() / acc_co\n",
    "    loss_history.append(train_loss)\n",
    "    acc_history.append(train_acc)\n",
    "    print(f'Epoch [{epoch + 1}/{n_epochs}], Loss: {train_loss:.4f}, Accuracy: {train_acc:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a75a80b-fe29-44e7-b8d5-6be2b8476949",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "for inputs, targets in tqdm(test_loader):\n",
    "    with torch.no_grad():\n",
    "        spk_rec, mem_rec = model(inputs, num_steps)\n",
    "        spk_count = torch.sum(spk_rec, dim=0)\n",
    "        outputs = F.softmax(spk_count)\n",
    "        \n",
    "    loss = criterion(outputs, targets)\n",
    "    epoch_loss += loss.item()\n",
    "print(f'Loss: {epoch_loss / len(test_loader):.4f}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "trch310",
   "language": "python",
   "name": "trch310"
  },
  "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}