Erdos_Cuts.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First take care of all imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import Linear\n",
    "from torch_geometric.datasets import TUDataset\n",
    "import time\n",
    "from torch import tensor\n",
    "from torch.optim import Adam\n",
    "from torch.optim import SGD\n",
    "from torch_geometric.data import DataLoader, DenseDataLoader as DenseLoader\n",
    "from math import ceil\n",
    "from torch.nn import Linear\n",
    "from torch.distributions import categorical\n",
    "from torch.distributions import Bernoulli\n",
    "import torch.nn\n",
    "%matplotlib inline\n",
    "from matplotlib import pyplot as plt\n",
    "import pygraphviz as pgv\n",
    "from torch_geometric.utils import convert as cnv\n",
    "from torch_geometric.utils import sparse as sp\n",
    "from torch_geometric.data import Data\n",
    "import pygraphviz as pgv\n",
    "from networkx.drawing.nx_agraph import graphviz_layout\n",
    "import networkx as nx\n",
    "from torch.utils.data.sampler import RandomSampler\n",
    "from torch.nn.functional import gumbel_softmax\n",
    "from torch.distributions import relaxed_categorical\n",
    "from torch_geometric.nn.inits import uniform\n",
    "from torch_geometric.nn.inits import glorot, zeros\n",
    "from torch.nn import Parameter\n",
    "from torch.nn import Sequential as Seq, Linear, ReLU\n",
    "from torch_geometric.nn import MessagePassing\n",
    "from torch_geometric.utils import degree\n",
    "from torch_geometric.nn import GINConv, GATConv\n",
    "from torch.nn import Parameter\n",
    "from torch.nn import Sequential as Seq, Linear, ReLU, LeakyReLU\n",
    "from torch_geometric.nn import MessagePassing\n",
    "from torch.nn import Linear, Sequential, ReLU, BatchNorm1d as BN\n",
    "from torch_geometric.data import Batch \n",
    "from torch_scatter import scatter_min, scatter_max, scatter_add, scatter_mean\n",
    "from torch import autograd\n",
    "from torch_geometric.utils import softmax, add_self_loops, remove_self_loops, segregate_self_loops, remove_isolated_nodes, contains_isolated_nodes, add_remaining_self_loops\n",
    "from models import cut_MPNN\n",
    "from modules_and_utils import derandomize_cut, GATAConv,get_diracs, total_var\n",
    "import scipy\n",
    "import scipy.io\n",
    "import GPUtil\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = [\"facebook\", \"sf\",\"twitter\"]\n",
    "curr_dataset= datasets[1] \n",
    "#set random seed\n",
    "rseed = 201\n",
    "\n",
    "if curr_dataset==\"facebook\":\n",
    "    datasetname = \"facebook_graphs\"\n",
    "    dataset = []\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    #set up facebook data \n",
    "    upper_limit = 15000\n",
    "    lower_limit = 0\n",
    "    path_to_facebook_dataset= \"%PATH\"\n",
    "    for file in os.listdir(path_to_facebook_dataset):\n",
    "            if file.endswith(\".mat\"):\n",
    "                adj_matrix = scipy.io.loadmat(path_to_facebook_dataset+str(file))\n",
    "                edge_index = from_scipy_sparse_matrix(adj_matrix['A'])[0]\n",
    "                x = torch.ones(adj_matrix['local_info'].shape[0])\n",
    "                if (adj_matrix['local_info'].shape[0] < lower_limit) or (adj_matrix['local_info'].shape[0] > upper_limit):\n",
    "                    continue\n",
    "                data_temp = Batch(x = x, edge_index = edge_index.long(), batch = torch.zeros_like(x).long())\n",
    "                data_proper = get_diracs(data_temp.to('cuda'), 1, sparse = True)\n",
    "                r,c = data_proper.edge_index\n",
    "                data = Batch(x = data_temp.x, edge_index = data_temp.edge_index)\n",
    "                degrees = degree(r, adj_matrix['local_info'].shape[0])\n",
    "                print(\"Graph specs: \")\n",
    "                print(\"number of nodes: \", adj_matrix['A'].shape[0])\n",
    "                print(\"average degree: \", degrees.mean(0))\n",
    "                print(\"total volume: \", data_proper.total_vol)\n",
    "                print(\"-------------\")\n",
    "                dataset += [data]\n",
    "\n",
    "elif curr_dataset==\"sf\":\n",
    "    datasetname = \"SF-295\"\n",
    "    dataset = TUDataset(root='/tmp/'+datasetname, name=datasetname)\n",
    "\n",
    "elif curr_dataset==\"twitter\":\n",
    "    path_to_twitter_dataset = \"%PATH\"\n",
    "    stored_dataset = open(path_to_twitter_dataset, 'rb')        \n",
    "    dataset = pickle.load(stored_dataset)   \n",
    "    dataset = [Data.from_dict(data) for data in dataset] \n",
    "\n",
    "\n",
    "dataset_scale = 0.1\n",
    "total_samples = int(np.floor(len(dataset)*dataset_scale))\n",
    "dataset = dataset[:total_samples]\n",
    "\n",
    "num_trainpoints = int(np.floor(0.6*len(dataset)))\n",
    "num_valpoints = int(np.floor(num_trainpoints/3))\n",
    "num_testpoints = len(dataset) - (num_trainpoints + num_valpoints)\n",
    "\n",
    "\n",
    "traindata= dataset[0:num_trainpoints]\n",
    "valdata = dataset[num_trainpoints:num_trainpoints + num_valpoints]\n",
    "\n",
    "testdata = dataset[num_trainpoints + num_valpoints:]\n",
    "\n",
    "batch_size = 32\n",
    "train_loader = DataLoader(traindata, batch_size, shuffle=True)\n",
    "test_loader = DataLoader(testdata, batch_size, shuffle=False)\n",
    "val_loader =  DataLoader(valdata, batch_size, shuffle=False)\n",
    "\n",
    "\n",
    "#set up random seed \n",
    "torch.manual_seed(rseed)\n",
    "np.random.seed(2)   \n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TRAIN MODEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(model, data_loader, recfield):\n",
    "    net.eval()\n",
    "    avg_loss = 0\n",
    "    for data in data_loader:\n",
    "        optimizer.zero_grad()\n",
    "        data = data.to(device)\n",
    "        data = get_diracs(data, 1, sparse = True, effective_volume_range=0.15, receptive_field = recfield)\n",
    "        data = data.to(device)\n",
    "        retdict = net(data)\n",
    "        avg_loss += retdict['loss'][0].item()/len(data_loader)\n",
    "\n",
    "    return avg_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Current epoch:  0\n",
      "torch.Size([1154])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/scratch/lts2/karalias/repoz/erdos_neural/modules_and_utils.py:278: UserWarning: This overload of nonzero is deprecated:\n",
      "\tnonzero()\n",
      "Consider using one of the following signatures instead:\n",
      "\tnonzero(*, bool as_tuple) (Triggered internally at  /opt/conda/conda-bld/pytorch_1607370172916/work/torch/csrc/utils/python_arg_parser.cpp:882.)\n",
      "  locationmatrix = diracmatrix.nonzero()\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([971])\n",
      "torch.Size([1041])\n",
      "torch.Size([1100])\n",
      "torch.Size([896])\n",
      "torch.Size([1106])\n",
      "torch.Size([975])\n",
      "torch.Size([1137])\n",
      "torch.Size([1114])\n",
      "torch.Size([1209])\n",
      "torch.Size([1017])\n",
      "torch.Size([1240])\n",
      "torch.Size([1002])\n",
      "torch.Size([1129])\n",
      "torch.Size([1009])\n",
      "torch.Size([1129])\n",
      "torch.Size([1099])\n",
      "torch.Size([1030])\n",
      "torch.Size([1064])\n",
      "torch.Size([1123])\n",
      "torch.Size([1037])\n",
      "torch.Size([1029])\n",
      "torch.Size([968])\n",
      "torch.Size([958])\n",
      "torch.Size([1032])\n",
      "torch.Size([1079])\n",
      "torch.Size([1173])\n",
      "torch.Size([1121])\n",
      "torch.Size([1064])\n",
      "torch.Size([1003])\n",
      "torch.Size([956])\n",
      "torch.Size([1098])\n",
      "torch.Size([1107])\n",
      "torch.Size([1049])\n",
      "torch.Size([1148])\n",
      "torch.Size([989])\n",
      "torch.Size([1187])\n",
      "torch.Size([1084])\n",
      "torch.Size([1002])\n",
      "torch.Size([1001])\n",
      "torch.Size([1196])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "ERROR:root:Internal Python error in the inspect module.\n",
      "Below is the traceback from this internal error.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Traceback (most recent call last):\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/interactiveshell.py\", line 3417, in run_code\n",
      "    exec(code_obj, self.user_global_ns, self.user_ns)\n",
      "  File \"<ipython-input-4-7d14a9b589d1>\", line 37, in <module>\n",
      "    retdict = net(data_prime)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/module.py\", line 727, in _call_impl\n",
      "    result = self.forward(*input, **kwargs)\n",
      "  File \"/mnt/scratch/lts2/karalias/repoz/erdos_neural/models.py\", line 118, in forward\n",
      "    x = x + conv(x, edge_index)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/module.py\", line 727, in _call_impl\n",
      "    result = self.forward(*input, **kwargs)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch_geometric/nn/conv/gin_conv.py\", line 70, in forward\n",
      "    return self.nn(out)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/module.py\", line 727, in _call_impl\n",
      "    result = self.forward(*input, **kwargs)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/container.py\", line 117, in forward\n",
      "    input = module(input)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/module.py\", line 727, in _call_impl\n",
      "    result = self.forward(*input, **kwargs)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/batchnorm.py\", line 131, in forward\n",
      "    return F.batch_norm(\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/functional.py\", line 2056, in batch_norm\n",
      "    return torch.batch_norm(\n",
      "KeyboardInterrupt\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/interactiveshell.py\", line 2044, in showtraceback\n",
      "    stb = value._render_traceback_()\n",
      "AttributeError: 'KeyboardInterrupt' object has no attribute '_render_traceback_'\n",
      "\n",
      "During handling of the above exception, another exception occurred:\n",
      "\n",
      "Traceback (most recent call last):\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/ultratb.py\", line 1169, in get_records\n",
      "    return _fixed_getinnerframes(etb, number_of_lines_of_context, tb_offset)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/ultratb.py\", line 316, in wrapped\n",
      "    return f(*args, **kwargs)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/ultratb.py\", line 350, in _fixed_getinnerframes\n",
      "    records = fix_frame_records_filenames(inspect.getinnerframes(etb, context))\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/inspect.py\", line 1503, in getinnerframes\n",
      "    frameinfo = (tb.tb_frame,) + getframeinfo(tb, context)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/inspect.py\", line 1461, in getframeinfo\n",
      "    filename = getsourcefile(frame) or getfile(frame)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/inspect.py\", line 708, in getsourcefile\n",
      "    if getattr(getmodule(object, filename), '__loader__', None) is not None:\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/inspect.py\", line 751, in getmodule\n",
      "    f = getabsfile(module)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/inspect.py\", line 721, in getabsfile\n",
      "    return os.path.normcase(os.path.abspath(_filename))\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/posixpath.py\", line 381, in abspath\n",
      "    return normpath(path)\n",
      "  File \"/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/posixpath.py\", line 360, in normpath\n",
      "    if (comp != dotdot or (not initial_slashes and not new_comps) or\n",
      "KeyboardInterrupt\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "object of type 'NoneType' has no len()",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
      "\u001b[0;32m<ipython-input-4-7d14a9b589d1>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     36\u001b[0m         \u001b[0mdata_prime\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdata_prime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 37\u001b[0;31m         \u001b[0mretdict\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata_prime\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     38\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mval\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mretdict\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitems\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    726\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 727\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    728\u001b[0m         for hook in itertools.chain(\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/repoz/erdos_neural/models.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, data, tvol)\u001b[0m\n\u001b[1;32m    117\u001b[0m             \u001b[0;32mif\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m>\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 118\u001b[0;31m                 \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mconv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0medge_index\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    119\u001b[0m                 \u001b[0mmask\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_mask\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmask\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0medge_index\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    726\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 727\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    728\u001b[0m         for hook in itertools.chain(\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch_geometric/nn/conv/gin_conv.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x, edge_index, size)\u001b[0m\n\u001b[1;32m     69\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 70\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     71\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    726\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 727\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    728\u001b[0m         for hook in itertools.chain(\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    116\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 117\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    118\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    726\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 727\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    728\u001b[0m         for hook in itertools.chain(\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/modules/batchnorm.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    130\u001b[0m         \"\"\"\n\u001b[0;32m--> 131\u001b[0;31m         return F.batch_norm(\n\u001b[0m\u001b[1;32m    132\u001b[0m             \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/torch/nn/functional.py\u001b[0m in \u001b[0;36mbatch_norm\u001b[0;34m(input, running_mean, running_var, weight, bias, training, momentum, eps)\u001b[0m\n\u001b[1;32m   2055\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2056\u001b[0;31m     return torch.batch_norm(\n\u001b[0m\u001b[1;32m   2057\u001b[0m         \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrunning_mean\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrunning_var\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: ",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mshowtraceback\u001b[0;34m(self, exc_tuple, filename, tb_offset, exception_only, running_compiled_code)\u001b[0m\n\u001b[1;32m   2043\u001b[0m                         \u001b[0;31m# in the engines. This should return a list of strings.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2044\u001b[0;31m                         \u001b[0mstb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_render_traceback_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2045\u001b[0m                     \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'KeyboardInterrupt' object has no attribute '_render_traceback_'",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mshowtraceback\u001b[0;34m(self, exc_tuple, filename, tb_offset, exception_only, running_compiled_code)\u001b[0m\n\u001b[1;32m   2044\u001b[0m                         \u001b[0mstb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_render_traceback_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2045\u001b[0m                     \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2046\u001b[0;31m                         stb = self.InteractiveTB.structured_traceback(etype,\n\u001b[0m\u001b[1;32m   2047\u001b[0m                                             value, tb, tb_offset=tb_offset)\n\u001b[1;32m   2048\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/ultratb.py\u001b[0m in \u001b[0;36mstructured_traceback\u001b[0;34m(self, etype, value, tb, tb_offset, number_of_lines_of_context)\u001b[0m\n\u001b[1;32m   1433\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1434\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtb\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1435\u001b[0;31m         return FormattedTB.structured_traceback(\n\u001b[0m\u001b[1;32m   1436\u001b[0m             self, etype, value, tb, tb_offset, number_of_lines_of_context)\n\u001b[1;32m   1437\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/ultratb.py\u001b[0m in \u001b[0;36mstructured_traceback\u001b[0;34m(self, etype, value, tb, tb_offset, number_of_lines_of_context)\u001b[0m\n\u001b[1;32m   1333\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mmode\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mverbose_modes\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1334\u001b[0m             \u001b[0;31m# Verbose modes need a full traceback\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1335\u001b[0;31m             return VerboseTB.structured_traceback(\n\u001b[0m\u001b[1;32m   1336\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0metype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtb_offset\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnumber_of_lines_of_context\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1337\u001b[0m             )\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/ultratb.py\u001b[0m in \u001b[0;36mstructured_traceback\u001b[0;34m(self, etype, evalue, etb, tb_offset, number_of_lines_of_context)\u001b[0m\n\u001b[1;32m   1190\u001b[0m         \u001b[0;34m\"\"\"Return a nice text document describing the traceback.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1191\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1192\u001b[0;31m         formatted_exception = self.format_exception_as_a_whole(etype, evalue, etb, number_of_lines_of_context,\n\u001b[0m\u001b[1;32m   1193\u001b[0m                                                                tb_offset)\n\u001b[1;32m   1194\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/ultratb.py\u001b[0m in \u001b[0;36mformat_exception_as_a_whole\u001b[0;34m(self, etype, evalue, etb, number_of_lines_of_context, tb_offset)\u001b[0m\n\u001b[1;32m   1148\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1149\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1150\u001b[0;31m         \u001b[0mlast_unique\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrecursion_repeat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfind_recursion\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0morig_etype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mevalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrecords\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1151\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1152\u001b[0m         \u001b[0mframes\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat_records\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrecords\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlast_unique\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrecursion_repeat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/mnt/scratch/lts2/karalias/envs/neuralll/lib/python3.8/site-packages/IPython/core/ultratb.py\u001b[0m in \u001b[0;36mfind_recursion\u001b[0;34m(etype, value, records)\u001b[0m\n\u001b[1;32m    449\u001b[0m     \u001b[0;31m# first frame (from in to out) that looks different.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    450\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mis_recursion_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0metype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrecords\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 451\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrecords\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    452\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    453\u001b[0m     \u001b[0;31m# Select filename, lineno, func_name to track frames with\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: object of type 'NoneType' has no len()"
     ]
    }
   ],
   "source": [
    "epochs=150\n",
    "numlayers=6\n",
    "elasticity = 0.25\n",
    "receptive_field= numlayers + 1\n",
    "val_losses = []\n",
    "\n",
    "#for sf/twitter\n",
    "#net =  cut_MPNN(dataset,numlayers, 64, 64,1, elasticity = elasticity)\n",
    "\n",
    "#for faceboook\n",
    "net =  cut_MPNN(dataset,6, 256, 24,1, elasticity = 0.25)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "lr_decay_step_size = 5\n",
    "lr_decay_factor = 0.95\n",
    "\n",
    "net.to(device).reset_parameters()\n",
    "optimizer = Adam(net.parameters(), lr=0.0001, weight_decay=0.00)\n",
    "net.train()\n",
    "retdict = {}\n",
    "for epoch in range(epochs):\n",
    "    print(\"Current epoch: \", epoch)\n",
    "    totalretdict = {}\n",
    "    count=0\n",
    "    #learning rate schedule\n",
    "    if epoch % lr_decay_step_size == 0:\n",
    "        for param_group in optimizer.param_groups:\n",
    "                    param_group['lr'] = lr_decay_factor * param_group['lr']\n",
    "   \n",
    "    net.train()\n",
    "    for data in train_loader:\n",
    "        count += 1 \n",
    "        optimizer.zero_grad(), \n",
    "        data = data.to(device)\n",
    "        data_prime = get_diracs(data, 1, sparse = True, effective_volume_range=0.15, receptive_field = receptive_field)\n",
    "        data = data.to('cpu')\n",
    "        data_prime = data_prime.to(device)  \n",
    "        retdict = net(data_prime)\n",
    "        for key,val in retdict.items():\n",
    "            if \"sequence\" in val[1]:\n",
    "                if key in totalretdict:\n",
    "                    totalretdict[key][0] += val[0].item()\n",
    "                else:\n",
    "                    totalretdict[key] = [val[0].item(),val[1]]\n",
    "        \n",
    "        if epoch > 0:\n",
    "                retdict[\"loss\"][0].backward()\n",
    "                torch.nn.utils.clip_grad_norm_(net.parameters(),1)\n",
    "                optimizer.step()                   \n",
    "\n",
    "        del data_prime"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EVALUATE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#initialize\n",
    "net.eval()\n",
    "rcuts =  {}\n",
    "rvols =  {}\n",
    "rconds = {}\n",
    "mcuts =  []\n",
    "mvols =  []\n",
    "mconds = []\n",
    "best_cuts =  []\n",
    "best_vols =  []\n",
    "best_conds = []\n",
    "count = 0\n",
    "\n",
    "#select number of diracs\n",
    "num_diracs = 10\n",
    "\n",
    "cuts = torch.zeros((num_testpoints, num_diracs))\n",
    "conds =  torch.zeros((num_testpoints, num_diracs))\n",
    "vols =  torch.zeros((num_testpoints, num_diracs))\n",
    "randtargets = torch.zeros((num_testpoints, num_diracs))\n",
    "best_sets = {}\n",
    "totalvols = []\n",
    "\n",
    "\n",
    "\n",
    "t_0 = time.time()\n",
    "with torch.no_grad():\n",
    "    for data2 in test_loader:\n",
    "        batch = data2.batch\n",
    "        \n",
    "        count += 1\n",
    "           \n",
    "        print(\"Batch count: \", count)\n",
    "        dirac_count = 0 \n",
    "        for dirac in range(num_diracs):\n",
    "            data2 = data2.to(device)            \n",
    "            data_new = get_diracs(data2, 1, sparse=True, effective_volume_range=0.2)            \n",
    "            \n",
    "            feasible_vols = (data_new.recfield_vol/data_new.total_vol)*0.85\n",
    "            target_vol = torch.rand_like(feasible_vols, device=device)*feasible_vols + 0.1\n",
    "            data_new = data_new.to(device) \n",
    "            retdict2 = net(data_new, target_vol)\n",
    "            netprobs = retdict2['output'][0]\n",
    "            batch_new = data_new.batch\n",
    "            num_graphs = batch_new.max().item() + 1\n",
    "            e_i = data_new.edge_index \n",
    "            r,c = e_i\n",
    "            deg = degree(r)\n",
    "            bestcond = torch.ones(num_graphs)\n",
    "            bestcut = 1000*torch.ones(num_graphs)\n",
    "            bestvol = torch.zeros(num_graphs)\n",
    "            outp =  derandomize_cut(data_new.to('cuda'), netprobs.cuda(), target_vol.cuda()*data_new.total_vol.cuda(), elasticity=0.25, draw=False)\n",
    "            tv_hard = total_var(outp, data_new.edge_index.cuda(), data_new.batch.cuda())\n",
    "            vol_hard = scatter_add(deg*outp, batch_new, 0, dim_size = batch_new.max().item()+1)\n",
    "            mycond = tv_hard/vol_hard\n",
    "            conds[(count-1)*batch_size:count*batch_size, dirac] = mycond\n",
    "            cuts[(count-1)*batch_size:count*batch_size, dirac] = tv_hard\n",
    "            vols[(count-1)*batch_size:count*batch_size, dirac] = vol_hard\n",
    "            randtargets[(count-1)*batch_size:count*batch_size, dirac] = target_vol.cpu()*data_new.total_vol.cpu()          \n",
    "            dirac_count += 1\n",
    "t_final = time.time() - t_0\n",
    "print(f\"average time per graph: {t_final/len(testdata)}\")   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Print out mean conductance +/- std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meanconds = conds.mean(0)\n",
    "print(f\"meanconds: {meanconds.mean()} +/- {meanconds.std()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:karalias-neuralll] *",
   "language": "python",
   "name": "conda-env-karalias-neuralll-py"
  },
  "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": 2
}