main_SBMs_node_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main Driver Notebook for Training Graph NNs on SBM Datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MODELS\n",
    "- GatedGCN \n",
    "- GCN \n",
    "- GAT \n",
    "- GraphSage \n",
    "- MLP\n",
    "- GIN\n",
    "- MoNet\n",
    "- RingGNN  \n",
    "- 3WLGNN  \n",
    "\n",
    "### DATASET\n",
    "- SBM_CLUSTER\n",
    "- SBM_PATTERN\n",
    "\n",
    "### TASK\n",
    "- Node Classification "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    IMPORTING LIBS\n",
    "\"\"\"\n",
    "import dgl\n",
    "\n",
    "import numpy as np\n",
    "import os\n",
    "import socket\n",
    "import time\n",
    "import random\n",
    "import glob\n",
    "import argparse, json\n",
    "import pickle\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from tensorboardX import SummaryWriter\n",
    "from tqdm import tqdm\n",
    "\n",
    "class DotDict(dict):\n",
    "    def __init__(self, **kwds):\n",
    "        self.update(kwds)\n",
    "        self.__dict__ = self"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "# \"\"\"\n",
    "#     AUTORELOAD IPYTHON EXTENSION FOR RELOADING IMPORTED MODULES\n",
    "# \"\"\"\n",
    "\n",
    "def in_ipynb():\n",
    "    try:\n",
    "        cfg = get_ipython().config \n",
    "        return True\n",
    "    except NameError:\n",
    "        return False\n",
    "    \n",
    "notebook_mode = in_ipynb()\n",
    "print(notebook_mode)\n",
    "\n",
    "if notebook_mode == True:\n",
    "    %load_ext autoreload\n",
    "    %autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    IMPORTING CUSTOM MODULES/METHODS\n",
    "\"\"\"\n",
    "\n",
    "from nets.SBMs_node_classification.load_net import gnn_model # import GNNs\n",
    "from data.data import LoadData # import dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    GPU Setup\n",
    "\"\"\"\n",
    "def gpu_setup(use_gpu, gpu_id):\n",
    "    os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "    os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id)  \n",
    "\n",
    "    if torch.cuda.is_available() and use_gpu:\n",
    "        print('cuda available with GPU:',torch.cuda.get_device_name(0))\n",
    "        device = torch.device(\"cuda\")\n",
    "    else:\n",
    "        print('cuda not available')\n",
    "        device = torch.device(\"cpu\")\n",
    "    return device\n",
    "\n",
    "\n",
    "# select GPU or CPU\n",
    "#use_gpu = True; gpu_id = 0; device = None # default GPU\n",
    "use_gpu = False; gpu_id = -1; device = None # CPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[I] Loading dataset SBM_PATTERN...\n",
      "train, test, val sizes : 10000 2000 2000\n",
      "[I] Finished loading.\n",
      "[I] Data load time: 18.0153s\n"
     ]
    }
   ],
   "source": [
    "# \"\"\"\n",
    "#     USER CONTROLS\n",
    "# \"\"\"\n",
    "if notebook_mode == True:\n",
    "    \n",
    "    #MODEL_NAME = '3WLGNN'\n",
    "    #MODEL_NAME = 'RingGNN'\n",
    "    MODEL_NAME = 'GatedGCN'\n",
    "    #MODEL_NAME = 'GCN'\n",
    "    #MODEL_NAME = 'GAT'\n",
    "    #MODEL_NAME = 'GraphSage'\n",
    "    #MODEL_NAME = 'MLP'\n",
    "    #MODEL_NAME = 'GIN'\n",
    "    #MODEL_NAME = 'MoNet'\n",
    "    \n",
    "\n",
    "    #DATASET_NAME = 'SBM_CLUSTER'\n",
    "    DATASET_NAME = 'SBM_PATTERN'\n",
    "\n",
    "    out_dir = 'out/SBMs_node_classification/'\n",
    "    root_log_dir = out_dir + 'logs/' + MODEL_NAME + \"_\" + DATASET_NAME + \"_\" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y')\n",
    "    root_ckpt_dir = out_dir + 'checkpoints/' + MODEL_NAME + \"_\" + DATASET_NAME + \"_\" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y')\n",
    "\n",
    "    dataset = LoadData(DATASET_NAME)\n",
    "    trainset, valset, testset = dataset.train, dataset.val, dataset.test\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#MODEL_NAME = 'RingGNN'\n",
    "MODEL_NAME = 'GatedGCN'\n",
    "#MODEL_NAME = 'GCN'\n",
    "#MODEL_NAME = 'GAT'\n",
    "#MODEL_NAME = 'GraphSage'\n",
    "#MODEL_NAME = 'MLP'\n",
    "#MODEL_NAME = 'DiffPool'\n",
    "#MODEL_NAME = 'GIN'\n",
    "#MODEL_NAME = 'MoNet'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"\"\"\n",
    "#     PARAMETERS\n",
    "# \"\"\"\n",
    "if notebook_mode == True:\n",
    "\n",
    "    n_heads = -1\n",
    "    edge_feat = False\n",
    "    pseudo_dim_MoNet = -1\n",
    "    kernel = -1\n",
    "    gnn_per_block = -1\n",
    "    embedding_dim = -1\n",
    "    pool_ratio = -1\n",
    "    n_mlp_GIN = -1\n",
    "    gated = False\n",
    "    self_loop = False\n",
    "    #self_loop = True\n",
    "    max_time = 12\n",
    "    pos_enc = True\n",
    "    #pos_enc = False\n",
    "    pos_enc_dim = 10\n",
    "    \n",
    "\n",
    "    if MODEL_NAME == 'GatedGCN':\n",
    "        seed=41; epochs=1000; batch_size=5; init_lr=5e-5; lr_reduce_factor=0.5; lr_schedule_patience=25; min_lr = 1e-6; weight_decay=0\n",
    "        L=4; hidden_dim=70; out_dim=hidden_dim; dropout=0.0; readout='mean'\n",
    "\n",
    "    if MODEL_NAME == 'GCN':\n",
    "        seed=41; epochs=1000; batch_size=5; init_lr=5e-5; lr_reduce_factor=0.5; lr_schedule_patience=25; min_lr = 1e-6; weight_decay=0\n",
    "        L=4; hidden_dim=146; out_dim=hidden_dim; dropout=0.0; readout='mean'\n",
    "\n",
    "    if MODEL_NAME == 'GAT':\n",
    "        seed=41; epochs=1000; batch_size=50; init_lr=5e-5; lr_reduce_factor=0.5; lr_schedule_patience=25; min_lr = 1e-6; weight_decay=0\n",
    "        L=4; n_heads=8; hidden_dim=19; out_dim=n_heads*hidden_dim; dropout=0.0; readout='mean'\n",
    "        print('True hidden dim:',out_dim)\n",
    "\n",
    "    if MODEL_NAME == 'GraphSage':\n",
    "        seed=41; epochs=1000; batch_size=50; init_lr=5e-5; lr_reduce_factor=0.5; lr_schedule_patience=25; min_lr = 1e-6; weight_decay=0\n",
    "        L=4; hidden_dim=108; out_dim=hidden_dim; dropout=0.0; readout='mean'\n",
    "\n",
    "    if MODEL_NAME == 'MLP':\n",
    "        seed=41; epochs=1000; batch_size=50; init_lr=5e-4; lr_reduce_factor=0.5; lr_schedule_patience=25; min_lr = 1e-6; weight_decay=0\n",
    "        gated=False; # MEAN\n",
    "        L=4; hidden_dim=150; out_dim=hidden_dim; dropout=0.0; readout='mean'\n",
    "        gated=True; # GATED\n",
    "        L=4; hidden_dim=135; out_dim=hidden_dim; dropout=0.0; readout='mean'\n",
    "        \n",
    "    if MODEL_NAME == 'DiffPool':\n",
    "        seed=41; epochs=1000; batch_size=50; init_lr=5e-4; lr_reduce_factor=0.5; lr_schedule_patience=25; min_lr = 1e-6; weight_decay=0\n",
    "        L=4; hidden_dim=32; out_dim=hidden_dim; dropout=0.0; readout='mean'\n",
    "        n_heads=8; gnn_per_block=3; embedding_dim=32; batch_size=128; pool_ratio=0.15\n",
    "\n",
    "    if MODEL_NAME == 'GIN':\n",
    "        seed=41; epochs=1000; batch_size=50; init_lr=5e-4; lr_reduce_factor=0.5; lr_schedule_patience=25; min_lr = 1e-6; weight_decay=0\n",
    "        L=4; hidden_dim=110; out_dim=hidden_dim; dropout=0.0; readout='mean'\n",
    "        n_mlp_GIN = 2; learn_eps_GIN=True; neighbor_aggr_GIN='sum'\n",
    "\n",
    "    if MODEL_NAME == 'MoNet':\n",
    "        seed=41; epochs=1000; batch_size=50; init_lr=5e-4; lr_reduce_factor=0.5; lr_schedule_patience=25; min_lr = 1e-6; weight_decay=0\n",
    "        L=4; hidden_dim=90; out_dim=hidden_dim; dropout=0.0; readout='mean'\n",
    "        pseudo_dim_MoNet=2; kernel=3;\n",
    "        \n",
    "    if MODEL_NAME == 'RingGNN':\n",
    "        seed=41; epochs=1000; batch_size=1; init_lr=5e-5; lr_reduce_factor=0.5; lr_schedule_patience=25; min_lr = 1e-6; weight_decay=0\n",
    "        #L=4; hidden_dim=145; out_dim=hidden_dim; dropout=0.0; readout='mean'\n",
    "        L=4; hidden_dim=25; out_dim=hidden_dim; dropout=0.0\n",
    "        \n",
    "    if MODEL_NAME == '3WLGNN':\n",
    "        seed=41; epochs=1000; batch_size=1; init_lr=5e-5; lr_reduce_factor=0.5; lr_schedule_patience=25; min_lr = 1e-6; weight_decay=0\n",
    "        L=3; hidden_dim=82; out_dim=hidden_dim; dropout=0.0\n",
    "      \n",
    "\n",
    "        \n",
    "    # generic new_params\n",
    "    net_params = {}\n",
    "    net_params['device'] = device\n",
    "    net_params['in_dim'] = torch.unique(trainset[0][0].ndata['feat'],dim=0).size(0) # node_dim (feat is an integer)\n",
    "    net_params['hidden_dim'] = hidden_dim\n",
    "    net_params['out_dim'] = out_dim\n",
    "    num_classes = torch.unique(trainset[0][1],dim=0).size(0)\n",
    "    net_params['n_classes'] = num_classes\n",
    "    net_params['L'] = L  # min L should be 2\n",
    "    net_params['readout'] = \"mean\"\n",
    "    net_params['layer_norm'] = True\n",
    "    net_params['batch_norm'] = True\n",
    "    net_params['in_feat_dropout'] = 0.0\n",
    "    net_params['dropout'] = 0.0\n",
    "    net_params['residual'] = True\n",
    "    net_params['edge_feat'] = edge_feat\n",
    "    net_params['self_loop'] = self_loop\n",
    "    \n",
    "    # for MLPNet \n",
    "    net_params['gated'] = gated  \n",
    "    \n",
    "    # for GAT\n",
    "    net_params['n_heads'] = n_heads \n",
    "    \n",
    "    # for graphsage\n",
    "    net_params['sage_aggregator'] = 'meanpool' \n",
    "\n",
    "    # specific for GIN\n",
    "    net_params['n_mlp_GIN'] = n_mlp_GIN\n",
    "    net_params['learn_eps_GIN'] = True\n",
    "    net_params['neighbor_aggr_GIN'] = 'sum'\n",
    "    \n",
    "    # specific for MoNet\n",
    "    net_params['pseudo_dim_MoNet'] = pseudo_dim_MoNet\n",
    "    net_params['kernel'] = kernel\n",
    "    \n",
    "    # specific for RingGNN\n",
    "    net_params['radius'] = 2\n",
    "    num_nodes = [dataset.train[i][0].number_of_nodes() for i in range(len(dataset.train))]\n",
    "    net_params['avg_node_num'] = int(np.ceil(np.mean(num_nodes)))\n",
    "    \n",
    "    # specific for 3WLGNN\n",
    "    net_params['depth_of_mlp'] = 2\n",
    "    \n",
    "    # specific for pos_enc_dim\n",
    "    net_params['pos_enc'] = pos_enc\n",
    "    net_params['pos_enc_dim'] = pos_enc_dim\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MODEL DETAILS:\n",
      "\n",
      "MODEL/Total parameters: GatedGCN 104773\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "    VIEWING MODEL CONFIG AND PARAMS\n",
    "\"\"\"\n",
    "def view_model_param(MODEL_NAME, net_params):\n",
    "    model = gnn_model(MODEL_NAME, net_params)\n",
    "    total_param = 0\n",
    "    print(\"MODEL DETAILS:\\n\")\n",
    "    #print(model)\n",
    "    for param in model.parameters():\n",
    "        # print(param.data.size())\n",
    "        total_param += np.prod(list(param.data.size()))\n",
    "    print('MODEL/Total parameters:', MODEL_NAME, total_param)\n",
    "    return total_param\n",
    "\n",
    "\n",
    "if notebook_mode == True:\n",
    "    view_model_param(MODEL_NAME, net_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    TRAINING CODE\n",
    "\"\"\"\n",
    "\n",
    "def train_val_pipeline(MODEL_NAME, dataset, params, net_params, dirs):\n",
    "    \n",
    "    start0 = time.time()\n",
    "    per_epoch_time = []\n",
    "    \n",
    "    DATASET_NAME = dataset.name\n",
    "    \n",
    "    if MODEL_NAME in ['GCN', 'GAT']:\n",
    "        if net_params['self_loop']:\n",
    "            print(\"[!] Adding graph self-loops for GCN/GAT models (central node trick).\")\n",
    "            dataset._add_self_loops()\n",
    "    \n",
    "    if MODEL_NAME in ['GatedGCN']:\n",
    "        if net_params['pos_enc']:\n",
    "            print(\"[!] Adding graph positional encoding.\")\n",
    "            dataset._add_positional_encodings(net_params['pos_enc_dim'])\n",
    "            print('Time PE:',time.time()-start0)\n",
    "        \n",
    "    trainset, valset, testset = dataset.train, dataset.val, dataset.test\n",
    "        \n",
    "    root_log_dir, root_ckpt_dir, write_file_name, write_config_file = dirs\n",
    "    device = net_params['device']\n",
    "    \n",
    "    # Write network and optimization hyper-parameters in folder config/\n",
    "    with open(write_config_file + '.txt', 'w') as f:\n",
    "        f.write(\"\"\"Dataset: {},\\nModel: {}\\n\\nparams={}\\n\\nnet_params={}\\n\\n\\nTotal Parameters: {}\\n\\n\"\"\"\\\n",
    "                .format(DATASET_NAME, MODEL_NAME, params, net_params, net_params['total_param']))\n",
    "        \n",
    "    log_dir = os.path.join(root_log_dir, \"RUN_\" + str(0))\n",
    "    writer = SummaryWriter(log_dir=log_dir)\n",
    "\n",
    "    # setting seeds\n",
    "    random.seed(params['seed'])\n",
    "    np.random.seed(params['seed'])\n",
    "    torch.manual_seed(params['seed'])\n",
    "    if device.type == 'cuda':\n",
    "        torch.cuda.manual_seed(params['seed'])\n",
    "    \n",
    "    print(\"Training Graphs: \", len(trainset))\n",
    "    print(\"Validation Graphs: \", len(valset))\n",
    "    print(\"Test Graphs: \", len(testset))\n",
    "    print(\"Number of Classes: \", net_params['n_classes'])\n",
    "\n",
    "    model = gnn_model(MODEL_NAME, net_params)\n",
    "    model = model.to(device)\n",
    "\n",
    "    optimizer = optim.Adam(model.parameters(), lr=params['init_lr'], weight_decay=params['weight_decay'])\n",
    "    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min',\n",
    "                                                     factor=params['lr_reduce_factor'],\n",
    "                                                     patience=params['lr_schedule_patience'],\n",
    "                                                     verbose=True)\n",
    "    \n",
    "    epoch_train_losses, epoch_val_losses = [], []\n",
    "    epoch_train_accs, epoch_val_accs = [], [] \n",
    "    \n",
    "    if MODEL_NAME in ['RingGNN', '3WLGNN']:\n",
    "        # import train functions specific for WL-GNNs\n",
    "        from train.train_SBMs_node_classification import train_epoch_dense as train_epoch, evaluate_network_dense as evaluate_network\n",
    "        \n",
    "        train_loader = DataLoader(trainset, shuffle=True, collate_fn=dataset.collate_dense_gnn)\n",
    "        val_loader = DataLoader(valset, shuffle=False, collate_fn=dataset.collate_dense_gnn)\n",
    "        test_loader = DataLoader(testset, shuffle=False, collate_fn=dataset.collate_dense_gnn)\n",
    "        \n",
    "    else:\n",
    "        # import train functions for all other GCNs\n",
    "        from train.train_SBMs_node_classification import train_epoch_sparse as train_epoch, evaluate_network_sparse as evaluate_network # import train functions\n",
    "        \n",
    "        train_loader = DataLoader(trainset, batch_size=params['batch_size'], shuffle=True, collate_fn=dataset.collate)\n",
    "        val_loader = DataLoader(valset, batch_size=params['batch_size'], shuffle=False, collate_fn=dataset.collate)\n",
    "        test_loader = DataLoader(testset, batch_size=params['batch_size'], shuffle=False, collate_fn=dataset.collate)\n",
    "        \n",
    "    # At any point you can hit Ctrl + C to break out of training early.\n",
    "    try:\n",
    "        with tqdm(range(params['epochs'])) as t:\n",
    "            for epoch in t:\n",
    "\n",
    "                t.set_description('Epoch %d' % epoch)\n",
    "\n",
    "                start = time.time()\n",
    "\n",
    "                if MODEL_NAME in ['RingGNN', '3WLGNN']: # since different batch training function for dense GNNs\n",
    "                    epoch_train_loss, epoch_train_acc, optimizer = train_epoch(model, optimizer, device, train_loader, epoch, params['batch_size'])\n",
    "                else:   # for all other models common train function\n",
    "                    epoch_train_loss, epoch_train_acc, optimizer = train_epoch(model, optimizer, device, train_loader, epoch)\n",
    "                    \n",
    "                epoch_val_loss, epoch_val_acc = evaluate_network(model, device, val_loader, epoch)\n",
    "                _, epoch_test_acc = evaluate_network(model, device, test_loader, epoch)        \n",
    "                \n",
    "                epoch_train_losses.append(epoch_train_loss)\n",
    "                epoch_val_losses.append(epoch_val_loss)\n",
    "                epoch_train_accs.append(epoch_train_acc)\n",
    "                epoch_val_accs.append(epoch_val_acc)\n",
    "\n",
    "                writer.add_scalar('train/_loss', epoch_train_loss, epoch)\n",
    "                writer.add_scalar('val/_loss', epoch_val_loss, epoch)\n",
    "                writer.add_scalar('train/_acc', epoch_train_acc, epoch)\n",
    "                writer.add_scalar('val/_acc', epoch_val_acc, epoch)\n",
    "                writer.add_scalar('test/_acc', epoch_test_acc, epoch)\n",
    "                writer.add_scalar('learning_rate', optimizer.param_groups[0]['lr'], epoch)\n",
    "\n",
    "                t.set_postfix(time=time.time()-start, lr=optimizer.param_groups[0]['lr'],\n",
    "                              train_loss=epoch_train_loss, val_loss=epoch_val_loss,\n",
    "                              train_acc=epoch_train_acc, val_acc=epoch_val_acc,\n",
    "                              test_acc=epoch_test_acc)\n",
    "\n",
    "                per_epoch_time.append(time.time()-start)\n",
    "\n",
    "                # Saving checkpoint\n",
    "                ckpt_dir = os.path.join(root_ckpt_dir, \"RUN_\")\n",
    "                if not os.path.exists(ckpt_dir):\n",
    "                    os.makedirs(ckpt_dir)\n",
    "                torch.save(model.state_dict(), '{}.pkl'.format(ckpt_dir + \"/epoch_\" + str(epoch)))\n",
    "\n",
    "                files = glob.glob(ckpt_dir + '/*.pkl')\n",
    "                for file in files:\n",
    "                    epoch_nb = file.split('_')[-1]\n",
    "                    epoch_nb = int(epoch_nb.split('.')[0])\n",
    "                    if epoch_nb < epoch-1:\n",
    "                        os.remove(file)\n",
    "\n",
    "                scheduler.step(epoch_val_loss)\n",
    "\n",
    "                if optimizer.param_groups[0]['lr'] < params['min_lr']:\n",
    "                    print(\"\\n!! LR SMALLER OR EQUAL TO MIN LR THRESHOLD.\")\n",
    "                    break\n",
    "                    \n",
    "                # Stop training after params['max_time'] hours\n",
    "                if time.time()-start0 > params['max_time']*3600:\n",
    "                    print('-' * 89)\n",
    "                    print(\"Max_time for training elapsed {:.2f} hours, so stopping\".format(params['max_time']))\n",
    "                    break\n",
    "    \n",
    "    except KeyboardInterrupt:\n",
    "        print('-' * 89)\n",
    "        print('Exiting from training early because of KeyboardInterrupt')\n",
    "    \n",
    "    \n",
    "    _, test_acc = evaluate_network(model, device, test_loader, epoch)\n",
    "    _, train_acc = evaluate_network(model, device, train_loader, epoch)\n",
    "    print(\"Test Accuracy: {:.4f}\".format(test_acc))\n",
    "    print(\"Train Accuracy: {:.4f}\".format(train_acc))\n",
    "    print(\"Convergence Time (Epochs): {:.4f}\".format(epoch))\n",
    "    print(\"TOTAL TIME TAKEN: {:.4f}s\".format(time.time()-start0))\n",
    "    print(\"AVG TIME PER EPOCH: {:.4f}s\".format(np.mean(per_epoch_time)))\n",
    "\n",
    "    writer.close()\n",
    "\n",
    "    \"\"\"\n",
    "        Write the results in out_dir/results folder\n",
    "    \"\"\"\n",
    "    with open(write_file_name + '.txt', 'w') as f:\n",
    "        f.write(\"\"\"Dataset: {},\\nModel: {}\\n\\nparams={}\\n\\nnet_params={}\\n\\n{}\\n\\nTotal Parameters: {}\\n\\n\n",
    "    FINAL RESULTS\\nTEST ACCURACY: {:.4f}\\nTRAIN ACCURACY: {:.4f}\\n\\n\n",
    "    Convergence Time (Epochs): {:.4f}\\nTotal Time Taken: {:.4f} hrs\\nAverage Time Per Epoch: {:.4f} s\\n\\n\\n\"\"\"\\\n",
    "          .format(DATASET_NAME, MODEL_NAME, params, net_params, model, net_params['total_param'],\n",
    "                  test_acc, train_acc, epoch, (time.time()-start0)/3600, np.mean(per_epoch_time)))\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Convert main_SBMs_node_classification.ipynb to main_SBMs_node_classification.py\n",
      "Clean main_SBMs_node_classification.py\n",
      "Done. \n",
      "[I] Loading dataset SBM_PATTERN...\n",
      "train, test, val sizes : 10000 2000 2000\n",
      "[I] Finished loading.\n",
      "[I] Data load time: 24.4756s\n",
      "cuda not available\n",
      "MODEL DETAILS:\n",
      "\n",
      "MODEL/Total parameters: GatedGCN 104773\n",
      "[!] Adding graph positional encoding.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vijay/miniconda3/envs/benchmark_gnn/lib/python3.7/site-packages/dgl/base.py:25: UserWarning: Currently adjacency_matrix() returns a matrix with destination as rows by default.  In 0.5 the result will have source as rows (i.e. transpose=True)\n",
      "  warnings.warn(msg, warn_type)\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-10-01282266494f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    236\u001b[0m     \u001b[0mcleaner_main\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'main_SBMs_node_classification'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    237\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 238\u001b[0;31m     \u001b[0mmain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mconfig\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    239\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    240\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<ipython-input-10-01282266494f>\u001b[0m in \u001b[0;36mmain\u001b[0;34m(notebook_mode, config)\u001b[0m\n\u001b[1;32m    197\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    198\u001b[0m     \u001b[0mnet_params\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'total_param'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mview_model_param\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mMODEL_NAME\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnet_params\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 199\u001b[0;31m     \u001b[0mtrain_val_pipeline\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mMODEL_NAME\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdataset\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnet_params\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdirs\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    200\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    201\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-9-592fec58e2c8>\u001b[0m in \u001b[0;36mtrain_val_pipeline\u001b[0;34m(MODEL_NAME, dataset, params, net_params, dirs)\u001b[0m\n\u001b[1;32m     18\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mnet_params\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'pos_enc'\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     19\u001b[0m             \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"[!] Adding graph positional encoding.\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 20\u001b[0;31m             \u001b[0mdataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_add_positional_encodings\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnet_params\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'pos_enc_dim'\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     21\u001b[0m             \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Time PE:'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtime\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mstart0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/graphdeeplearning/benchmarking-gnns-dev/data/SBMs.py\u001b[0m in \u001b[0;36m_add_positional_encodings\u001b[0;34m(self, pos_enc_dim)\u001b[0m\n\u001b[1;32m    242\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    243\u001b[0m         \u001b[0;31m# Graph positional encoding v/ Laplacian eigenvectors\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 244\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mpositional_encoding\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpos_enc_dim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mg\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\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    245\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mpositional_encoding\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpos_enc_dim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mg\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    246\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtest\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mpositional_encoding\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpos_enc_dim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mg\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtest\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/graphdeeplearning/benchmarking-gnns-dev/data/SBMs.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    242\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    243\u001b[0m         \u001b[0;31m# Graph positional encoding v/ Laplacian eigenvectors\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 244\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mpositional_encoding\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpos_enc_dim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mg\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\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    245\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mpositional_encoding\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpos_enc_dim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mg\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    246\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtest\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mpositional_encoding\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpos_enc_dim\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mg\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtest\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph_lists\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/graphdeeplearning/benchmarking-gnns-dev/data/SBMs.py\u001b[0m in \u001b[0;36mpositional_encoding\u001b[0;34m(g, pos_enc_dim)\u001b[0m\n\u001b[1;32m    139\u001b[0m     \u001b[0;31m# Eigenvectors with scipy\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    140\u001b[0m     \u001b[0;31m#EigVal, EigVec = sp.linalg.eigs(L, k=pos_enc_dim+1, which='SR')\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 141\u001b[0;31m     \u001b[0mEigVal\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mEigVec\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlinalg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meigs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mL\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpos_enc_dim\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwhich\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'SR'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtol\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1e-2\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# for 40 PEs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    142\u001b[0m     \u001b[0mEigVec\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mEigVec\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mEigVal\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margsort\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;31m# increasing order\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    143\u001b[0m     \u001b[0mg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mndata\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'pos_enc'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_numpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mabs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mEigVec\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[0mpos_enc_dim\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[0mfloat\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~/miniconda3/envs/benchmark_gnn/lib/python3.7/site-packages/scipy/sparse/linalg/eigen/arpack/arpack.py\u001b[0m in \u001b[0;36meigs\u001b[0;34m(A, k, M, sigma, which, v0, ncv, maxiter, tol, return_eigenvectors, Minv, OPinv, OPpart)\u001b[0m\n\u001b[1;32m   1345\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0m_ARPACK_LOCK\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1346\u001b[0m         \u001b[0;32mwhile\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconverged\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1347\u001b[0;31m             \u001b[0mparams\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0miterate\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   1348\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1349\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mextract\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mreturn_eigenvectors\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/benchmark_gnn/lib/python3.7/site-packages/scipy/sparse/linalg/eigen/arpack/arpack.py\u001b[0m in \u001b[0;36miterate\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    725\u001b[0m                                     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtol\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mresid\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mv\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0miparam\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    726\u001b[0m                                     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mipntr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mworkd\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mworkl\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                                     self.info)\n\u001b[0m\u001b[1;32m    728\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    729\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mido\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtol\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mresid\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mv\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0miparam\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mipntr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minfo\u001b[0m \u001b[0;34m=\u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "def main(notebook_mode=False,config=None):\n",
    "    \n",
    "    \"\"\"\n",
    "        USER CONTROLS\n",
    "    \"\"\"\n",
    "    \n",
    "    # terminal mode\n",
    "    if notebook_mode==False:\n",
    "        \n",
    "        parser = argparse.ArgumentParser()\n",
    "        parser.add_argument('--config', help=\"Please give a config.json file with training/model/data/param details\")\n",
    "        parser.add_argument('--gpu_id', help=\"Please give a value for gpu id\")\n",
    "        parser.add_argument('--model', help=\"Please give a value for model name\")\n",
    "        parser.add_argument('--dataset', help=\"Please give a value for dataset name\")\n",
    "        parser.add_argument('--out_dir', help=\"Please give a value for out_dir\")\n",
    "        parser.add_argument('--seed', help=\"Please give a value for seed\")\n",
    "        parser.add_argument('--epochs', help=\"Please give a value for epochs\")\n",
    "        parser.add_argument('--batch_size', help=\"Please give a value for batch_size\")\n",
    "        parser.add_argument('--init_lr', help=\"Please give a value for init_lr\")\n",
    "        parser.add_argument('--lr_reduce_factor', help=\"Please give a value for lr_reduce_factor\")\n",
    "        parser.add_argument('--lr_schedule_patience', help=\"Please give a value for lr_schedule_patience\")\n",
    "        parser.add_argument('--min_lr', help=\"Please give a value for min_lr\")\n",
    "        parser.add_argument('--weight_decay', help=\"Please give a value for weight_decay\")\n",
    "        parser.add_argument('--print_epoch_interval', help=\"Please give a value for print_epoch_interval\")    \n",
    "        parser.add_argument('--L', help=\"Please give a value for L\")\n",
    "        parser.add_argument('--hidden_dim', help=\"Please give a value for hidden_dim\")\n",
    "        parser.add_argument('--out_dim', help=\"Please give a value for out_dim\")\n",
    "        parser.add_argument('--residual', help=\"Please give a value for residual\")\n",
    "        parser.add_argument('--edge_feat', help=\"Please give a value for edge_feat\")\n",
    "        parser.add_argument('--readout', help=\"Please give a value for readout\")\n",
    "        parser.add_argument('--kernel', help=\"Please give a value for kernel\")\n",
    "        parser.add_argument('--n_heads', help=\"Please give a value for n_heads\")\n",
    "        parser.add_argument('--gated', help=\"Please give a value for gated\")\n",
    "        parser.add_argument('--in_feat_dropout', help=\"Please give a value for in_feat_dropout\")\n",
    "        parser.add_argument('--dropout', help=\"Please give a value for dropout\")\n",
    "        parser.add_argument('--layer_norm', help=\"Please give a value for layer_norm\")\n",
    "        parser.add_argument('--batch_norm', help=\"Please give a value for batch_norm\")\n",
    "        parser.add_argument('--sage_aggregator', help=\"Please give a value for sage_aggregator\")\n",
    "        parser.add_argument('--data_mode', help=\"Please give a value for data_mode\")\n",
    "        parser.add_argument('--num_pool', help=\"Please give a value for num_pool\")\n",
    "        parser.add_argument('--gnn_per_block', help=\"Please give a value for gnn_per_block\")\n",
    "        parser.add_argument('--embedding_dim', help=\"Please give a value for embedding_dim\")\n",
    "        parser.add_argument('--pool_ratio', help=\"Please give a value for pool_ratio\")\n",
    "        parser.add_argument('--linkpred', help=\"Please give a value for linkpred\")\n",
    "        parser.add_argument('--cat', help=\"Please give a value for cat\")\n",
    "        parser.add_argument('--self_loop', help=\"Please give a value for self_loop\")\n",
    "        parser.add_argument('--max_time', help=\"Please give a value for max_time\")\n",
    "        parser.add_argument('--pos_enc_dim', help=\"Please give a value for pos_enc_dim\")\n",
    "        parser.add_argument('--pos_enc', help=\"Please give a value for pos_enc\")\n",
    "        args = parser.parse_args()\n",
    "        with open(args.config) as f:\n",
    "            config = json.load(f)\n",
    "            \n",
    "\n",
    "        # device\n",
    "        if args.gpu_id is not None:\n",
    "            config['gpu']['id'] = int(args.gpu_id)\n",
    "            config['gpu']['use'] = True\n",
    "        device = gpu_setup(config['gpu']['use'], config['gpu']['id'])\n",
    "\n",
    "        # model, dataset, out_dir\n",
    "        if args.model is not None:\n",
    "            MODEL_NAME = args.model\n",
    "        else:\n",
    "            MODEL_NAME = config['model']\n",
    "        if args.dataset is not None:\n",
    "            DATASET_NAME = args.dataset\n",
    "        else:\n",
    "            DATASET_NAME = config['dataset']\n",
    "        dataset = LoadData(DATASET_NAME)\n",
    "        if args.out_dir is not None:\n",
    "            out_dir = args.out_dir\n",
    "        else:\n",
    "            out_dir = config['out_dir']\n",
    "\n",
    "        # parameters\n",
    "        params = config['params']\n",
    "        if args.seed is not None:\n",
    "            params['seed'] = int(args.seed)\n",
    "        if args.epochs is not None:\n",
    "            params['epochs'] = int(args.epochs)\n",
    "        if args.batch_size is not None:\n",
    "            params['batch_size'] = int(args.batch_size)\n",
    "        if args.init_lr is not None:\n",
    "            params['init_lr'] = float(args.init_lr)\n",
    "        if args.lr_reduce_factor is not None:\n",
    "            params['lr_reduce_factor'] = float(args.lr_reduce_factor)\n",
    "        if args.lr_schedule_patience is not None:\n",
    "            params['lr_schedule_patience'] = int(args.lr_schedule_patience)\n",
    "        if args.min_lr is not None:\n",
    "            params['min_lr'] = float(args.min_lr)\n",
    "        if args.weight_decay is not None:\n",
    "            params['weight_decay'] = float(args.weight_decay)\n",
    "        if args.print_epoch_interval is not None:\n",
    "            params['print_epoch_interval'] = int(args.print_epoch_interval)\n",
    "        if args.max_time is not None:\n",
    "            params['max_time'] = float(args.max_time)\n",
    "\n",
    "        # network parameters\n",
    "        net_params = config['net_params']\n",
    "        net_params['device'] = device\n",
    "        net_params['gpu_id'] = config['gpu']['id']\n",
    "        net_params['batch_size'] = params['batch_size']\n",
    "        if args.L is not None:\n",
    "            net_params['L'] = int(args.L)\n",
    "        if args.hidden_dim is not None:\n",
    "            net_params['hidden_dim'] = int(args.hidden_dim)\n",
    "        if args.out_dim is not None:\n",
    "            net_params['out_dim'] = int(args.out_dim)   \n",
    "        if args.residual is not None:\n",
    "            net_params['residual'] = True if args.residual=='True' else False\n",
    "        if args.edge_feat is not None:\n",
    "            net_params['edge_feat'] = True if args.edge_feat=='True' else False\n",
    "        if args.readout is not None:\n",
    "            net_params['readout'] = args.readout\n",
    "        if args.kernel is not None:\n",
    "            net_params['kernel'] = int(args.kernel)\n",
    "        if args.n_heads is not None:\n",
    "            net_params['n_heads'] = int(args.n_heads)\n",
    "        if args.gated is not None:\n",
    "            net_params['gated'] = True if args.gated=='True' else False\n",
    "        if args.in_feat_dropout is not None:\n",
    "            net_params['in_feat_dropout'] = float(args.in_feat_dropout)\n",
    "        if args.dropout is not None:\n",
    "            net_params['dropout'] = float(args.dropout)\n",
    "        if args.layer_norm is not None:\n",
    "            net_params['layer_norm'] = True if args.layer_norm=='True' else False\n",
    "        if args.batch_norm is not None:\n",
    "            net_params['batch_norm'] = True if args.batch_norm=='True' else False\n",
    "        if args.sage_aggregator is not None:\n",
    "            net_params['sage_aggregator'] = args.sage_aggregator\n",
    "        if args.data_mode is not None:\n",
    "            net_params['data_mode'] = args.data_mode\n",
    "        if args.num_pool is not None:\n",
    "            net_params['num_pool'] = int(args.num_pool)\n",
    "        if args.gnn_per_block is not None:\n",
    "            net_params['gnn_per_block'] = int(args.gnn_per_block)\n",
    "        if args.embedding_dim is not None:\n",
    "            net_params['embedding_dim'] = int(args.embedding_dim)\n",
    "        if args.pool_ratio is not None:\n",
    "            net_params['pool_ratio'] = float(args.pool_ratio)\n",
    "        if args.linkpred is not None:\n",
    "            net_params['linkpred'] = True if args.linkpred=='True' else False\n",
    "        if args.cat is not None:\n",
    "            net_params['cat'] = True if args.cat=='True' else False\n",
    "        if args.self_loop is not None:\n",
    "            net_params['self_loop'] = True if args.self_loop=='True' else False\n",
    "        if args.pos_enc is not None:\n",
    "            net_params['pos_enc'] = True if args.pos_enc=='True' else False\n",
    "        if args.pos_enc_dim is not None:\n",
    "            net_params['pos_enc_dim'] = int(args.pos_enc_dim)\n",
    "\n",
    "            \n",
    "    # notebook mode\n",
    "    if notebook_mode:\n",
    "        \n",
    "        # parameters\n",
    "        params = config['params']\n",
    "        \n",
    "        # dataset\n",
    "        DATASET_NAME = config['dataset']\n",
    "        dataset = LoadData(DATASET_NAME)\n",
    "        \n",
    "        # device\n",
    "        device = gpu_setup(config['gpu']['use'], config['gpu']['id'])\n",
    "        out_dir = config['out_dir']\n",
    "        \n",
    "        # GNN model\n",
    "        MODEL_NAME = config['model']\n",
    "        \n",
    "        # network parameters\n",
    "        net_params = config['net_params']\n",
    "        net_params['device'] = device\n",
    "        net_params['gpu_id'] = config['gpu']['id']\n",
    "        net_params['batch_size'] = params['batch_size']\n",
    "        \n",
    "            \n",
    "    # SBM\n",
    "    net_params['in_dim'] = torch.unique(dataset.train[0][0].ndata['feat'],dim=0).size(0) # node_dim (feat is an integer)\n",
    "    net_params['n_classes'] = torch.unique(dataset.train[0][1],dim=0).size(0)\n",
    "    \n",
    "    if MODEL_NAME == 'RingGNN':\n",
    "        num_nodes = [dataset.train[i][0].number_of_nodes() for i in range(len(dataset.train))]\n",
    "        net_params['avg_node_num'] = int(np.ceil(np.mean(num_nodes)))\n",
    "\n",
    "    root_log_dir = out_dir + 'logs/' + MODEL_NAME + \"_\" + DATASET_NAME + \"_GPU\" + str(config['gpu']['id']) + \"_\" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y')\n",
    "    root_ckpt_dir = out_dir + 'checkpoints/' + MODEL_NAME + \"_\" + DATASET_NAME + \"_GPU\" + str(config['gpu']['id']) + \"_\" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y')\n",
    "    write_file_name = out_dir + 'results/result_' + MODEL_NAME + \"_\" + DATASET_NAME + \"_GPU\" + str(config['gpu']['id']) + \"_\" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y')\n",
    "    write_config_file = out_dir + 'configs/config_' + MODEL_NAME + \"_\" + DATASET_NAME + \"_GPU\" + str(config['gpu']['id']) + \"_\" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y')\n",
    "    dirs = root_log_dir, root_ckpt_dir, write_file_name, write_config_file\n",
    "\n",
    "    if not os.path.exists(out_dir + 'results'):\n",
    "        os.makedirs(out_dir + 'results')\n",
    "        \n",
    "    if not os.path.exists(out_dir + 'configs'):\n",
    "        os.makedirs(out_dir + 'configs')\n",
    "\n",
    "    net_params['total_param'] = view_model_param(MODEL_NAME, net_params)\n",
    "    train_val_pipeline(MODEL_NAME, dataset, params, net_params, dirs)\n",
    "\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "if notebook_mode==True:\n",
    "    \n",
    "    config = {}\n",
    "    # gpu config\n",
    "    gpu = {}\n",
    "    gpu['use'] = use_gpu\n",
    "    gpu['id'] = gpu_id\n",
    "    config['gpu'] = gpu\n",
    "    # GNN model, dataset, out_dir\n",
    "    config['model'] = MODEL_NAME\n",
    "    config['dataset'] = DATASET_NAME\n",
    "    config['out_dir'] = out_dir\n",
    "    # parameters\n",
    "    params = {}\n",
    "    params['seed'] = seed\n",
    "    params['epochs'] = epochs\n",
    "    params['batch_size'] = batch_size\n",
    "    params['init_lr'] = init_lr\n",
    "    params['lr_reduce_factor'] = lr_reduce_factor \n",
    "    params['lr_schedule_patience'] = lr_schedule_patience\n",
    "    params['min_lr'] = min_lr\n",
    "    params['weight_decay'] = weight_decay\n",
    "    params['print_epoch_interval'] = 5\n",
    "    params['max_time'] = max_time\n",
    "    config['params'] = params\n",
    "    # network parameters\n",
    "    config['net_params'] = net_params\n",
    "    \n",
    "    # convert to .py format\n",
    "    from utils.cleaner_main import *\n",
    "    cleaner_main('main_SBMs_node_classification')\n",
    "    \n",
    "    main(True,config)\n",
    "    \n",
    "else:\n",
    "    \n",
    "    main()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}