things_image_encoder.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0678cb22-26ed-4b31-8f9c-01df86152ef5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/hdd/yuchen\n"
     ]
    }
   ],
   "source": [
    "%cd /hdd/yuchen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e6efe314-dd30-40f3-a29e-7f72efaa9c08",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "import glob\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from torch.utils.data import Dataset, BatchSampler\n",
    "from sklearn.model_selection import train_test_split\n",
    "from tqdm.notebook import tqdm\n",
    "import random\n",
    "import nltk, re, string, collections\n",
    "import scipy \n",
    "import IProgress \n",
    "from ipywidgets import FloatProgress\n",
    "from nilearn import connectome\n",
    "from sklearn.preprocessing import Normalizer, OrdinalEncoder, OneHotEncoder, StandardScaler\n",
    "from torch.nn import Linear, ReLU, Dropout\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCNConv, TopKPooling, BatchNorm, global_mean_pool\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.loader import DataLoader\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7784657d-f060-432a-af91-671b7d40eae3",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 42\n",
    "random.seed(seed)\n",
    "os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed(seed)\n",
    "torch.backends.cudnn.benchmark = False\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "33c4a696-089d-4a22-b678-4a696ad38acf",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_dir = '/hdd/yuchen/cs292/test/nsdfeat/init_latent/'\n",
    "files = glob.glob(save_dir + '*')\n",
    "t = [int(f.split('/')[-1].split('.')[0]) for f in files]\n",
    "t.sort()\n",
    "feats_te_avg = []\n",
    "for file_idx in t: feats_te_avg.append(np.load(f'{save_dir}{file_idx:06}.npy'))\n",
    "feats_te_avg = np.array(feats_te_avg)\n",
    "\n",
    "save_dir = '/hdd/yuchen/cs292/train/nsdfeat/init_latent/'\n",
    "files = glob.glob(save_dir + '*')\n",
    "t = [int(f.split('/')[-1].split('.')[0]) for f in files]\n",
    "t.sort()\n",
    "feats_tr = []\n",
    "for file_idx in t: feats_tr.append(np.load(f'{save_dir}{file_idx:06}.npy'))\n",
    "feats_tr = np.array(feats_tr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "725c995a-c465-4fda-8e87-f22f3b900724",
   "metadata": {},
   "outputs": [],
   "source": [
    "from os.path import join as pjoin\n",
    "\n",
    "betas_csv_dir = '/hdd/yuchen/cs292/fmri/betas_csv/'\n",
    "sub = '01'\n",
    "    \n",
    "data_file = pjoin(betas_csv_dir, f'sub-{sub}_ResponseData.h5')\n",
    "responses = pd.read_hdf(data_file)\n",
    "vox_f = pjoin(betas_csv_dir, f'sub-{sub}_VoxelMetadata.csv')\n",
    "voxdata = pd.read_csv(vox_f)\n",
    "stim_f = pjoin(betas_csv_dir, f'sub-{sub}_StimulusMetadata.csv')\n",
    "stimdata = pd.read_csv(stim_f)\n",
    "\n",
    "train_idx = stimdata[stimdata['trial_type'] == 'train'].index.tolist()\n",
    "test_idx = stimdata[stimdata['trial_type'] == 'test'].index.tolist()\n",
    "    \n",
    "train_fmri = responses[train_idx]\n",
    "test_fmri = responses[test_idx]\n",
    "    \n",
    "train_labels = stimdata[stimdata['trial_type'] == 'train']['stimulus']\n",
    "test_labels = stimdata[stimdata['trial_type'] == 'test']['stimulus']\n",
    "    \n",
    "roi_idx = voxdata[(voxdata['V1'] == 1) ]['voxel_id'].tolist()\n",
    "train_fmri_temp = train_fmri.iloc[roi_idx]\n",
    "test_fmri_temp = test_fmri.iloc[roi_idx]\n",
    "\n",
    "roi_neuron = {'V1': train_fmri_temp.shape[0]}\n",
    "\n",
    "for roi in ['V2', 'V3', 'hV4']:\n",
    "    roi_idx = voxdata[(voxdata[roi] == 1)]['voxel_id'].tolist()\n",
    "    train_fmri_temp = pd.concat([train_fmri_temp, train_fmri.iloc[roi_idx]])\n",
    "    test_fmri_temp = pd.concat([test_fmri_temp, test_fmri.iloc[roi_idx]])\n",
    "    roi_neuron[roi] = train_fmri.iloc[roi_idx].shape[0]\n",
    "    \n",
    "train_fmri = train_fmri_temp\n",
    "test_fmri = test_fmri_temp\n",
    "\n",
    "# roi_idx = voxdata[(voxdata['V1'] == 1) | (voxdata['V2'] == 1) | (voxdata['V3'] == 1) | (voxdata['hV4'] == 1) ]['voxel_id'].tolist()\n",
    "# train_fmri = train_fmri.iloc[roi_idx]\n",
    "# test_fmri = test_fmri.iloc[roi_idx]\n",
    "\n",
    "test = stimdata[stimdata['trial_type'] == 'test']\n",
    "test_fmri_avg = []\n",
    "test_labels = test['stimulus'].unique()\n",
    "for stim in test_labels:\n",
    "    repeated_trial = test[test['stimulus'] == stim].index.tolist()\n",
    "    test_fmri_avg.append(test_fmri[repeated_trial].mean(axis=1).to_numpy())\n",
    "test_fmri_avg = np.array(test_fmri_avg)\n",
    "\n",
    "del responses, voxdata, stimdata\n",
    "# train_fmri, test_fmri =  train_fmri.to_numpy(), test_fmri.to_numpy()\n",
    "train_img_label_all, test_img_label_all = train_labels.tolist(), test_labels.tolist()\n",
    "\n",
    "betas_tr = train_fmri.T\n",
    "betas_te_avg = test_fmri_avg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ac1ee82f-b8ca-457d-a2d4-77ecb28ccd81",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'V1': 1049, 'V2': 774, 'V3': 762, 'hV4': 613}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "roi_neuron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3402250a-126b-4f07-bdbe-6bd136bf6115",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100, 3198)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_fmri_avg.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d5f3bb46-a313-43a8-890b-eba7d35bfc50",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "feats_tr (8640, 6400)\n",
      "betas_tr (8640, 3198)\n",
      "feats_te_avg (100, 6400)\n",
      "betas_te_avg (100, 3198)\n"
     ]
    }
   ],
   "source": [
    "print('feats_tr', feats_tr.shape)\n",
    "print('betas_tr', betas_tr.shape)\n",
    "\n",
    "print('feats_te_avg', feats_te_avg.shape)\n",
    "print('betas_te_avg', betas_te_avg.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fafd8421-9704-4e84-80fc-a32cda7ade94",
   "metadata": {},
   "outputs": [],
   "source": [
    "betas_tr = np.array(betas_tr)\n",
    "betas_te_avg_ = np.array(betas_te_avg)\n",
    "\n",
    "betas_tr_ = (betas_tr - betas_tr.mean()) / betas_tr.std()\n",
    "betas_te_avg_ = (betas_te_avg - betas_te_avg.mean()) / betas_te_avg.std()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "06b8a638-80da-499a-9ef1-8766a730ca0b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(8640, 3198)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "betas_tr_.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e180a496-c301-4b91-b54a-7573d02ff5ca",
   "metadata": {},
   "source": [
    "## goal: 1.8787439"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "389b5d54-aacc-4db9-9c1f-98ff7af7823d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse, os\n",
    "import numpy as np\n",
    "from himalaya.backend import set_backend\n",
    "from himalaya.ridge import RidgeCV\n",
    "from himalaya.scoring import correlation_score\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "alpha = [0.000001,0.00001,0.0001,0.001,0.01, 0.1, 1]\n",
    "\n",
    "ridge = RidgeCV(alphas=alpha)\n",
    "\n",
    "preprocess_pipeline = make_pipeline(\n",
    "        StandardScaler(with_mean=True, with_std=True),\n",
    "    )\n",
    "pipeline = make_pipeline(\n",
    "        preprocess_pipeline,\n",
    "        ridge,\n",
    ")    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7914f467-bc01-4c6b-bd24-312df879bfb2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.74321836"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = betas_tr\n",
    "Y = feats_tr\n",
    "X_te = betas_te_avg\n",
    "pipeline.fit(X, Y)\n",
    "scores = pipeline.predict(X_te)\n",
    "from sklearn.metrics import mean_squared_error\n",
    "mean_squared_error(scores,feats_te_avg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "39246964-b95c-45ef-83bd-4980dd451014",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save('/hdd/yuchen/things_baseline_latent_init.npy',scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da88f155-b5ed-4fa5-9e1e-b399a746db76",
   "metadata": {},
   "source": [
    "## gcn best: 0.55675"
   ]
  },
  {
   "cell_type": "raw",
   "id": "89005bec-c620-4078-a34b-7a55dc0c1e04",
   "metadata": {},
   "source": [
    "ideas inspired by:\n",
    "https://www.mdpi.com/2076-3425/12/10/1394\n",
    "- Decoding Visual fMRI Stimuli from Human Brain Based on Graph Convolutional Neural Network\n",
    "https://www.bendsawyer.com/wp-content/uploads/2023/07/Saeidi-et-al-Decoding-Task-Based-fMRI-Data-with-Graph-Neural-Networks-Considering-Individual-Differences.pdf\n",
    "- Decoding Task-Based fMRI Data with Graph Neural Networks, Considering Individual Differences\n",
    "\n",
    "code adapted from:\n",
    "https://github.com/krzysztoffiok/gnn-classification-pipeline/blob/master/gnn_fw/models.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fa2b706b-ff9e-4ad1-bd36-df596ac7e453",
   "metadata": {},
   "outputs": [],
   "source": [
    "betas_tr_temp = betas_tr_[:, :roi_neuron['V1']]\n",
    "num_feature = 1375\n",
    "zeros = np.zeros((betas_tr_temp.shape[0], num_feature))\n",
    "zeros[:, :betas_tr_temp.shape[-1]] = betas_tr_temp\n",
    "betas_tr_temp = zeros  # (8640, 1573)\n",
    "\n",
    "for roi in [2,3]:\n",
    "    temp = betas_tr_[:, sum(list(roi_neuron.values())[:roi-1]):sum(list(roi_neuron.values())[:roi])]\n",
    "    if roi == 3: temp = betas_tr_[:, sum(list(roi_neuron.values())[:roi-1]):]\n",
    "    zeros = np.zeros((betas_tr_.shape[0], num_feature))\n",
    "    zeros[:, :temp.shape[-1]] = temp\n",
    "    betas_tr_temp = np.hstack([betas_tr_temp, zeros])\n",
    "\n",
    "betas_te_avg_temp = betas_te_avg_[:, :roi_neuron['V1']]\n",
    "zeros = np.zeros((betas_te_avg_temp.shape[0], num_feature))\n",
    "zeros[:, :betas_te_avg_temp.shape[-1]] = betas_te_avg_temp\n",
    "betas_te_avg_temp = zeros\n",
    "# for roi in [2,3,4]:\n",
    "for roi in [2,3]:\n",
    "    temp = betas_te_avg_[:, sum(list(roi_neuron.values())[:roi-1]):sum(list(roi_neuron.values())[:roi])]\n",
    "    if roi == 3: temp = betas_te_avg_[:, sum(list(roi_neuron.values())[:roi-1]):]\n",
    "    zeros = np.zeros((betas_te_avg_.shape[0], num_feature))\n",
    "    zeros[:, :temp.shape[-1]] = temp\n",
    "    betas_te_avg_temp = np.hstack([betas_te_avg_temp, zeros])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a2e973ce-a298-4bd6-9245-52827caece03",
   "metadata": {},
   "outputs": [],
   "source": [
    "betas_tr_ = betas_tr_temp\n",
    "betas_te_avg_ = betas_te_avg_temp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7785af68-f500-4e10-8b37-0643b35dd7d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_roi = 3\n",
    "\n",
    "num_neurons = betas_tr_.shape[-1]\n",
    "num_pad = (num_neurons // num_roi + 1) * num_roi - num_neurons\n",
    "\n",
    "betas_tr = np.pad(betas_tr_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "betas_te_avg = np.pad(betas_te_avg_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "73e6d383-91be-41ef-bd97-0ff64bb0bc83",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/hdd/yuchen/anaconda3/envs/ldm/lib/python3.8/site-packages/nilearn/connectome/connectivity_matrices.py:509: FutureWarning: The default strategy for standardize is currently 'zscore' which incorrectly uses population std to calculate sample zscores. The new strategy 'zscore_sample' corrects this behavior by using the sample std. In release 0.13, the default strategy will be replaced by the new strategy and the 'zscore' option will be removed. Please use 'zscore_sample' instead.\n",
      "  covariances_std = [\n"
     ]
    }
   ],
   "source": [
    "num_neurons_each_roi = int(betas_te_avg.shape[-1]/num_roi)\n",
    "\n",
    "conn_measure = connectome.ConnectivityMeasure(kind='correlation')\n",
    "\n",
    "function_con = betas_tr.reshape(-1, num_roi, num_neurons_each_roi)\n",
    "\n",
    "connectivity = conn_measure.fit_transform([np.mean(function_con, axis=-1)])\n",
    "for i in range(len(connectivity[0])):\n",
    "    connectivity[0][i, i] = 0  # setting the i-th row to zero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "de5db6dc-10a8-4b04-b2f1-1fd87895f868",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 3, 3)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "connectivity.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "40ee759e-4e0e-4dd2-903a-0ef1e6ffe576",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 6)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adjmatrix = np.where(abs(connectivity[0]) > 0.5)\n",
    "edge_index = np.array(adjmatrix)  # edge index matrix of shape [2, num_edges]\n",
    "\n",
    "edge_index.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "33455771-3966-47d2-bd97-bf35f967253b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(6,)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "edge_attr = []\n",
    "for idx in range(edge_index.shape[1]):\n",
    "    edge_attr.append(connectivity[0][edge_index[0,idx], edge_index[1,idx]])\n",
    "edge_attr = np.array(edge_attr)\n",
    "edge_attr.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "da977361-3ae4-48a8-b19f-f0d554014529",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 0, 1, 1, 2, 2],\n",
       "       [1, 2, 0, 2, 0, 1]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "edge_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ef50f958-34d9-4b3d-92cd-1e3b9cdebe96",
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_index_temp = []\n",
    "edge_attr_temp = []\n",
    "for i in range(edge_attr.shape[0]):\n",
    "    edge = edge_index.T[i]\n",
    "    if edge[0] < edge[1]:\n",
    "        edge_index_temp.append(edge)\n",
    "        edge_attr_temp.append(edge_attr[i])\n",
    "edge_attr = np.array(edge_attr_temp)\n",
    "edge_index = np.array(edge_index_temp).T\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "de0db021-bc55-46cb-9f18-c0515471db6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_trial_tr, num_trials_te = feats_tr.shape[0], feats_te_avg.shape[0]\n",
    "\n",
    "train_dataset = []\n",
    "\n",
    "for i in range(num_trial_tr):\n",
    "    train_dataset.append(Data(x=torch.tensor(betas_tr.reshape(-1, num_roi, num_neurons_each_roi)[i], dtype=torch.float), \n",
    "            edge_index=torch.tensor(edge_index, dtype=torch.long), \n",
    "            edge_attr=torch.tensor(edge_attr, dtype=torch.float),\n",
    "            y=torch.tensor(feats_tr[i], dtype=torch.float)\n",
    "            ))\n",
    "    \n",
    "test_dataset = []\n",
    "for i in range(num_trials_te):\n",
    "    test_dataset.append(Data(x=torch.tensor(betas_te_avg.reshape(-1, num_roi, num_neurons_each_roi)[i], dtype=torch.float), \n",
    "            edge_index=torch.tensor(edge_index, dtype=torch.long), \n",
    "            edge_attr=torch.tensor(edge_attr, dtype=torch.float),\n",
    "            y=torch.tensor(feats_te_avg[i], dtype=torch.float)\n",
    "            ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "15bd82d5-b3c7-4f5f-900f-92aef341a1a7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(8640, 100)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_dataset), len(test_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "20b580f4-8cac-42d4-9d7b-f7814758fddf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Data(x=[3, 1376], edge_index=[2, 3], edge_attr=[3], y=[6400]),\n",
       " Data(x=[3, 1376], edge_index=[2, 3], edge_attr=[3], y=[6400]))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset[0], test_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "3710c1eb-b263-4a19-94b7-e79b44b1c557",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=64, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "9ae53258-f5a4-48f6-a535-6399098bcde3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import GCNConv, TopKPooling, BatchNorm,ChebConv, SAGEConv,GATConv\n",
    "from torch_geometric.nn import global_mean_pool, global_max_pool,global_add_pool\n",
    "\n",
    "class GNN(torch.nn.Module):\n",
    "    def __init__(self, num_neurons_each_roi):\n",
    "        super(GNN, self).__init__()\n",
    "        \n",
    "        self.conv1 = ChebConv(num_neurons_each_roi, 128, K=4)\n",
    "\n",
    "        self.conv = ChebConv(128, 128, K=3)\n",
    "        self.bn = BatchNorm(128)\n",
    "\n",
    "        nnconv_channel = 2\n",
    "        self.nnconv = nn.Sequential(\n",
    "            nn.Conv1d(1, nnconv_channel, kernel_size=1),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv1d(nnconv_channel, nnconv_channel, kernel_size=3, padding = 1),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.downsample = nn.Conv1d(nnconv_channel, 1, kernel_size=1)\n",
    "\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(128, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 6400)\n",
    "        )\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch,edge_attr = data.x, data.edge_index, data.batch, data.edge_attr\n",
    "        batch_size = batch.max() + 1\n",
    "\n",
    "        x = F.relu(self.conv1(x, edge_index, edge_attr))\n",
    "        \n",
    "        for i in range(1):\n",
    "            a = self.conv(x, edge_index, edge_attr)\n",
    "            a = self.bn(a)\n",
    "            x = x + F.leaky_relu(a)\n",
    "        \n",
    "        x = global_mean_pool(x, batch) # print(x.shape) torch.Size([64, 256])\n",
    "\n",
    "        x = x.unsqueeze(1)\n",
    "        x = self.nnconv(x)\n",
    "        for j in range(1):\n",
    "            x = x + self.net(x)\n",
    "        x = self.downsample(x)\n",
    "        x = x.flatten(start_dim=1)\n",
    "        x = self.fc(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "a7044222-7cce-44de-891b-481e01a5d892",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[epoch 0] train loss 0.72386 val loss 0.71847\n",
      "[epoch 1] train loss 0.71281 val loss 0.70301\n",
      "[epoch 2] train loss 0.70046 val loss 0.69713\n",
      "[epoch 3] train loss 0.68225 val loss 0.69702\n",
      "[epoch 4] train loss 0.66077 val loss 0.68736\n",
      "[epoch 5] train loss 0.63916 val loss 0.71032\n",
      "[epoch 6] train loss 0.62045 val loss 0.70825\n",
      "[epoch 7] train loss 0.60451 val loss 0.69806\n",
      "[epoch 8] train loss 0.5905 val loss 0.7187\n",
      "[epoch 9] train loss 0.57821 val loss 0.73498\n",
      "[epoch 10] train loss 0.56881 val loss 0.72134\n",
      "end\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda:0'\n",
    "model = GNN(num_neurons_each_roi=num_neurons_each_roi).to(device)\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0015)\n",
    "\n",
    "num_epochs = 100\n",
    "patience = 5\n",
    "ct = 0\n",
    "best_loss_val = np.inf\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss_train = 0\n",
    "    model.train()\n",
    "    for batch in train_dataloader:\n",
    "        batch = batch.to(device)\n",
    "        batch_num = batch.batch.max() + 1\n",
    "        feats = batch['y'].reshape(batch_num,-1).to(device)\n",
    "        pred = model(batch)\n",
    "        \n",
    "        loss = loss_fn(pred, feats)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss_train += loss.item()\n",
    "    print('[epoch {}] train loss {}'.format(epoch, round(total_loss_train/len(train_dataloader), 5)), end = ' ')\n",
    "\n",
    "    model.eval()\n",
    "    total_loss_val = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in test_dataloader:\n",
    "            batch = batch.to(device)\n",
    "            batch_num = batch.batch.max() + 1\n",
    "            feats = batch['y'].reshape(batch_num,-1).to(device)\n",
    "            pred = model(batch)\n",
    "            loss = loss_fn(pred, feats)\n",
    "            total_loss_val += loss.item()\n",
    "        total_loss_val = total_loss_val/len(test_dataloader)\n",
    "        print('val loss {}'.format(round(total_loss_val, 5)))\n",
    "        \n",
    "        if total_loss_val < best_loss_val:\n",
    "            ct = 0\n",
    "            best_loss_val = total_loss_val\n",
    "            torch.save(model.state_dict(), '/hdd/yuchen/modelweights_things_visual_gcn_roi.pth')\n",
    "        else:\n",
    "            ct += 1\n",
    "            if ct > patience:\n",
    "                print('end')\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "bdfaf03c-30b6-454a-80e3-299a04ad6f4b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6873591840267181"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_loss_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "3f27f073-3a5d-4b98-9395-93b119dfb157",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = 'cuda:0'\n",
    "model = GNN(num_neurons_each_roi=num_neurons_each_roi).to(device)\n",
    "model.load_state_dict(torch.load('/hdd/yuchen/modelweights_things_visual_gcn_roi.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "fbb9afec-12d7-4c29-9c00-106dab3c4e59",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "lst = []\n",
    "lst_gt = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for batch in test_dataloader:\n",
    "        batch = batch.to(device)\n",
    "        batch_num = batch.batch.max() + 1\n",
    "        pred = model(batch)\n",
    "        batch_num = batch.batch.max() + 1\n",
    "        feats = batch['y'].reshape(batch_num,-1).to(device)\n",
    "        \n",
    "        lst.append(pred.cpu().detach().numpy())\n",
    "        lst_gt.append(feats.cpu().detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "7be9ee0a-8e9a-447d-9283-2f3162c84c88",
   "metadata": {},
   "outputs": [],
   "source": [
    "lst = np.vstack(lst)\n",
    "lst_gt = np.vstack(lst_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "15e3273d-fd0a-4d17-bfbb-0834707b3943",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.68060076"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.metrics import mean_squared_error\n",
    "mean_squared_error(lst,lst_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "19838a3b-5d2c-4686-a5fa-a6a57bdbcecb",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save('/hdd/yuchen/things_temp_lst_latent_init.npy', lst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c66a1391-bb64-4f3c-a9e9-37dbe0c732a6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "80b8067b-d8e8-42d0-88de-71dd6cef735f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, betas, feats):\n",
    "        self.betas = torch.tensor(betas, dtype=torch.float32)\n",
    "        self.feats = torch.tensor(feats, dtype=torch.float32)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.betas)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {'betas': self.betas[idx], 'feats': self.feats[idx]}\n",
    "\n",
    "train_dataset = CustomDataset(betas_tr_, feats_tr)\n",
    "test_dataset = CustomDataset(betas_te_avg_, feats_te_avg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "68dd501c-457b-417f-969d-09fe479dcf41",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([3198]), torch.Size([6400]))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset[0]['betas'].shape, train_dataset[0]['feats'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1006689e-817d-4dcf-9b1f-977fbf3ca882",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 256  \n",
    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "for batch in train_dataloader:\n",
    "    break\n",
    "num_neurons = batch['betas'].shape[-1]\n",
    "num_features = batch['feats'].shape[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1ca98026-8aee-4cbe-b9f0-f3f8b1035bb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self, num_neurons, num_features):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        \n",
    "        self.num_neurons = num_neurons\n",
    "        channel = 64\n",
    "        self.conv1d = nn.Conv1d(1, channel, kernel_size=1)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv1d(channel, channel, kernel_size=11, padding = 5),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv1d(channel, channel, kernel_size=11, padding = 5),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.downsample = nn.Conv1d(channel, 1, kernel_size=1)\n",
    "\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(num_neurons, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 6400)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1d(x)\n",
    "        for i in range(2):\n",
    "            x = x + self.net(x)\n",
    "        x = self.downsample(x)\n",
    "        x = x.reshape(-1, self.num_neurons)\n",
    "        x = self.fc(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3db402b6-4067-4701-bd80-a4d69b1639f9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[epoch 0] train loss 0.73409 val loss 0.70004\n",
      "[epoch 1] train loss 0.71778 val loss 0.69203\n",
      "[epoch 2] train loss 0.70781 val loss 0.69262\n",
      "[epoch 3] train loss 0.69221 val loss 0.69515\n",
      "[epoch 4] train loss 0.66643 val loss 0.69906\n",
      "[epoch 5] train loss 0.64036 val loss 0.70593\n",
      "[epoch 6] train loss 0.61582 val loss 0.74341\n",
      "[epoch 7] train loss 0.59286 val loss 0.75947\n",
      "end\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda:0'\n",
    "\n",
    "model = NeuralNetwork(num_neurons = num_neurons, num_features = num_features).to(device)\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)\n",
    "\n",
    "num_epochs = 100\n",
    "patience = 5\n",
    "ct = 0\n",
    "best_loss_val = np.inf\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss_train = 0\n",
    "    model.train()\n",
    "    for batch in train_dataloader:\n",
    "        betas = batch['betas'].to(device).unsqueeze(1)\n",
    "        feats = batch['feats'].to(device)\n",
    "        \n",
    "        pred = model(betas)\n",
    "        loss = loss_fn(pred, feats)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss_train += loss.item()\n",
    "    print('[epoch {}] train loss {}'.format(epoch, round(total_loss_train/len(train_dataloader), 5)), end = ' ')\n",
    "\n",
    "    model.eval()\n",
    "    total_loss_val = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in test_dataloader:\n",
    "            betas = batch['betas'].to(device).unsqueeze(1)\n",
    "            feats = batch['feats'].to(device)\n",
    "            pred = model(betas)\n",
    "            loss = loss_fn(pred, feats)\n",
    "            total_loss_val += loss.item()\n",
    "        total_loss_val = total_loss_val/len(test_dataloader)\n",
    "        print('val loss {}'.format(round(total_loss_val, 5)))\n",
    "        \n",
    "        if total_loss_val < best_loss_val:\n",
    "            ct = 0\n",
    "            best_loss_val = total_loss_val\n",
    "            torch.save(model.state_dict(), '/hdd/yuchen/modelweights_visual_things_textnn.pth')\n",
    "        else:\n",
    "            ct += 1\n",
    "            if ct > patience:\n",
    "                print('end')\n",
    "                break"
   ]
  },
  {
   "cell_type": "raw",
   "id": "ff6d852a-a8ad-4bad-9855-c3be8ea5f0cd",
   "metadata": {},
   "source": [
    "lr0.001: 0.56826; lr0.0005: 0.56561\n",
    "    2 resblock: 0.56314; 1resblock: 0.56571\n",
    "        conv_channnel32: 0.56449; conv_channnel128: 0.56253 (very slow)\n",
    "        remove last conv in each resblock: 0.5635\n",
    "        kernal=5: 0.56329; kernal=11: 0.56288; kernal=13: 0.5633823126554489\n",
    "            num_neurons->512&1024: 0.5636\n",
    "            num_neurons->256: 0.5682\n",
    "            num_neurons->1024: 0.5634\n",
    "            num_neurons->512&512: 0.56112\n",
    "            num_neurons->512&256: bad\n",
    "                conv without padding: kernal=11: 0.56968; kernal=5: 0.5665; kernal=3: 0.56189\n",
    "                gelu in resblock: 0.5616\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3f962a50-af0d-4541-88d0-4ce9015ceae1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6920348405838013"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_loss_val"
   ]
  },
  {
   "cell_type": "raw",
   "id": "dd18e84a-129e-4a46-88fa-5f701f027cd5",
   "metadata": {},
   "source": [
    "Idea inspired by: https://arxiv.org/pdf/2302.12971.pdf\n",
    "\"BrainCLIP: Bridging Brain and Visual-Linguistic Representation Via CLIP for Generic Natural Visual Stimulus Decoding\"\n",
    "\n",
    "encoder of the VAE reconstructed from appendix figure 1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b6917528-52f5-4737-bbd0-59b19e9ad93a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, betas, feats):\n",
    "        self.betas = torch.tensor(betas, dtype=torch.float32)\n",
    "        self.feats = torch.tensor(feats, dtype=torch.float32)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.betas)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {'betas': self.betas[idx], 'feats': self.feats[idx]}\n",
    "\n",
    "train_dataset = CustomDataset(betas_tr_, feats_tr)\n",
    "test_dataset = CustomDataset(betas_te_avg_, feats_te_avg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "863db648-1fab-4513-975c-78b4b4ee688e",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 256  \n",
    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "for batch in train_dataloader:\n",
    "    break\n",
    "num_neurons = batch['betas'].shape[-1]\n",
    "num_features = batch['feats'].shape[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6b0f84d3-bc11-4611-a621-9ed9b5c95fc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "class FBL(nn.Module):\n",
    "    def __init__(self, in_features, out_features):\n",
    "        super(FBL, self).__init__()\n",
    "        self.fc = nn.Linear(in_features, out_features)\n",
    "        self.bn = nn.BatchNorm1d(out_features)\n",
    "        self.act = nn.LeakyReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc(x)\n",
    "        x = self.bn(x)\n",
    "        return self.act(x)\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dims, z_dim):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.ln1 = nn.Linear(input_dim, 512)\n",
    "        self.relu = nn.LeakyReLU()\n",
    "        self.ln2 = nn.Linear(512, 256)\n",
    "        self.fc_mu = nn.Linear(256, z_dim)\n",
    "        self.fc_log_var = nn.Linear(256, z_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.ln1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.ln2(x)\n",
    "        mu = self.fc_mu(x)\n",
    "        log_var = self.fc_log_var(x)\n",
    "        return mu, log_var\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, z_dim, hidden_dims, output_dim):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.fbl1 = FBL(z_dim, hidden_dims)\n",
    "        self.fbl2 = FBL(hidden_dims, hidden_dims)\n",
    "        self.fbl3 = FBL(hidden_dims, hidden_dims)\n",
    "        self.final_fc = nn.Linear(hidden_dims, output_dim)\n",
    "\n",
    "    def forward(self, z):\n",
    "        z = self.fbl1(z)\n",
    "        for i in range(1):\n",
    "            z = self.fbl2(z)\n",
    "        z = self.fbl3(z)\n",
    "        return self.final_fc(z)\n",
    "\n",
    "class VAE(nn.Module):\n",
    "    def __init__(self, fMRI_input_dim, hidden_dims, z_dim, output_dim):\n",
    "        super(VAE, self).__init__()\n",
    "        self.fMRI_encoder = Encoder(fMRI_input_dim, hidden_dims, z_dim)\n",
    "        self.decoder = Decoder(z_dim, hidden_dims, output_dim)\n",
    "\n",
    "    def reparameterize(self, mu, log_var):\n",
    "        std = torch.exp(0.5*log_var)\n",
    "        eps = torch.randn_like(std)\n",
    "        return mu + eps*std\n",
    "\n",
    "    def forward(self, fMRI):\n",
    "        fMRI_mu, fMRI_log_var = self.fMRI_encoder(fMRI)\n",
    "        z_fMRI = self.reparameterize(fMRI_mu, fMRI_log_var)\n",
    "        fMRI_embedding = self.decoder(z_fMRI)\n",
    "\n",
    "        return fMRI_embedding, fMRI_mu, fMRI_log_var, z_fMRI\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f81de1d7-370a-4d48-be07-e842cff6fae6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[epoch 0] train loss 1.92623 val loss 0.78619\n",
      "[epoch 1] train loss 1.4294 val loss 0.75204\n",
      "[epoch 2] train loss 1.39404 val loss 0.73553\n",
      "[epoch 3] train loss 1.3761 val loss 0.73762\n",
      "[epoch 4] train loss 1.35226 val loss 0.73101\n",
      "[epoch 5] train loss 1.33352 val loss 0.7342\n",
      "[epoch 6] train loss 1.32215 val loss 0.72603\n",
      "[epoch 7] train loss 1.30154 val loss 0.74174\n",
      "[epoch 8] train loss 1.28702 val loss 0.75498\n",
      "[epoch 9] train loss 1.26782 val loss 0.7313\n",
      "[epoch 10] train loss 1.25246 val loss 0.74651\n",
      "[epoch 11] train loss 1.24498 val loss 0.74528\n",
      "[epoch 12] train loss 1.23276 val loss 0.77179\n",
      "end\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda:0'\n",
    "\n",
    "fMRI_input_dim = num_neurons  \n",
    "hidden_dims = 4096 \n",
    "z_dim = num_features  \n",
    "output_dim = num_neurons  \n",
    "\n",
    "model = VAE(fMRI_input_dim, hidden_dims, z_dim, output_dim).to(device)\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)\n",
    "\n",
    "num_epochs = 100\n",
    "patience = 5\n",
    "ct = 0\n",
    "best_loss_val = np.inf\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss_train = 0\n",
    "    model.train()\n",
    "    for batch in train_dataloader:\n",
    "        betas = batch['betas'].to(device)\n",
    "        feats = batch['feats'].to(device)\n",
    "        fMRI_embedding, fMRI_mu, fMRI_log_var, z_fMRI = model(betas)\n",
    "        recon_loss_fMRI = loss_fn(fMRI_embedding, betas)\n",
    "        kld_fMRI = loss_fn(feats, z_fMRI)\n",
    "        loss = recon_loss_fMRI + kld_fMRI\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss_train += loss.item()\n",
    "    print('[epoch {}] train loss {}'.format(epoch, round(total_loss_train/len(train_dataloader), 5)), end = ' ')\n",
    "        \n",
    "        \n",
    "    model.eval()\n",
    "    total_loss_val = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in test_dataloader:\n",
    "            betas = batch['betas'].to(device)\n",
    "            feats = batch['feats'].to(device)\n",
    "\n",
    "            fMRI_embedding, fMRI_mu, fMRI_log_var, z_fMRI = model(betas)\n",
    "        \n",
    "            loss = loss_fn(feats, z_fMRI)\n",
    "            total_loss_val += loss.item()\n",
    "        \n",
    "        total_loss_val = total_loss_val/len(test_dataloader)\n",
    "        print('val loss {}'.format(round(total_loss_val, 5)))\n",
    "        \n",
    "        if total_loss_val < best_loss_val:\n",
    "            ct = 0\n",
    "            best_loss_val = total_loss_val\n",
    "            torch.save(model.state_dict(), '/hdd/yuchen/modelweights_visual_things_vae.pth')\n",
    "        else:\n",
    "            ct += 1\n",
    "            if ct > patience:\n",
    "                print('end')\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "2f3352fd-64b4-478c-85ce-972927d6940b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7260345816612244"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_loss_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "564f0737-c3aa-4435-94e2-feb48d675741",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d969aba0-cc4d-420a-a53d-d8bdefe7662e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6971274161290577"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import argparse, os\n",
    "import numpy as np\n",
    "from himalaya.backend import set_backend\n",
    "from himalaya.ridge import RidgeCV\n",
    "from himalaya.scoring import correlation_score\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "alpha = [0.000001,0.00001,0.0001,0.001,0.01, 0.1, 1]\n",
    "\n",
    "ridge = RidgeCV(alphas=alpha)\n",
    "\n",
    "preprocess_pipeline = make_pipeline(\n",
    "        StandardScaler(with_mean=True, with_std=True),\n",
    "    )\n",
    "pipeline = make_pipeline(\n",
    "        preprocess_pipeline,\n",
    "        ridge,\n",
    ")    \n",
    "\n",
    "X = betas_tr\n",
    "Y = feats_tr\n",
    "X_te = betas_te_avg\n",
    "pipeline.fit(X, Y)\n",
    "scores = pipeline.predict(X_te)\n",
    "from sklearn.metrics import mean_squared_error\n",
    "mean_squared_error(scores,feats_te_avg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0e9e610-3c35-40d7-8fa7-4ce1f13f3ee6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "18ef401d-08f5-46c9-bde6-f35596f78e1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/hdd/yuchen/anaconda3/envs/ldm/lib/python3.8/site-packages/nilearn/connectome/connectivity_matrices.py:509: FutureWarning: The default strategy for standardize is currently 'zscore' which incorrectly uses population std to calculate sample zscores. The new strategy 'zscore_sample' corrects this behavior by using the sample std. In release 0.13, the default strategy will be replaced by the new strategy and the 'zscore' option will be removed. Please use 'zscore_sample' instead.\n",
      "  covariances_std = [\n"
     ]
    }
   ],
   "source": [
    "num_roi = 4\n",
    "num_neurons = betas_tr_.shape[-1]\n",
    "num_pad = (num_neurons // num_roi + 1) * num_roi - num_neurons\n",
    "\n",
    "betas_tr = np.pad(betas_tr_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "betas_tr_avg = np.pad(betas_tr_avg_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "betas_te = np.pad(betas_te_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "betas_te_avg = np.pad(betas_te_avg_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "\n",
    "num_neurons_each_roi = int(betas_te_avg.shape[-1]/num_roi)\n",
    "conn_measure = connectome.ConnectivityMeasure(kind='correlation')\n",
    "function_con = betas_tr.reshape(-1, num_roi, num_neurons_each_roi)\n",
    "connectivity = conn_measure.fit_transform([np.mean(function_con, axis=-1)])\n",
    "for i in range(len(connectivity[0])):\n",
    "    connectivity[0][i, i] = 0  # Setting the i-th row to zero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0a9b3a32-1b49-4b9d-93ca-a3bb7fd8bd0a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 4)\n"
     ]
    }
   ],
   "source": [
    "adjmatrix = np.where(abs(connectivity[0]) > 0.8)\n",
    "edge_index = np.array(adjmatrix)  # Edge index matrix of shape [2, num_edges]\n",
    "print(edge_index.shape)\n",
    "\n",
    "edge_attr = []\n",
    "for idx in range(edge_index.shape[1]):\n",
    "    edge_attr.append(connectivity[0][edge_index[0,idx], edge_index[1,idx]])\n",
    "edge_attr = np.array(edge_attr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d65ddb67-a48d-4659-bb15-d1e72657c199",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_trial_tr, num_trials_te = feats_tr.shape[0], feats_te_avg.shape[0]\n",
    "\n",
    "train_dataset = []\n",
    "for i in range(num_trial_tr):\n",
    "    train_dataset.append(Data(x=torch.tensor(betas_tr.reshape(-1, num_roi, num_neurons_each_roi)[i], dtype=torch.float), \n",
    "            edge_index=torch.tensor(edge_index, dtype=torch.long), \n",
    "            edge_attr=torch.tensor(edge_attr, dtype=torch.float),\n",
    "            y=torch.tensor(feats_tr[i], dtype=torch.float)\n",
    "            ))\n",
    "    \n",
    "test_dataset = []\n",
    "for i in range(num_trials_te):\n",
    "    test_dataset.append(Data(x=torch.tensor(betas_te_avg.reshape(-1, num_roi, num_neurons_each_roi)[i], dtype=torch.float), \n",
    "            edge_index=torch.tensor(edge_index, dtype=torch.long), \n",
    "            edge_attr=torch.tensor(edge_attr, dtype=torch.float),\n",
    "            y=torch.tensor(feats_te_avg[i], dtype=torch.float)\n",
    "            ))\n",
    "\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=64, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0dadecd5-48d4-4df6-b59f-8a393b6e8ba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import GCNConv, TopKPooling, BatchNorm,ChebConv, SAGEConv,GATConv\n",
    "from torch_geometric.nn import global_mean_pool, global_max_pool,global_add_pool\n",
    "\n",
    "class GCNWithPooling(torch.nn.Module):\n",
    "    def __init__(self, num_neurons_each_roi):\n",
    "        super(GCNWithPooling, self).__init__()\n",
    "        \n",
    "        self.conv1 = ChebConv(num_neurons_each_roi, 256, K=4)\n",
    "\n",
    "        self.dropout = nn.Dropout(p=0.1)\n",
    "        self.conv = ChebConv(256, 256, K=3)\n",
    "        self.bn = BatchNorm(256)\n",
    "\n",
    "        nnconv_channel = 16\n",
    "        self.nnconv = nn.Sequential(\n",
    "            nn.Conv1d(1, nnconv_channel, kernel_size=1),\n",
    "            nn.ReLU(), \n",
    "        )\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv1d(nnconv_channel, nnconv_channel, kernel_size=3, padding = 1),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.downsample = nn.Conv1d(nnconv_channel, 1, kernel_size=1)\n",
    "\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(256, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 6400)\n",
    "        )\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch,edge_attr = data.x, data.edge_index, data.batch, data.edge_attr\n",
    "        batch_size = batch.max() + 1\n",
    "\n",
    "        x = F.relu(self.conv1(x, edge_index, edge_attr))\n",
    "\n",
    "        for i in range(1):\n",
    "            a = self.conv(x, edge_index, edge_attr)\n",
    "            a = self.bn(a)\n",
    "            x = x + F.relu(a)\n",
    "\n",
    "        x = global_mean_pool(x, batch) # print(x.shape) torch.Size([64, 256])\n",
    "\n",
    "        x = x.unsqueeze(1)\n",
    "        x = self.nnconv(x)\n",
    "        \n",
    "        for j in range(2):\n",
    "            x = x + self.net(x)\n",
    "        x = self.downsample(x)\n",
    "        x = x.flatten(start_dim=1)\n",
    "        x = self.fc(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "2edb8c5d-afaa-40d5-ac55-50c16d2ab01f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[epoch 0] train loss 0.65145 val loss 0.62354\n",
      "[epoch 1] train loss 0.64557 val loss 0.62117\n",
      "[epoch 2] train loss 0.64482 val loss 0.62083\n",
      "[epoch 3] train loss 0.64456 val loss 0.62064\n",
      "[epoch 4] train loss 0.64436 val loss 0.62031\n",
      "[epoch 5] train loss 0.64402 val loss 0.61988\n",
      "[epoch 6] train loss 0.64311 val loss 0.61861\n",
      "[epoch 7] train loss 0.64154 val loss 0.6169\n",
      "[epoch 8] train loss 0.63933 val loss 0.61434\n",
      "[epoch 9] train loss 0.63718 val loss 0.61329\n",
      "[epoch 10] train loss 0.63523 val loss 0.61178\n",
      "[epoch 11] train loss 0.63353 val loss 0.61143\n",
      "[epoch 12] train loss 0.63188 val loss 0.61121\n",
      "[epoch 13] train loss 0.63019 val loss 0.61096\n",
      "[epoch 14] train loss 0.62867 val loss 0.61115\n",
      "[epoch 15] train loss 0.62739 val loss 0.611\n",
      "[epoch 16] train loss 0.62583 val loss 0.6112\n",
      "[epoch 17] train loss 0.62411 val loss 0.61157\n",
      "[epoch 18] train loss 0.6224 val loss 0.61154\n",
      "[epoch 19] train loss 0.62063 val loss 0.61303\n",
      "end\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda:1'\n",
    "model = GCNWithPooling(num_neurons_each_roi=num_neurons_each_roi).to(device)\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.00001)\n",
    "\n",
    "num_epochs = 100\n",
    "patience = 5\n",
    "ct = 0\n",
    "best_loss_val = np.inf\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss_train = 0\n",
    "    model.train()\n",
    "    for batch in train_dataloader:\n",
    "        batch = batch.to(device)\n",
    "        batch_num = batch.batch.max() + 1\n",
    "        feats = batch['y'].reshape(batch_num,-1).to(device)\n",
    "        pred = model(batch)\n",
    "        \n",
    "        loss = loss_fn(pred, feats)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss_train += loss.item()\n",
    "    print('[epoch {}] train loss {}'.format(epoch, round(total_loss_train/len(train_dataloader), 5)), end = ' ')\n",
    "\n",
    "    model.eval()\n",
    "    total_loss_val = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in test_dataloader:\n",
    "            batch = batch.to(device)\n",
    "            batch_num = batch.batch.max() + 1\n",
    "            feats = batch['y'].reshape(batch_num,-1).to(device)\n",
    "            pred = model(batch)\n",
    "            loss = loss_fn(pred, feats)\n",
    "            total_loss_val += loss.item()\n",
    "        total_loss_val = total_loss_val/len(test_dataloader)\n",
    "        print('val loss {}'.format(round(total_loss_val, 5)))\n",
    "        \n",
    "        if total_loss_val < best_loss_val:\n",
    "            ct = 0\n",
    "            best_loss_val = total_loss_val\n",
    "            torch.save(model.state_dict(), '/hdd/yuchen/modelweights_visual_gcn_test.pth')\n",
    "        else:\n",
    "            ct += 1\n",
    "            if ct > patience:\n",
    "                print('end')\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34532e39-6366-4904-bbda-f6f01911acd1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad5eba82-3818-4266-861c-a63600c17451",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "797ca880-1891-448e-b634-0e2457634bfe",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93a571a4-3b48-4639-92b9-a181b17d208c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4563000-839e-47b4-af2f-f217ff37c8ac",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7318bf-595e-4f98-87a6-75dcf44c9e1c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbce9413-34a6-484e-a18c-3873f289ec19",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "25344d33-c3d8-4f1e-8f64-d3fbcdace130",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 6400])"
      ]
     },
     "execution_count": 181,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = model(batch)\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "243a89e2-b94b-4206-9144-4b90a2f765ad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1480, 64])"
      ]
     },
     "execution_count": 151,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(dataset[0]).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "id": "569d8a18-f8a6-4e42-95d1-cbe4b1950bdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "id": "bba313f1-d979-4e96-a0de-2bfbf53df258",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 1480, 64])"
      ]
     },
     "execution_count": 169,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = model(batch)\n",
    "out.reshape(16, -1, 64).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "id": "a538617e-4f55-4fc5-9698-0464de17d274",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 140,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_loader = DataLoader(dataset, batch_size=1)\n",
    "len(train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "id": "934214e5-79cd-4af4-ba57-3df915ea4569",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.loader import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "199989bd-ca74-4650-a787-c2cc286749a3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "512ff401-1d56-47cb-bc58-4925013398bd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "06095492-27a1-4bca-ba7a-5e1ad767b259",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCNConv\n",
    "from torch_geometric.datasets import Planetoid\n",
    "\n",
    "# Load a dataset\n",
    "dataset = Planetoid(root='/tmp/Cora', name='Cora')\n",
    "data = dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "54158bca-c6d6-484d-973e-d02760dee2f0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[2708, 1433], edge_index=[2, 10556], y=[2708], train_mask=[2708], val_mask=[2708], test_mask=[2708])"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "56bacd45-18fe-4cf2-8bd9-84919941d8c8",
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'GlobalStorage' object has no attribute 'num_classes'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[55], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mdata\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnum_classes\u001b[49m\n",
      "File \u001b[0;32m/hdd/yuchen/anaconda3/envs/ldm/lib/python3.8/site-packages/torch_geometric/data/data.py:482\u001b[0m, in \u001b[0;36mData.__getattr__\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m    476\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m_store\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__dict__\u001b[39m:\n\u001b[1;32m    477\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\n\u001b[1;32m    478\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThe \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdata\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m object was created by an older version of PyG. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    479\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mIf this error occurred while loading an already existing \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    480\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdataset, remove the \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mprocessed/\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m directory in the dataset\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124ms \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    481\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mroot folder and try again.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 482\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mgetattr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_store\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/hdd/yuchen/anaconda3/envs/ldm/lib/python3.8/site-packages/torch_geometric/data/storage.py:87\u001b[0m, in \u001b[0;36mBaseStorage.__getattr__\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m     85\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m[key]\n\u001b[1;32m     86\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m:\n\u001b[0;32m---> 87\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m(\n\u001b[1;32m     88\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m object has no attribute \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mkey\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     89\u001b[0m     ) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'GlobalStorage' object has no attribute 'num_classes'"
     ]
    }
   ],
   "source": [
    "data.num_classes"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}