cross_trial_EEG.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff7feea2-e04a-40d4-991f-88894bbfb7d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.11.0+cu113\n",
      "11.3\n",
      "weight_decay 0\n",
      "learning_rate 0.001\n",
      "depth 3\n",
      "side_depth 3\n",
      "random seed: 778539\n",
      "(76800, 81, 128, 1)\n",
      "(76800, 4, 128, 1)\n",
      "(76800, 4, 128, 1)\n",
      "total data (1280, 60, 89, 128)\n",
      "X_test shape (128, 60, 89, 128)\n",
      "for_train 1024\n",
      "for_valid 128\n",
      "cls_token\n",
      "side_cls_token\n",
      "to_patch_embedding.0.weight\n",
      "to_patch_embedding.0.bias\n",
      "to_patch_side_embedding.0.weight\n",
      "to_patch_side_embedding.0.bias\n",
      "to_patch_side_embedding.1.weight\n",
      "to_patch_side_embedding.1.bias\n",
      "position_enc.pos_table\n",
      "side_position_enc.pos_table\n",
      "transformer.layers.0.0.norm1.weight\n",
      "transformer.layers.0.0.norm1.bias\n",
      "transformer.layers.0.0.norm2.weight\n",
      "transformer.layers.0.0.norm2.bias\n",
      "transformer.layers.0.0.fn.to_qkv.weight\n",
      "transformer.layers.0.1.norm1.weight\n",
      "transformer.layers.0.1.norm1.bias\n",
      "transformer.layers.0.1.norm2.weight\n",
      "transformer.layers.0.1.norm2.bias\n",
      "transformer.layers.0.1.fn.net.0.weight\n",
      "transformer.layers.0.1.fn.net.0.bias\n",
      "transformer.layers.0.1.fn.net.3.weight\n",
      "transformer.layers.0.1.fn.net.3.bias\n",
      "transformer.layers.1.0.norm1.weight\n",
      "transformer.layers.1.0.norm1.bias\n",
      "transformer.layers.1.0.norm2.weight\n",
      "transformer.layers.1.0.norm2.bias\n",
      "transformer.layers.1.0.fn.to_qkv.weight\n",
      "transformer.layers.1.1.norm1.weight\n",
      "transformer.layers.1.1.norm1.bias\n",
      "transformer.layers.1.1.norm2.weight\n",
      "transformer.layers.1.1.norm2.bias\n",
      "transformer.layers.1.1.fn.net.0.weight\n",
      "transformer.layers.1.1.fn.net.0.bias\n",
      "transformer.layers.1.1.fn.net.3.weight\n",
      "transformer.layers.1.1.fn.net.3.bias\n",
      "transformer.layers.2.0.norm1.weight\n",
      "transformer.layers.2.0.norm1.bias\n",
      "transformer.layers.2.0.norm2.weight\n",
      "transformer.layers.2.0.norm2.bias\n",
      "transformer.layers.2.0.fn.to_qkv.weight\n",
      "transformer.layers.2.1.norm1.weight\n",
      "transformer.layers.2.1.norm1.bias\n",
      "transformer.layers.2.1.norm2.weight\n",
      "transformer.layers.2.1.norm2.bias\n",
      "transformer.layers.2.1.fn.net.0.weight\n",
      "transformer.layers.2.1.fn.net.0.bias\n",
      "transformer.layers.2.1.fn.net.3.weight\n",
      "transformer.layers.2.1.fn.net.3.bias\n",
      "transformer.side_layers.0.0.norm1.weight\n",
      "transformer.side_layers.0.0.norm1.bias\n",
      "transformer.side_layers.0.0.norm2.weight\n",
      "transformer.side_layers.0.0.norm2.bias\n",
      "transformer.side_layers.0.0.fn.to_qkv.weight\n",
      "transformer.side_layers.0.1.norm1.weight\n",
      "transformer.side_layers.0.1.norm1.bias\n",
      "transformer.side_layers.0.1.norm2.weight\n",
      "transformer.side_layers.0.1.norm2.bias\n",
      "transformer.side_layers.0.1.fn.net.0.weight\n",
      "transformer.side_layers.0.1.fn.net.0.bias\n",
      "transformer.side_layers.0.1.fn.net.3.weight\n",
      "transformer.side_layers.0.1.fn.net.3.bias\n",
      "transformer.side_layers.1.0.norm1.weight\n",
      "transformer.side_layers.1.0.norm1.bias\n",
      "transformer.side_layers.1.0.norm2.weight\n",
      "transformer.side_layers.1.0.norm2.bias\n",
      "transformer.side_layers.1.0.fn.to_qkv.weight\n",
      "transformer.side_layers.1.1.norm1.weight\n",
      "transformer.side_layers.1.1.norm1.bias\n",
      "transformer.side_layers.1.1.norm2.weight\n",
      "transformer.side_layers.1.1.norm2.bias\n",
      "transformer.side_layers.1.1.fn.net.0.weight\n",
      "transformer.side_layers.1.1.fn.net.0.bias\n",
      "transformer.side_layers.1.1.fn.net.3.weight\n",
      "transformer.side_layers.1.1.fn.net.3.bias\n",
      "transformer.side_layers.2.0.norm1.weight\n",
      "transformer.side_layers.2.0.norm1.bias\n",
      "transformer.side_layers.2.0.norm2.weight\n",
      "transformer.side_layers.2.0.norm2.bias\n",
      "transformer.side_layers.2.0.fn.to_qkv.weight\n",
      "transformer.side_layers.2.1.norm1.weight\n",
      "transformer.side_layers.2.1.norm1.bias\n",
      "transformer.side_layers.2.1.norm2.weight\n",
      "transformer.side_layers.2.1.norm2.bias\n",
      "transformer.side_layers.2.1.fn.net.0.weight\n",
      "transformer.side_layers.2.1.fn.net.0.bias\n",
      "transformer.side_layers.2.1.fn.net.3.weight\n",
      "transformer.side_layers.2.1.fn.net.3.bias\n",
      "transformer.cross_modal.to_qkv.weight\n",
      "transformer.cross_modal.to_qkv_side.weight\n",
      "transformer.cross_prenorm.norm1.weight\n",
      "transformer.cross_prenorm.norm1.bias\n",
      "transformer.cross_prenorm.norm2.weight\n",
      "transformer.cross_prenorm.norm2.bias\n",
      "transformer.cross_prenorm.fn.to_qkv.weight\n",
      "transformer.cross_prenorm.fn.to_qkv_side.weight\n",
      "transformer.cross_post.norm1.weight\n",
      "transformer.cross_post.norm1.bias\n",
      "transformer.cross_post.norm2.weight\n",
      "transformer.cross_post.norm2.bias\n",
      "transformer.cross_post.fn.net.0.weight\n",
      "transformer.cross_post.fn.net.0.bias\n",
      "transformer.cross_post.fn.net.3.weight\n",
      "transformer.cross_post.fn.net.3.bias\n",
      "transformer.back_cross_modal.to_qkv.weight\n",
      "transformer.back_cross_modal.to_qkv_side.weight\n",
      "transformer.back_cross_prenorm.norm1.weight\n",
      "transformer.back_cross_prenorm.norm1.bias\n",
      "transformer.back_cross_prenorm.norm2.weight\n",
      "transformer.back_cross_prenorm.norm2.bias\n",
      "transformer.back_cross_prenorm.fn.to_qkv.weight\n",
      "transformer.back_cross_prenorm.fn.to_qkv_side.weight\n",
      "transformer.back_cross_post.norm1.weight\n",
      "transformer.back_cross_post.norm1.bias\n",
      "transformer.back_cross_post.norm2.weight\n",
      "transformer.back_cross_post.norm2.bias\n",
      "transformer.back_cross_post.fn.net.0.weight\n",
      "transformer.back_cross_post.fn.net.0.bias\n",
      "transformer.back_cross_post.fn.net.3.weight\n",
      "transformer.back_cross_post.fn.net.3.bias\n",
      "laynorm.weight\n",
      "laynorm.bias\n",
      "side_laynorm.weight\n",
      "side_laynorm.bias\n",
      "mlp_head.0.weight\n",
      "mlp_head.0.bias\n",
      "['gpu']\n",
      "start training...\n",
      "Train Epoch: 1 [0/61440 (0%)]\tLoss: 0.683018,Accuracy: 39/64 (61%)\n",
      " \n",
      "Train Epoch: 1 [3200/61440 (5%)]\tLoss: 0.660205,Accuracy: 39/64 (61%)\n",
      " \n",
      "Train Epoch: 1 [6400/61440 (10%)]\tLoss: 0.659286,Accuracy: 38/64 (59%)\n",
      " \n",
      "Train Epoch: 1 [9600/61440 (16%)]\tLoss: 0.657536,Accuracy: 40/64 (62%)\n",
      " \n",
      "Train Epoch: 1 [12800/61440 (21%)]\tLoss: 0.683526,Accuracy: 35/64 (55%)\n",
      " \n",
      "Train Epoch: 1 [16000/61440 (26%)]\tLoss: 0.691736,Accuracy: 37/64 (58%)\n",
      " \n",
      "Train Epoch: 1 [19200/61440 (31%)]\tLoss: 0.635577,Accuracy: 43/64 (67%)\n",
      " \n",
      "Train Epoch: 1 [22400/61440 (36%)]\tLoss: 0.636642,Accuracy: 43/64 (67%)\n",
      " \n",
      "Train Epoch: 1 [25600/61440 (42%)]\tLoss: 0.680033,Accuracy: 37/64 (58%)\n",
      " \n",
      "Train Epoch: 1 [28800/61440 (47%)]\tLoss: 0.654159,Accuracy: 40/64 (62%)\n",
      " \n",
      "Train Epoch: 1 [32000/61440 (52%)]\tLoss: 0.611147,Accuracy: 41/64 (64%)\n",
      " \n",
      "Train Epoch: 1 [35200/61440 (57%)]\tLoss: 0.627352,Accuracy: 39/64 (61%)\n",
      " \n",
      "Train Epoch: 1 [38400/61440 (62%)]\tLoss: 0.555791,Accuracy: 47/64 (73%)\n",
      " \n",
      "Train Epoch: 1 [41600/61440 (68%)]\tLoss: 0.614310,Accuracy: 40/64 (62%)\n",
      " \n",
      "Train Epoch: 1 [44800/61440 (73%)]\tLoss: 0.588146,Accuracy: 38/64 (59%)\n",
      " \n",
      "Train Epoch: 1 [48000/61440 (78%)]\tLoss: 0.591043,Accuracy: 41/64 (64%)\n",
      " \n",
      "Train Epoch: 1 [51200/61440 (83%)]\tLoss: 0.556286,Accuracy: 48/64 (75%)\n",
      " \n",
      "Train Epoch: 1 [54400/61440 (89%)]\tLoss: 0.510632,Accuracy: 44/64 (69%)\n",
      " \n",
      "Train Epoch: 1 [57600/61440 (94%)]\tLoss: 0.555848,Accuracy: 52/64 (81%)\n",
      " \n",
      "Train Epoch: 1 [60800/61440 (99%)]\tLoss: 0.681113,Accuracy: 37/64 (58%)\n",
      " \n",
      "training in this epoch completed \n",
      "['gpu']\n",
      "test loss: tensor(0.0104, device='cuda:0')\n",
      "\n",
      "in Fold 0,Test set: Avg. loss: 0.0104, Accuracy: 4977/7680 (65%)\n",
      "\n",
      "['gpu']\n",
      "start training...\n",
      "Train Epoch: 2 [0/61440 (0%)]\tLoss: 0.522879,Accuracy: 48/64 (75%)\n",
      " \n",
      "Train Epoch: 2 [3200/61440 (5%)]\tLoss: 0.453939,Accuracy: 52/64 (81%)\n",
      " \n",
      "Train Epoch: 2 [6400/61440 (10%)]\tLoss: 0.491863,Accuracy: 50/64 (78%)\n",
      " \n",
      "Train Epoch: 2 [9600/61440 (16%)]\tLoss: 0.578731,Accuracy: 46/64 (72%)\n",
      " \n",
      "Train Epoch: 2 [12800/61440 (21%)]\tLoss: 0.593377,Accuracy: 44/64 (69%)\n",
      " \n",
      "Train Epoch: 2 [16000/61440 (26%)]\tLoss: 0.548489,Accuracy: 48/64 (75%)\n",
      " \n",
      "Train Epoch: 2 [19200/61440 (31%)]\tLoss: 0.633190,Accuracy: 42/64 (66%)\n",
      " \n",
      "Train Epoch: 2 [22400/61440 (36%)]\tLoss: 0.579380,Accuracy: 43/64 (67%)\n",
      " \n",
      "Train Epoch: 2 [25600/61440 (42%)]\tLoss: 0.430091,Accuracy: 51/64 (80%)\n",
      " \n",
      "Train Epoch: 2 [28800/61440 (47%)]\tLoss: 0.461118,Accuracy: 54/64 (84%)\n",
      " \n",
      "Train Epoch: 2 [32000/61440 (52%)]\tLoss: 0.454543,Accuracy: 47/64 (73%)\n",
      " \n",
      "Train Epoch: 2 [35200/61440 (57%)]\tLoss: 0.572193,Accuracy: 48/64 (75%)\n",
      " \n",
      "Train Epoch: 2 [38400/61440 (62%)]\tLoss: 0.446543,Accuracy: 47/64 (73%)\n",
      " \n",
      "Train Epoch: 2 [41600/61440 (68%)]\tLoss: 0.564814,Accuracy: 43/64 (67%)\n",
      " \n",
      "Train Epoch: 2 [44800/61440 (73%)]\tLoss: 0.428877,Accuracy: 50/64 (78%)\n",
      " \n",
      "Train Epoch: 2 [48000/61440 (78%)]\tLoss: 0.424809,Accuracy: 51/64 (80%)\n",
      " \n",
      "Train Epoch: 2 [51200/61440 (83%)]\tLoss: 0.503380,Accuracy: 45/64 (70%)\n",
      " \n",
      "Train Epoch: 2 [54400/61440 (89%)]\tLoss: 0.434982,Accuracy: 50/64 (78%)\n",
      " \n",
      "Train Epoch: 2 [57600/61440 (94%)]\tLoss: 0.487190,Accuracy: 47/64 (73%)\n",
      " \n",
      "Train Epoch: 2 [60800/61440 (99%)]\tLoss: 0.480961,Accuracy: 50/64 (78%)\n",
      " \n",
      "training in this epoch completed \n",
      "['gpu']\n",
      "test loss: tensor(0.0112, device='cuda:0')\n",
      "\n",
      "in Fold 0,Test set: Avg. loss: 0.0112, Accuracy: 4632/7680 (60%)\n",
      "\n",
      "['gpu']\n",
      "start training...\n",
      "Train Epoch: 3 [0/61440 (0%)]\tLoss: 0.469798,Accuracy: 51/64 (80%)\n",
      " \n",
      "Train Epoch: 3 [3200/61440 (5%)]\tLoss: 0.483222,Accuracy: 47/64 (73%)\n",
      " \n",
      "Train Epoch: 3 [6400/61440 (10%)]\tLoss: 0.512862,Accuracy: 46/64 (72%)\n",
      " \n",
      "Train Epoch: 3 [9600/61440 (16%)]\tLoss: 0.454500,Accuracy: 49/64 (77%)\n",
      " \n",
      "Train Epoch: 3 [12800/61440 (21%)]\tLoss: 0.375542,Accuracy: 53/64 (83%)\n",
      " \n",
      "Train Epoch: 3 [16000/61440 (26%)]\tLoss: 0.402184,Accuracy: 52/64 (81%)\n",
      " \n",
      "Train Epoch: 3 [19200/61440 (31%)]\tLoss: 0.478279,Accuracy: 48/64 (75%)\n",
      " \n",
      "Train Epoch: 3 [22400/61440 (36%)]\tLoss: 0.487637,Accuracy: 45/64 (70%)\n",
      " \n",
      "Train Epoch: 3 [25600/61440 (42%)]\tLoss: 0.370084,Accuracy: 53/64 (83%)\n",
      " \n",
      "Train Epoch: 3 [28800/61440 (47%)]\tLoss: 0.435331,Accuracy: 54/64 (84%)\n",
      " \n",
      "Train Epoch: 3 [32000/61440 (52%)]\tLoss: 0.495096,Accuracy: 49/64 (77%)\n",
      " \n",
      "Train Epoch: 3 [35200/61440 (57%)]\tLoss: 0.471906,Accuracy: 50/64 (78%)\n",
      " \n",
      "Train Epoch: 3 [38400/61440 (62%)]\tLoss: 0.494177,Accuracy: 50/64 (78%)\n",
      " \n",
      "Train Epoch: 3 [41600/61440 (68%)]\tLoss: 0.407762,Accuracy: 50/64 (78%)\n",
      " \n",
      "Train Epoch: 3 [44800/61440 (73%)]\tLoss: 0.307536,Accuracy: 55/64 (86%)\n",
      " \n",
      "Train Epoch: 3 [48000/61440 (78%)]\tLoss: 0.349146,Accuracy: 55/64 (86%)\n",
      " \n",
      "Train Epoch: 3 [51200/61440 (83%)]\tLoss: 0.458422,Accuracy: 49/64 (77%)\n",
      " \n",
      "Train Epoch: 3 [54400/61440 (89%)]\tLoss: 0.402915,Accuracy: 54/64 (84%)\n",
      " \n",
      "Train Epoch: 3 [57600/61440 (94%)]\tLoss: 0.387979,Accuracy: 50/64 (78%)\n",
      " \n",
      "Train Epoch: 3 [60800/61440 (99%)]\tLoss: 0.397949,Accuracy: 52/64 (81%)\n",
      " \n",
      "training in this epoch completed \n",
      "['gpu']\n",
      "test loss: tensor(0.0140, device='cuda:0')\n",
      "\n",
      "in Fold 0,Test set: Avg. loss: 0.0140, Accuracy: 4900/7680 (64%)\n",
      "\n",
      "['gpu']\n",
      "start training...\n",
      "Train Epoch: 4 [0/61440 (0%)]\tLoss: 0.452233,Accuracy: 51/64 (80%)\n",
      " \n",
      "Train Epoch: 4 [3200/61440 (5%)]\tLoss: 0.410104,Accuracy: 51/64 (80%)\n",
      " \n",
      "Train Epoch: 4 [6400/61440 (10%)]\tLoss: 0.364374,Accuracy: 50/64 (78%)\n",
      " \n",
      "Train Epoch: 4 [9600/61440 (16%)]\tLoss: 0.482164,Accuracy: 52/64 (81%)\n",
      " \n",
      "Train Epoch: 4 [12800/61440 (21%)]\tLoss: 0.306840,Accuracy: 56/64 (88%)\n",
      " \n",
      "Train Epoch: 4 [16000/61440 (26%)]\tLoss: 0.374909,Accuracy: 51/64 (80%)\n",
      " \n",
      "Train Epoch: 4 [19200/61440 (31%)]\tLoss: 0.233417,Accuracy: 59/64 (92%)\n",
      " \n",
      "Train Epoch: 4 [22400/61440 (36%)]\tLoss: 0.339291,Accuracy: 55/64 (86%)\n",
      " \n",
      "Train Epoch: 4 [25600/61440 (42%)]\tLoss: 0.348197,Accuracy: 51/64 (80%)\n",
      " \n",
      "Train Epoch: 4 [28800/61440 (47%)]\tLoss: 0.341603,Accuracy: 53/64 (83%)\n",
      " \n",
      "Train Epoch: 4 [32000/61440 (52%)]\tLoss: 0.252194,Accuracy: 59/64 (92%)\n",
      " \n",
      "Train Epoch: 4 [35200/61440 (57%)]\tLoss: 0.396798,Accuracy: 49/64 (77%)\n",
      " \n",
      "Train Epoch: 4 [38400/61440 (62%)]\tLoss: 0.341260,Accuracy: 54/64 (84%)\n",
      " \n",
      "Train Epoch: 4 [41600/61440 (68%)]\tLoss: 0.429531,Accuracy: 51/64 (80%)\n",
      " \n",
      "Train Epoch: 4 [44800/61440 (73%)]\tLoss: 0.363788,Accuracy: 52/64 (81%)\n",
      " \n",
      "Train Epoch: 4 [48000/61440 (78%)]\tLoss: 0.381737,Accuracy: 54/64 (84%)\n",
      " \n",
      "Train Epoch: 4 [51200/61440 (83%)]\tLoss: 0.409842,Accuracy: 51/64 (80%)\n",
      " \n",
      "Train Epoch: 4 [54400/61440 (89%)]\tLoss: 0.413279,Accuracy: 50/64 (78%)\n",
      " \n",
      "Train Epoch: 4 [57600/61440 (94%)]\tLoss: 0.414945,Accuracy: 54/64 (84%)\n",
      " \n",
      "Train Epoch: 4 [60800/61440 (99%)]\tLoss: 0.289569,Accuracy: 55/64 (86%)\n",
      " \n",
      "training in this epoch completed \n",
      "['gpu']\n",
      "test loss: tensor(0.0141, device='cuda:0')\n",
      "\n",
      "in Fold 0,Test set: Avg. loss: 0.0141, Accuracy: 4530/7680 (59%)\n",
      "\n",
      "['gpu']\n",
      "start training...\n",
      "Train Epoch: 5 [0/61440 (0%)]\tLoss: 0.394072,Accuracy: 52/64 (81%)\n",
      " \n",
      "Train Epoch: 5 [3200/61440 (5%)]\tLoss: 0.424650,Accuracy: 51/64 (80%)\n",
      " \n",
      "Train Epoch: 5 [6400/61440 (10%)]\tLoss: 0.326269,Accuracy: 56/64 (88%)\n",
      " \n",
      "Train Epoch: 5 [9600/61440 (16%)]\tLoss: 0.352084,Accuracy: 53/64 (83%)\n",
      " \n",
      "Train Epoch: 5 [12800/61440 (21%)]\tLoss: 0.489015,Accuracy: 50/64 (78%)\n",
      " \n",
      "Train Epoch: 5 [16000/61440 (26%)]\tLoss: 0.239433,Accuracy: 60/64 (94%)\n",
      " \n",
      "Train Epoch: 5 [19200/61440 (31%)]\tLoss: 0.249988,Accuracy: 58/64 (91%)\n",
      " \n",
      "Train Epoch: 5 [22400/61440 (36%)]\tLoss: 0.188931,Accuracy: 59/64 (92%)\n",
      " \n",
      "Train Epoch: 5 [25600/61440 (42%)]\tLoss: 0.323083,Accuracy: 52/64 (81%)\n",
      " \n",
      "Train Epoch: 5 [28800/61440 (47%)]\tLoss: 0.243330,Accuracy: 58/64 (91%)\n",
      " \n",
      "Train Epoch: 5 [32000/61440 (52%)]\tLoss: 0.211208,Accuracy: 58/64 (91%)\n",
      " \n",
      "Train Epoch: 5 [35200/61440 (57%)]\tLoss: 0.408270,Accuracy: 52/64 (81%)\n",
      " \n",
      "Train Epoch: 5 [38400/61440 (62%)]\tLoss: 0.192918,Accuracy: 61/64 (95%)\n",
      " \n",
      "Train Epoch: 5 [41600/61440 (68%)]\tLoss: 0.189200,Accuracy: 60/64 (94%)\n",
      " \n"
     ]
    }
   ],
   "source": [
    "# cross attention on channel 's softmax on dim=0\n",
    "import torch.nn.functional as F\n",
    "from torch import nn\n",
    "import torch\n",
    "import torchvision\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from einops import rearrange, repeat\n",
    "from einops.layers.torch import Rearrange\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import sys\n",
    "print(torch.__version__)\n",
    "print(torch.version.cuda)\n",
    "import math\n",
    "from sklearn import preprocessing\n",
    "from torch.autograd import Variable\n",
    "import pytorch_warmup as warmup\n",
    "import torch.nn.functional as F\n",
    "import random\n",
    "from sklearn.model_selection import KFold, train_test_split\n",
    "weight_decay=0\n",
    "learning_rate=1e-3\n",
    "depth=3\n",
    "side_depth=3\n",
    "print(\"weight_decay\",weight_decay)\n",
    "print(\"learning_rate\",learning_rate)\n",
    "print(\"depth\",depth)\n",
    "print(\"side_depth\",side_depth)\n",
    "'''\n",
    "WD=np.int(sys.argv[1])\n",
    "if WD == 2:\n",
    "    weight_decay=1e-2\n",
    "\n",
    "if WD == 3:\n",
    "    weight_decay=5e-2\n",
    "\n",
    "if WD == 4:\n",
    "    weight_decay=1e-1\n",
    "'''\n",
    "# seed setting\n",
    "seed = random.randint(0, 999999)\n",
    "seed = 778539\n",
    "\n",
    "print(\"random seed:\", seed)\n",
    "\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "# Define the number of folds for cross-validation\n",
    "num_folds = 9\n",
    "\n",
    "# Shuffle and split the dataset into k folds\n",
    "kf = KFold(n_splits=num_folds, shuffle=True, random_state=seed)\n",
    "\n",
    "\n",
    "# used as position embedding for PPS data\n",
    "class PositionalEncoding(nn.Module):\n",
    "\n",
    "    def __init__(self, d_hid, n_position=200):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "\n",
    "        self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))\n",
    "\n",
    "    def _get_sinusoid_encoding_table(self, n_position, d_hid):\n",
    "        ''' Sinusoid position encoding table '''\n",
    "\n",
    "        def get_position_angle_vec(position):\n",
    "            return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]\n",
    "\n",
    "        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])\n",
    "        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i\n",
    "        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1\n",
    "\n",
    "        return torch.FloatTensor(sinusoid_table).unsqueeze(0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x + self.pos_table[:, :x.size(1)].clone().detach()\n",
    "\n",
    "\n",
    "# used as position embedding for EEG data\n",
    "class PositionalEncoding2d(nn.Module):\n",
    "\n",
    "    def __init__(self, d_hid, d_wid, n_position=200):\n",
    "        super(PositionalEncoding2d, self).__init__()\n",
    "\n",
    "        self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid, d_wid))\n",
    "\n",
    "    def _get_sinusoid_encoding_table(self, n_position, d_hid, d_wid):\n",
    "\n",
    "        def get_position_angle_vec(position):\n",
    "\n",
    "            return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]\n",
    "\n",
    "        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])\n",
    "\n",
    "        sinusoid_table_x = np.zeros_like(sinusoid_table)\n",
    "        sinusoid_table_y = np.zeros_like(sinusoid_table)\n",
    "        # horizontal position coding list\n",
    "        sinusoid_table_x[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # for position as 2i, sin() is used\n",
    "        sinusoid_table_x[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # for position as 2i+1, cos() is used\n",
    "        # vertical position coding list\n",
    "        sinusoid_table_y[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # for position as 2j, sin() is used\n",
    "        sinusoid_table_y[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # for position as 2j+1, cos() is used\n",
    "\n",
    "        sinusoid_table = np.zeros((n_position, d_hid, d_wid))\n",
    "\n",
    "        for i in range(d_hid):\n",
    "            for j in range(d_wid):\n",
    "                for k in range(n_position):\n",
    "                    # multiply two lists to get position encoding table for EEG data\n",
    "                    sinusoid_table[k][i][j] = sinusoid_table_x[k, i] * sinusoid_table_y[k, j]\n",
    "\n",
    "        return torch.FloatTensor(sinusoid_table).unsqueeze(0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x + self.pos_table[:, :x.size(1)].clone().detach()\n",
    "\n",
    "\n",
    "# this module could implement LayerNorm for single-modal data and  multimodal data (seperately).\n",
    "# the Norm function is invoked before the latter \"fn\" function.\n",
    "class PreNorm(nn.Module):\n",
    "    def __init__(self, dim, side_embedding_dim, fn):\n",
    "        super().__init__()\n",
    "        self.norm1 = nn.LayerNorm(dim)\n",
    "        self.norm2 = nn.LayerNorm(side_embedding_dim)\n",
    "        self.fn = fn\n",
    "\n",
    "    def forward(self, inp, **kwargs):\n",
    "\n",
    "        if len(inp) == 2:\n",
    "            x = inp[0]\n",
    "\n",
    "            y = inp[1]\n",
    "\n",
    "            return self.fn((self.norm1(x), self.norm2(y)), **kwargs)\n",
    "        else:\n",
    "            x = inp\n",
    "\n",
    "            return self.fn(self.norm1(x), **kwargs)\n",
    "\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, dim, hidden_dim, dropout=0.):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(dim, hidden_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.Linear(hidden_dim, dim),\n",
    "            nn.Dropout(dropout)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    def __init__(self, dim, heads=8, dim_head=64, dropout=0.):\n",
    "        super().__init__()\n",
    "        inner_dim = dim_head * heads\n",
    "        project_out = not (heads == 1 and dim_head == dim)\n",
    "\n",
    "        self.heads = heads\n",
    "        self.scale = dim_head ** -0.5\n",
    "\n",
    "        self.attend = nn.Softmax(dim=-1)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)\n",
    "\n",
    "        self.to_out = nn.Sequential(\n",
    "\n",
    "            nn.Dropout(dropout)\n",
    "        ) if project_out else nn.Identity()\n",
    "\n",
    "    def forward(self, x):\n",
    "        qkv = self.to_qkv(x).chunk(3, dim=-1)\n",
    "        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)\n",
    "\n",
    "        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale\n",
    "\n",
    "        attn = self.attend(dots)\n",
    "        attn = self.dropout(attn)\n",
    "\n",
    "        out = torch.matmul(attn, v)\n",
    "        out = rearrange(out, 'b h n d -> b n (h d)')\n",
    "        return self.to_out(out)\n",
    "\n",
    "\n",
    "### the TACO Cross Attention block that aims at modeling temporal-spatial dependencies simultaneously ###\n",
    "class Cross_Attention(nn.Module):\n",
    "    def __init__(self, dim, side_embedding_dim, heads, dim_head, side_dim_head, cross_heads, cross_dim_head,\n",
    "                 dropout=0.):\n",
    "        super().__init__()\n",
    "        inner_dim = cross_dim_head * cross_heads\n",
    "\n",
    "        project_out = not (cross_heads == 1 and cross_dim_head == dim)\n",
    "\n",
    "        self.heads = cross_heads\n",
    "        self.scale = (dim_head * side_dim_head) ** -0.5  # defined by head_dim of EEG and side_head_dim of PPS\n",
    "\n",
    "        self.attend_chan = nn.Softmax(dim=-1)\n",
    "        self.attend_token = nn.Softmax(dim=-2)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)\n",
    "\n",
    "        self.to_qkv_side = nn.Linear(side_embedding_dim, inner_dim * 3, bias=False)\n",
    "\n",
    "        self.to_out = nn.Sequential(\n",
    "\n",
    "            nn.Dropout(dropout)\n",
    "        ) if project_out else nn.Identity()\n",
    "\n",
    "    def forward(self, inp):\n",
    "        x = inp[0]\n",
    "        y = inp[1]\n",
    "\n",
    "        qkv = self.to_qkv(x).chunk(3, dim=-1)\n",
    "        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)\n",
    "\n",
    "        qkv_side = self.to_qkv_side(y).chunk(3, dim=-1)\n",
    "        qs, ks, vs = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv_side)\n",
    "\n",
    "        sig_len = x.size()[1]\n",
    "        # print(\"sig_len\",sig_len)\n",
    "\n",
    "        chan_dots = torch.matmul(q.transpose(-1, -2), ks) * self.scale\n",
    "        token_dots = torch.matmul(qs, k.transpose(-1, -2)) * ((sig_len) ** -0.5)\n",
    "\n",
    "        chan_attn = self.attend_chan(chan_dots)\n",
    "        chan_attn = self.dropout(chan_attn)\n",
    "\n",
    "        chan_out = torch.matmul(v, chan_attn)\n",
    "\n",
    "        token_attn = self.attend_token(token_dots)\n",
    "        token_attn = self.dropout(token_attn)\n",
    "\n",
    "        out = torch.matmul(token_attn, chan_out)\n",
    "        out = rearrange(out, 'b h n d -> b n (h d)')\n",
    "        qs_res = rearrange(qs, 'b h n d -> b n (h d)')\n",
    "        v_res = rearrange(v, 'b h n d -> b n (h d)')\n",
    "        return v_res, out\n",
    "\n",
    "\n",
    "class Transformer(nn.Module):\n",
    "    def __init__(self, embedding_dim, side_embedding_dim, depth, side_depth, heads, side_heads, cross_heads, cross_dim_head,\n",
    "                 dim_head, side_dim_head, mlp_dim, side_mlp_dim, dropout=0.):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList([])\n",
    "        self.side_layers = nn.ModuleList([])\n",
    "        self.cross_modal = Cross_Attention(dim=embedding_dim, side_embedding_dim=side_embedding_dim, heads=heads,\n",
    "                                           dim_head=dim_head, side_dim_head=side_dim_head, cross_heads=cross_heads,\n",
    "                                           cross_dim_head=cross_dim_head, dropout=0.)\n",
    "        self.cross_prenorm = PreNorm(embedding_dim, side_embedding_dim, self.cross_modal)\n",
    "        self.cross_post = PreNorm(embedding_dim, side_embedding_dim, FeedForward(embedding_dim, mlp_dim, dropout=dropout))\n",
    "        \n",
    "        self.back_cross_modal = Cross_Attention(dim=side_embedding_dim, side_embedding_dim=embedding_dim, heads=heads,\n",
    "                                           dim_head=dim_head, side_dim_head=side_dim_head, cross_heads=cross_heads,\n",
    "                                           cross_dim_head=cross_dim_head, dropout=0.)\n",
    "        self.back_cross_prenorm = PreNorm(side_embedding_dim, embedding_dim, self.back_cross_modal)\n",
    "        self.back_cross_post = PreNorm(embedding_dim, embedding_dim, FeedForward(embedding_dim, mlp_dim, dropout=dropout))\n",
    "\n",
    "        for _ in range(depth):\n",
    "            self.layers.append(nn.ModuleList([\n",
    "                PreNorm(embedding_dim, side_embedding_dim, Attention(embedding_dim, heads=heads, dim_head=dim_head, dropout=dropout)),\n",
    "                PreNorm(embedding_dim, side_embedding_dim, FeedForward(embedding_dim, mlp_dim, dropout=dropout))\n",
    "            ]))\n",
    "\n",
    "        for _ in range(side_depth):\n",
    "            self.side_layers.append(nn.ModuleList([\n",
    "                PreNorm(side_embedding_dim, side_embedding_dim,\n",
    "                        Attention(side_embedding_dim, heads=side_heads, dim_head=side_dim_head, dropout=dropout)),\n",
    "                PreNorm(side_embedding_dim, side_embedding_dim,\n",
    "                        FeedForward(side_embedding_dim, side_mlp_dim, dropout=dropout))\n",
    "            ]))\n",
    "\n",
    "    def forward(self, inp):\n",
    "\n",
    "        x = inp[0]  # EEG modal\n",
    "        y = inp[1]  # PPS modal\n",
    "        for attn, ff in self.layers:\n",
    "            x = attn(x) + x\n",
    "            x = ff(x) + x\n",
    "        for attn, ff in self.side_layers:\n",
    "            y = attn(y) + y\n",
    "            y = ff(y) + y\n",
    "        \n",
    "        # seperately pre-normalize x and y, then conduct TACO cross attention on x and y,\n",
    "        # return \"out\" as fusion output and \"vs\",the multi-head value matrix as residual term\n",
    "        vs, out = self.cross_prenorm((x, y))\n",
    "        # residual connection\n",
    "        k = vs + out\n",
    "        # pre-normalize embedding k, and conduct Feedforward on it , finally conduct residual connection on k and output of Feedforward\n",
    "        z = k + self.cross_post(k)\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        return z\n",
    "\n",
    "\n",
    "class ViT(nn.Module):\n",
    "    def __init__(self, *, num_classes=2, height=9, width=9, embedding_dim=1024, side_dim=8, side_mid_dim=64,\n",
    "                 side_embedding_dim=128, depth=8, side_depth=4, heads=8, side_heads=4, cross_heads=8, dim_head=64,\n",
    "                 side_dim_head=32, cross_dim_head=64, mlp_dim=2048, side_mlp_dim=256, pool='cls', dropout=0.,\n",
    "                 emb_dropout=0.):\n",
    "        super().__init__()\n",
    "\n",
    "        self.side_dim = side_dim\n",
    "        self.height = height\n",
    "        self.width = width\n",
    "        patch_dim = height * width\n",
    "        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'\n",
    "\n",
    "        self.to_patch_embedding = nn.Sequential(\n",
    "\n",
    "            nn.Linear(patch_dim, embedding_dim)\n",
    "        )\n",
    "        self.to_patch_side_embedding = nn.Sequential(\n",
    "\n",
    "            nn.Linear(side_dim, side_mid_dim),\n",
    "            nn.Linear(side_mid_dim, side_embedding_dim)\n",
    "        )\n",
    "        self.position_enc = PositionalEncoding2d(height, width, n_position=200)\n",
    "        self.side_position_enc = PositionalEncoding(side_embedding_dim, n_position=200)\n",
    "        self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim))\n",
    "        self.side_cls_token = nn.Parameter(torch.randn(1, 1, side_embedding_dim))\n",
    "        self.dropout = nn.Dropout(emb_dropout)\n",
    "        self.transformer = Transformer(embedding_dim, side_embedding_dim, depth, side_depth, heads, side_heads,\n",
    "                                       cross_heads, cross_dim_head, dim_head, side_dim_head, mlp_dim, side_mlp_dim,\n",
    "                                       dropout=dropout)\n",
    "        self.pool = pool\n",
    "        self.to_latent = nn.Identity()\n",
    "        self.laynorm = nn.LayerNorm(embedding_dim)\n",
    "        self.side_laynorm = nn.LayerNorm(side_embedding_dim)\n",
    "        self.mlp_head = nn.Sequential(\n",
    "            nn.Linear((embedding_dim), num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, img):\n",
    "        img = img.permute(0, 2, 1)\n",
    "        y = img[:, :, -self.side_dim:]\n",
    "        img = img[:, :, :-self.side_dim]\n",
    "\n",
    "        # reshape the 81 channel to 9X9 matrix\n",
    "        x = img.reshape((img.shape[0], img.shape[1], self.height, self.width))\n",
    "        # 2D position encoding\n",
    "        x = self.position_enc(x)\n",
    "        # reshape 9X9 matrix back into 81 channels\n",
    "        x = x.reshape(x.shape[0], x.shape[1], self.height * self.width)\n",
    "\n",
    "        # Linear embedding for EEG and PPS\n",
    "        x = self.to_patch_embedding(x)\n",
    "        y = self.to_patch_side_embedding(y)\n",
    "\n",
    "        # concatenate class token as the first Token in sequence\n",
    "        b, n, _ = x.shape\n",
    "        cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b=b)\n",
    "        x = torch.cat((cls_tokens, x), dim=1)\n",
    "        x = self.dropout(x)\n",
    "        x = self.laynorm(x)\n",
    "\n",
    "        b_s, n_s, _s = y.shape\n",
    "        side_cls_tokens = repeat(self.side_cls_token, '1 1 d -> b 1 d', b=b)\n",
    "        y = torch.cat((side_cls_tokens, y), dim=1)\n",
    "        y = self.side_position_enc(y)\n",
    "        y = self.dropout(y)\n",
    "        y = self.side_laynorm(y)\n",
    "\n",
    "        # seperate transformer encoders and TACO Cross Attention block\n",
    "        x = self.transformer((x, y))\n",
    "        # integrating emotion-related information from multi tokens in Class Token\n",
    "        x = x.mean(dim=1) if self.pool == 'mean' else x[:, 0,:]#b n c \n",
    "\n",
    "        x = self.to_latent(x)\n",
    "        return self.mlp_head(x)\n",
    "\n",
    "\n",
    "def train(epoch):\n",
    "    Net = net\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.Adam(Net.parameters(), lr=learning_rate,weight_decay=weight_decay)\n",
    "    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=-1)\n",
    "    warmup_scheduler = warmup.UntunedLinearWarmup(optimizer)\n",
    "\n",
    "    '''cuda accelerate'''\n",
    "    device = ['gpu' if torch.cuda.is_available() else 'cpu']\n",
    "    print(device)\n",
    "    if device[0] == 'gpu':\n",
    "        criterion.cuda()\n",
    "        Net.cuda()\n",
    "\n",
    "    print('start training...')\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "\n",
    "        correct = 0\n",
    "        if device[0] == 'gpu':\n",
    "            data = data.cuda()\n",
    "            target = target.cuda()\n",
    "\n",
    "            data = data.to(torch.float32)\n",
    "            target = target.to(torch.float32)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        output = Net(data)\n",
    "        loss = criterion(output, target.long())\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        with warmup_scheduler.dampening():\n",
    "            lr_scheduler.step()\n",
    "        pred = output.data.max(1, keepdim=True)[1]\n",
    "        correct += pred.eq(target.data.view_as(pred)).sum()\n",
    "\n",
    "        if batch_idx % log_interval == 0:\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f},Accuracy: {}/{} ({:.0f}%)\\n '.format(\n",
    "                epoch, batch_idx * len(data), len(train_loader.dataset),\n",
    "                       100. * batch_idx / len(train_loader), loss.item(), correct, train_size,\n",
    "                       100. * correct / train_size))\n",
    "            train_losses.append(loss.item())\n",
    "            train_counter.append(\n",
    "                (batch_idx * train_size) + ((epoch - 1) * len(train_loader.dataset)))\n",
    "    print('training in this epoch completed ')\n",
    "\n",
    "\n",
    "def test(real_test):\n",
    "    network = net\n",
    "    network.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    '''cuda加速'''\n",
    "    device = ['gpu' if torch.cuda.is_available() else 'cpu']\n",
    "    print(device)\n",
    "    if device == 'gpu':\n",
    "        criterion.cuda()\n",
    "        network.cuda()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (data, target) in enumerate(test_loader):\n",
    "\n",
    "            data = data.to(torch.float32)\n",
    "            target = target.to(torch.float32)\n",
    "\n",
    "            if device[0] == 'gpu':\n",
    "                data = data.cuda()\n",
    "                target = target.cuda()\n",
    "                data = data.to(torch.float32)\n",
    "                target = target.to(torch.float32)\n",
    "\n",
    "            output = network(data)\n",
    "\n",
    "            test_loss += criterion(output, target.long())\n",
    "            pred = output.data.max(1, keepdim=True)[1]\n",
    "\n",
    "            correct += pred.eq(target.data.view_as(pred)).sum()\n",
    "\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    print(\"test loss:\", test_loss)\n",
    "    test_losses.append(test_loss)\n",
    "    if real_test==True:\n",
    "        print(\"running in test dataset with model trained on the last fold split:\")\n",
    "        print('\\nin Test set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(test_loss, correct,\n",
    "                                                                                         len(test_loader.dataset),\n",
    "                                                                                         100. * correct / len(\n",
    "                                                                                             test_loader.dataset)))\n",
    "    if real_test==False:\n",
    "        print('\\nin Fold {},Test set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(fold,\n",
    "                                                                                         test_loss, correct,\n",
    "                                                                                         len(test_loader.dataset),\n",
    "                                                                                         100. * correct / len(\n",
    "                                                                                             test_loader.dataset)))\n",
    "\n",
    "    #conf_matrix = confusion_matrix(pred, target.data.view_as(pred))\n",
    "    #conf_matrix = conf_matrix.cpu()\n",
    "    acc_this_fold = 100. * correct / len(test_loader.dataset)\n",
    "    if real_test == False:\n",
    "        # '''4.save model'''\n",
    "        torch.save(network, 'deap_' + str(fold) + '.pt')\n",
    "    return acc_this_fold\n",
    "\n",
    "\n",
    "def confusion_matrix(preds, labels):\n",
    "    for p, t in zip(preds.long(), labels.long()):\n",
    "        conf_matrix[p, t] += 1\n",
    "    return conf_matrix\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # get original data\n",
    "    data_eeg = np.load(\"/root/autodl-tmp/datalist_real.npy\")\n",
    "    data_eog = np.load(\"/root/autodl-tmp/datalist_eog_deap.npy\")\n",
    "    data_emg = np.load(\"/root/autodl-tmp/datalist_emg_deap.npy\")\n",
    "    print(data_eeg.shape)# (32 40 60) , 81 128\n",
    "    print(data_eog.shape)\n",
    "    print(data_emg.shape)\n",
    "\n",
    "    data_eeg = np.reshape(data_eeg, (1280,60, 81, 128))\n",
    "    data_eog = np.reshape(data_eog, (1280,60, 4, 128))\n",
    "    data_emg = np.reshape(data_emg, (1280,60, 4, 128))\n",
    "\n",
    "    data = np.concatenate((data_eeg, data_eog), axis=2)\n",
    "    data = np.concatenate((data, data_emg), axis=2)\n",
    "    data = np.squeeze(data)\n",
    "    print(\"total data\", data.shape)\n",
    "    y = np.load(\"/root/autodl-tmp/dataset/arousal.npy\")\n",
    "    y = np.reshape(y, (1280,60))\n",
    "\n",
    "    y = y.astype(np.float64)\n",
    "    data = data.astype(np.float64)\n",
    "\n",
    "    # take test dataset out of whole dataset\n",
    "   \n",
    "    data, X_test, y, y_test = train_test_split(data, y, test_size=0.1, random_state=seed)# which is 1/32 \n",
    "    print(\"X_test shape\",X_test.shape)\n",
    "    \n",
    "    X = torch.Tensor(data)\n",
    "    y = torch.Tensor(y)\n",
    "    X_test = torch.Tensor(X_test)\n",
    "    y_test = torch.Tensor(y_test)\n",
    "    train_size =64\n",
    "    test_size = 64\n",
    "    accuracy_log = []\n",
    "\n",
    "    for fold, (train_indices, valid_indices) in enumerate(kf.split(X)):\n",
    "        train_data, validation_data = X[train_indices], X[valid_indices]\n",
    "        train_labels, validation_labels = y[train_indices], y[valid_indices]\n",
    "        print(\"for_train\", len(train_data))\n",
    "        print(\"for_valid\", len(validation_data))\n",
    "        \n",
    "        \n",
    "        # reshape datas into shape: samples, channel ,timestamp\n",
    "        train_data = np.reshape(train_data, (61440, 89, 128))#30*2400\n",
    "        train_labels = np.reshape(train_labels, (61440,))\n",
    "\n",
    "\n",
    "        \n",
    "        validation_data=np.reshape(validation_data,(7680,89,128))\n",
    "        validation_labels=np.reshape(validation_labels,(7680,))\n",
    "\n",
    "\n",
    "        train_dataset = TensorDataset(train_data, train_labels)\n",
    "        test_dataset = TensorDataset(validation_data, validation_labels)\n",
    "        train_loader = DataLoader(train_dataset, batch_size=train_size, shuffle=True)\n",
    "        test_loader = DataLoader(test_dataset, batch_size=test_size, shuffle=False)\n",
    "\n",
    "        log_interval = 50\n",
    "        n_epochs=100\n",
    "        train_losses = []\n",
    "        train_counter = []\n",
    "        test_losses = []\n",
    "        test_counter = [i*len(train_loader.dataset) for i in range(n_epochs + 1)]\n",
    "\n",
    "        net=ViT(\n",
    "            num_classes = 2,\n",
    "            height=9,\n",
    "            width=9,\n",
    "            embedding_dim=128,\n",
    "            side_dim=8,\n",
    "            side_mid_dim=64,#for linear embedding \n",
    "            side_embedding_dim=128,#for linear embedding \n",
    "\n",
    "\n",
    "            depth = depth,\n",
    "            side_depth=side_depth,\n",
    "            heads = 4,\n",
    "            side_heads=4,\n",
    "            cross_heads=4,\n",
    "\n",
    "            dim_head=32,\n",
    "            side_dim_head=32,\n",
    "            cross_dim_head=32,\n",
    "\n",
    "\n",
    "            mlp_dim = 128,# used in Feedforward \n",
    "            side_mlp_dim=128,# used in Feedforward  \n",
    "            dropout = 0.2,\n",
    "            emb_dropout = 0.2,\n",
    "            \n",
    "        )\n",
    "        for name in net.state_dict():\n",
    "          print(name)\n",
    "\n",
    "\n",
    "        for epoch in range(1, n_epochs + 1):\n",
    "          train(epoch)\n",
    "          acc_this_fold=test(real_test=False)\n",
    "        accuracy_log.append(acc_this_fold)\n",
    "\n",
    "        #rint(\"conf_mat\",conf_matrix)\n",
    "        test_losses=torch.tensor(test_losses,device=\"cpu\")\n",
    "        import matplotlib.pyplot as plt\n",
    "        fig = plt.figure()\n",
    "        plt.plot(train_counter, train_losses, color='blue')\n",
    "        print(np.array(test_counter[1:]).shape)\n",
    "        print(np.array(test_losses).shape)\n",
    "        plt.scatter(np.array(test_counter[1:]), np.array(test_losses), color='red')\n",
    "        plt.legend(['Train Loss', 'Test Loss'], loc='upper right')\n",
    "        plt.xlabel('number of training examples seen')\n",
    "        plt.ylabel('negative log likelihood loss')\n",
    "        plt.savefig(\"train loss of fold\"+str(fold))\n",
    "        plt.show()\n",
    "   \n",
    "    print(\"acc in 31 fold cross validation:\", accuracy_log)\n",
    "    print(\"average acc in 31-fold cross subject validation\",torch.mean(accuracy_log))\n",
    "    \n",
    "    \n",
    "\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}