{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\"\n",
    "import sys\n",
    "sys.path.append(\"../\")\n",
    "import torch\n",
    "assert torch.cuda.is_available()\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import torch.nn as nn\n",
    "from tqdm.auto import tqdm\n",
    "import json\n",
    "from train import eval\n",
    "from Utils import load\n",
    "import networkx as nx\n",
    "import scipy\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pruners = np.array(\n",
    "    [\n",
    "        'cs',\n",
    "        'grasp',\n",
    "        'sf',\n",
    "        'mica_erk_shuffle',\n",
    "        'mica_igq_shuffle',\n",
    "        'mica_cs_rand_effective_comp',\n",
    "        'mica_grasp_rand_effective_comp',\n",
    "        'mica_sf_rand_effective_comp',\n",
    "        'mica_erk_rand_effective_comp',\n",
    "        'mica_igq_rand_effective_comp'\n",
    "    ],\n",
    "    dtype=object\n",
    "    )\n",
    "\n",
    "colors = {\n",
    "    'cs'         : plt.cm.tab10(9),\n",
    "    'grasp'      : plt.cm.tab10(8),\n",
    "    'sf'         : plt.cm.tab10(6),\n",
    "    'mica_cs_rand'    : plt.cm.tab10(9),\n",
    "    'mica_grasp_rand' : plt.cm.tab10(8),\n",
    "    'mica_sf'    : plt.cm.tab10(6),\n",
    "    'mica_sf_rand' : plt.cm.tab10(6),\n",
    "    'mica_sf_min'  : plt.cm.tab10(6),\n",
    "    'mica_sf_max'  : plt.cm.tab10(6),\n",
    "    'mica_erk'   : plt.cm.tab10(4),\n",
    "    'mica_erk_rand' : plt.cm.tab10(4),\n",
    "    'mica_igq'   : plt.cm.tab10(3),\n",
    "    'mica_igq_rand' : plt.cm.tab10(3),\n",
    "    'mica_igq_min'  : plt.cm.tab10(8),\n",
    "    'mica_igq_max'  : plt.cm.tab10(7),\n",
    "}\n",
    "\n",
    "markers = {\n",
    "    'cs'         : \"^\", \n",
    "    'grasp'      : \"h\", \n",
    "    'sf'         : \"H\", \n",
    "    'mica_cs_rand'    : \"^\", \n",
    "    'mica_grasp_rand' : \"h\", \n",
    "    'mica_sf_rand'    : \"H\",\n",
    "    'mica_sf_min'    : \"H\",\n",
    "    'mica_sf_max'    : \"H\",\n",
    "    'mica_erk'   : \"p\", \n",
    "    'mica_erk_rand' : \"p\", \n",
    "    'mica_igq'   : \"*\",\n",
    "    'mica_igq_rand' : \"*\",\n",
    "    'mica_igq_min'  : \"<\",\n",
    "    'mica_igq_max'  : \">\",\n",
    "}\n",
    "\n",
    "pruner_names = {\n",
    "    'cs'                 : 'RbPI-SNIP',\n",
    "    'grasp'              : 'RbPI-GraSP',\n",
    "    'sf'                 : 'RbPI-SynFlow',\n",
    "    'mica_cs_shuffle'    : 'RPI-SNIP',\n",
    "    'mica_grasp_shuffle' : 'RPI-GraSP',\n",
    "    'mica_sf_shuffle'    : 'RPI-SynFlow',\n",
    "    'mica_erk_shuffle'   : 'RPI-ERK',\n",
    "    'mica_igq_shuffle'   : 'RPI-IGQ',\n",
    "    'mica_cs_rand_effective_comp'    : '$\\\\bf{MiCA}$-$\\\\bf{SNIP}$',\n",
    "    'mica_grasp_rand_effective_comp' : '$\\\\bf{MiCA}$-$\\\\bf{GraSP}$',\n",
    "    'mica_sf_rand_effective_comp'    : '$\\\\bf{MiCA}$-$\\\\bf{SynFlow}$',\n",
    "    'mica_erk_rand_effective_comp'   : '$\\\\bf{MiCA}$-$\\\\bf{ERK}$',\n",
    "    'mica_igq_rand_effective_comp'   : '$\\\\bf{MiCA}$-$\\\\bf{IGQ}$',\n",
    "    }\n",
    "\n",
    "pepochs = {\n",
    "    'cs'         : '1',\n",
    "    'grasp'      : '1',\n",
    "    'sf'         : '100',\n",
    "    'mica_cs_rand'    : '1',\n",
    "    'mica_grasp_rand' : '1',\n",
    "    'mica_sf'    : '1',\n",
    "    'mica_sf_rand' : '1',\n",
    "    'mica_sf_max'  : '1',\n",
    "    'mica_sf_min'  : '1',\n",
    "    'mica_erk'   : '1',\n",
    "    'mica_erk_rand'   : '1',\n",
    "    'mica_igq'   : '1',\n",
    "    'mica_igq_rand'  : '1',\n",
    "    'mica_igq_min'   : '1',\n",
    "    'mica_igq_max'   : '1',\n",
    "}\n",
    "\n",
    "for pruner in pruners:\n",
    "    if '_shuffle' in pruner:\n",
    "        colors[pruner]  = colors[pruner.replace('_shuffle', '')]\n",
    "        markers[pruner]  = markers[pruner.replace('_shuffle', '')]\n",
    "        pepochs[pruner] = pepochs[pruner.replace('_shuffle', '')]\n",
    "    elif '_effective_comp' in pruner:\n",
    "        colors[pruner]  = colors[pruner.replace('_effective_comp', '')]\n",
    "        markers[pruner]  = markers[pruner.replace('_effective_comp', '')]\n",
    "        pepochs[pruner] = pepochs[pruner.replace('_effective_comp', '')]\n",
    "    elif '_rand' in pruner:\n",
    "        colors[pruner]  = colors[pruner.replace('_rand', '')]\n",
    "        markers[pruner]  = markers[pruner.replace('_rand', '')]\n",
    "        pepochs[pruner] = pepochs[pruner.replace('_rand', '')]\n",
    "\n",
    "\n",
    "runs = [0, 1, 2]\n",
    "max_run_num = len(runs)\n",
    "\n",
    "print(pruners)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_experimentid(pruner, ratio, pepochs, run_num, expid):\n",
    "    if '_shuffle' in pruner:\n",
    "        base_dir = f'../Results/data/singleshot/{expid}_shuffle'\n",
    "        p_name = pruner.replace('_shuffle', '')\n",
    "    elif '_min' in pruner:\n",
    "        if '_effective_comp' in pruner:\n",
    "            base_dir = f'../Results/data/singleshot/{expid}_min_effective_comp'\n",
    "            p_name = pruner.replace('_min_effective_comp', '')\n",
    "        else:\n",
    "            base_dir = f'../Results/data/singleshot/{expid}_min'\n",
    "            p_name = pruner.replace('_min', '')\n",
    "    elif '_max' in pruner:\n",
    "        if '_effective_comp' in pruner:\n",
    "            base_dir = f'../Results/data/singleshot/{expid}_max_effective_comp'\n",
    "            p_name = pruner.replace('_max_effective_comp', '')\n",
    "        else:\n",
    "            base_dir = f'../Results/data/singleshot/{expid}_max'\n",
    "            p_name = pruner.replace('_max', '')\n",
    "    elif '_rand' in pruner:\n",
    "        if '_effective_comp' in pruner:\n",
    "            base_dir = f'../Results/data/singleshot/{expid}_rand_effective_comp'\n",
    "            p_name = pruner.replace('_rand_effective_comp', '')\n",
    "        else:\n",
    "            base_dir = f'../Results/data/singleshot/{expid}_rand'\n",
    "            p_name = pruner.replace('_rand', '')\n",
    "    else:\n",
    "        base_dir = f'../Results/data/singleshot/{expid}'\n",
    "        p_name = pruner\n",
    "    b_dir  = f'{base_dir}'\n",
    "    experimentid =  f'{b_dir}/' \\\n",
    "                    f'{p_name}-{ratio}-{pepochs}/' \\\n",
    "                    f'run_{run_num}/' \n",
    "    return experimentid\n",
    "\n",
    "def zero_to_nan(values):\n",
    "    return [float('nan') if x==0 else x for x in values]\n",
    "\n",
    "def get_ordered_list(l, order):\n",
    "    out = []\n",
    "    for i in order:\n",
    "        out.append(l[i])\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expid = 'cifar10_lottery_resnet20'\n",
    "ratios = [\n",
    "    '0.5', '1.0', '1.5', '2.0', \n",
    "    '2.5', '3.0', '3.5', '4.0', \n",
    "    ]\n",
    "resnet_baselines = []\n",
    "for r, run_num in enumerate(runs):\n",
    "    experimentid      = get_experimentid('mica_igq_max', '0.0', '1', run_num, expid)\n",
    "    train_result_file = experimentid + 'post-train.pkl'\n",
    "    train_result      = pd.read_pickle(train_result_file)\n",
    "    resnet_baselines.append(np.max(train_result['test_top1_acc']))\n",
    "resnet_baseline = np.mean(resnet_baselines)\n",
    "print(resnet_baseline)\n",
    "\n",
    "resnet_x_ratios           = {}\n",
    "resnet_effective_x_ratios = {}\n",
    "resnet_accs               = {}\n",
    "\n",
    "for l, pruner in enumerate(pruners):\n",
    "    resnet_x_ratios[pruner]           = np.zeros((len(ratios)))\n",
    "    resnet_effective_x_ratios[pruner] = np.zeros((len(ratios)))\n",
    "    resnet_accs[pruner]               = np.zeros((len(ratios)))\n",
    "    for i, ratio in enumerate(tqdm(ratios, desc=pruner)): \n",
    "        x_ratio   = []\n",
    "        effective_x_ratio = []\n",
    "        acc       = []\n",
    "        n_errors  = 0\n",
    "        for r, run_num in enumerate(runs):\n",
    "            experimentid      = get_experimentid(pruner, ratio, pepochs[pruner], run_num, expid)\n",
    "            model_file        = experimentid + 'model.pt'\n",
    "            train_result_file = experimentid + 'post-train.pkl'\n",
    "            compression_file  = experimentid + 'compression.pkl'\n",
    "            try:\n",
    "                train_result      = pd.read_pickle(train_result_file)\n",
    "            except:\n",
    "                n_errors += 1\n",
    "                continue\n",
    "            compression       = pd.read_pickle(compression_file)\n",
    "            prunable          = compression['prunable']==True\n",
    "            weight            = compression['param']=='weight'\n",
    "            modules           = compression[np.logical_and(prunable, weight)]['module']\n",
    "            densities         = np.array(compression[np.logical_and(prunable, weight)]['density'])\n",
    "            shape             = compression[np.logical_and(prunable,weight)]['shape']\n",
    "            e                 = np.array([np.prod(s) for s in shape]) * (densities)\n",
    "            \n",
    "            if np.all(e == 0):                \n",
    "                actual_comp = float('nan')\n",
    "            else:\n",
    "                actual_comp = np.sum(np.array([np.prod(s) for s in shape])) / np.sum(e)\n",
    "            x_ratio.append(actual_comp)\n",
    "            acc.append(np.max(train_result['test_top1_acc']))\n",
    "\n",
    "            with open(experimentid + 'args.json') as f:\n",
    "                args = json.load(f)\n",
    "            input_shape, num_classes = load.dimension(args['dataset']) \n",
    "            loss = nn.CrossEntropyLoss()\n",
    "            if i == 0 and r == 0:\n",
    "                test_loader, _, _ = load.dataloader(\n",
    "                    args['dataset'], args['test_batch_size'], args['test_batch_size'],\n",
    "                    False, args['workers'], 0, 0)\n",
    "                dataset = args['dataset']\n",
    "                test_batch_size = args['test_batch_size']\n",
    "                workers = args['workers']\n",
    "                (data, _) = next(iter(test_loader))\n",
    "                input_dim = list(data[0,:].shape)\n",
    "                input = torch.ones([1] + input_dim).to('cuda')\n",
    "            assert args['seed'] == run_num * 1000 # adhoc\n",
    "            assert args['run_number'] == str(run_num)\n",
    "            assert dataset == args['dataset']\n",
    "            assert test_batch_size == args['test_batch_size']\n",
    "            assert workers == args['workers']\n",
    "\n",
    "            model = load.model(args['model'], args['model_class'])(\n",
    "                input_shape, num_classes, args['dense_classifier'], args['pretrained']).to('cuda')\n",
    "            model.load_state_dict(torch.load(model_file))\n",
    "            print(f'Forward :')\n",
    "            model.set_effective_sparsity(is_forward=True, x=input)\n",
    "            print(f'Backward :')\n",
    "            model.set_effective_sparsity(is_forward=False, x=None)\n",
    "\n",
    "            effective_densities = []\n",
    "            set_modules = set(modules.unique())\n",
    "            for n, m in model.named_modules():\n",
    "                if hasattr(m, 'weight_mask') and n in set_modules:\n",
    "                    effective_densities.append((m.weight_mask.sum() / m.weight_mask.numel()).to('cpu'))\n",
    "            effective_densities = np.array(effective_densities)\n",
    "            e = np.array([np.prod(s) for s in shape]) * (effective_densities)\n",
    "            if np.all(e == 0):                \n",
    "                actual_comp = float('nan')\n",
    "            else:\n",
    "                actual_comp = np.sum(np.array([np.prod(s) for s in shape])) / np.sum(e)\n",
    "            effective_x_ratio.append(actual_comp)\n",
    "        if n_errors == len(runs):\n",
    "            pass\n",
    "        elif n_errors == 0:\n",
    "            resnet_x_ratios[pruner][i]           = np.mean(x_ratio)    \n",
    "            resnet_effective_x_ratios[pruner][i] = np.mean(effective_x_ratio)\n",
    "            resnet_accs[pruner][i]               = np.mean(acc)  \n",
    "        else:\n",
    "            raise ValueError\n",
    "    resnet_x_ratios[pruner]           = resnet_x_ratios[pruner][resnet_x_ratios[pruner] != 0] \n",
    "    resnet_effective_x_ratios[pruner] = resnet_effective_x_ratios[pruner][resnet_effective_x_ratios[pruner] != 0]\n",
    "    resnet_accs[pruner]               = resnet_accs[pruner][resnet_accs[pruner] != 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expid = 'cifar10_lottery_vgg16_bn'\n",
    "ratios = [\n",
    "    '0.5', '1.0', '1.5', \n",
    "    '2.0', '2.5', '3.0', \n",
    "    '3.5', '4.0', '4.5',\n",
    "    '5.0', '5.5', '6.0'\n",
    "    ]\n",
    "vgg_baselines = []\n",
    "for r, run_num in enumerate(runs):\n",
    "    experimentid      = get_experimentid('mica_igq_min', '0.0', '1', run_num, expid)\n",
    "    train_result_file = experimentid + 'post-train.pkl'\n",
    "    train_result      = pd.read_pickle(train_result_file)\n",
    "    vgg_baselines.append(np.max(train_result['test_top1_acc']))\n",
    "vgg_baseline = np.mean(vgg_baselines)\n",
    "print(vgg_baseline)\n",
    "\n",
    "vgg_x_ratios           = {}\n",
    "vgg_effective_x_ratios = {}\n",
    "vgg_accs               = {}\n",
    "\n",
    "for l, pruner in enumerate(pruners):\n",
    "    vgg_x_ratios[pruner]           = np.zeros((len(ratios)))\n",
    "    vgg_effective_x_ratios[pruner] = np.zeros((len(ratios)))\n",
    "    vgg_accs[pruner]               = np.zeros((len(ratios)))\n",
    "    for i, ratio in enumerate(tqdm(ratios, desc=pruner)): \n",
    "        x_ratio   = []\n",
    "        effective_x_ratio = []\n",
    "        acc       = []\n",
    "        n_errors  = 0\n",
    "        for r, run_num in enumerate(runs):\n",
    "            experimentid      = get_experimentid(pruner, ratio, pepochs[pruner], run_num, expid)\n",
    "            model_file        = experimentid + 'model.pt'\n",
    "            train_result_file = experimentid + 'post-train.pkl'\n",
    "            compression_file  = experimentid + 'compression.pkl'\n",
    "            try:\n",
    "                train_result      = pd.read_pickle(train_result_file)\n",
    "            except:\n",
    "                n_errors += 1\n",
    "                continue\n",
    "            compression       = pd.read_pickle(compression_file)\n",
    "            prunable          = compression['prunable']==True\n",
    "            weight            = compression['param']=='weight'\n",
    "            modules           = compression[np.logical_and(prunable, weight)]['module']\n",
    "            densities         = np.array(compression[np.logical_and(prunable, weight)]['density'])\n",
    "            shape             = compression[np.logical_and(prunable,weight)]['shape']\n",
    "            e                 = np.array([np.prod(s) for s in shape]) * (densities)\n",
    "            \n",
    "            if np.all(e == 0):                \n",
    "                actual_comp = float('nan')\n",
    "            else:\n",
    "                actual_comp = np.sum(np.array([np.prod(s) for s in shape])) / np.sum(e)\n",
    "            x_ratio.append(actual_comp)\n",
    "            acc.append(np.max(train_result['test_top1_acc']))\n",
    "\n",
    "            with open(experimentid + 'args.json') as f:\n",
    "                args = json.load(f)\n",
    "            input_shape, num_classes = load.dimension(args['dataset']) \n",
    "            loss = nn.CrossEntropyLoss()\n",
    "            if i == 0 and r == 0:\n",
    "                test_loader, _, _ = load.dataloader(\n",
    "                    args['dataset'], args['test_batch_size'], args['test_batch_size'],\n",
    "                    False, args['workers'], 0, 0)\n",
    "                dataset = args['dataset']\n",
    "                test_batch_size = args['test_batch_size']\n",
    "                workers = args['workers']\n",
    "                (data, _) = next(iter(test_loader))\n",
    "                input_dim = list(data[0,:].shape)\n",
    "                input = torch.ones([1] + input_dim).to('cuda')\n",
    "            assert args['seed'] == run_num * 1000 # adhoc\n",
    "            assert args['run_number'] == str(run_num)\n",
    "            assert dataset == args['dataset']\n",
    "            assert test_batch_size == args['test_batch_size']\n",
    "            assert workers == args['workers']\n",
    "\n",
    "            model = load.model(args['model'], args['model_class'])(\n",
    "                input_shape, num_classes, args['dense_classifier'], args['pretrained']).to('cuda')\n",
    "            model.load_state_dict(torch.load(model_file))\n",
    "            print(f'Forward :')\n",
    "            model.set_effective_sparsity(is_forward=True, x=input)\n",
    "            print(f'Backward :')\n",
    "            model.set_effective_sparsity(is_forward=False, x=None)\n",
    "\n",
    "            effective_densities = []\n",
    "            set_modules = set(modules.unique())\n",
    "            for n, m in model.named_modules():\n",
    "                if hasattr(m, 'weight_mask') and n in set_modules:\n",
    "                    effective_densities.append((m.weight_mask.sum() / m.weight_mask.numel()).to('cpu'))\n",
    "            effective_densities = np.array(effective_densities)\n",
    "            e = np.array([np.prod(s) for s in shape]) * (effective_densities)\n",
    "            if np.all(e == 0):                \n",
    "                actual_comp = float('nan')\n",
    "            else:\n",
    "                actual_comp = np.sum(np.array([np.prod(s) for s in shape])) / np.sum(e)\n",
    "            effective_x_ratio.append(actual_comp)\n",
    "        if n_errors == len(runs):\n",
    "            pass\n",
    "        elif n_errors == 0:\n",
    "            vgg_x_ratios[pruner][i]           = np.mean(x_ratio)    \n",
    "            vgg_effective_x_ratios[pruner][i] = np.mean(effective_x_ratio)\n",
    "            vgg_accs[pruner][i]               = np.mean(acc)  \n",
    "        else:\n",
    "            raise ValueError\n",
    "    vgg_x_ratios[pruner]           = vgg_x_ratios[pruner][vgg_x_ratios[pruner] != 0] \n",
    "    vgg_effective_x_ratios[pruner] = vgg_effective_x_ratios[pruner][vgg_effective_x_ratios[pruner] != 0]\n",
    "    vgg_accs[pruner]               = vgg_accs[pruner][vgg_accs[pruner] != 0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MiCA vs Random Pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "lw=4\n",
    "markersize=30\n",
    "fontsize=35\n",
    "plt.rcParams[\"font.family\"] = \"Segoe UI\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mica_erk_rand_effective_comp\n",
      "vgg\n",
      "[3.16227995e+00 1.00000094e+01 3.16229545e+01 9.99985334e+01\n",
      " 3.16232955e+02 1.10579503e+03 4.78280038e+04            nan\n",
      "            nan            nan            nan            nan]\n",
      "[92.46       91.78       89.47333333 87.00666667 82.58666667 76.84333333\n",
      " 39.30333333 10.         10.         10.         10.         10.        ]\n",
      "mica_igq_rand_effective_comp\n",
      "vgg\n",
      "[3.16228112e+00 1.00000640e+01 3.16231584e+01 1.00004648e+02\n",
      " 3.16543622e+02 1.16415660e+03 1.79843963e+04]\n",
      "[92.37       91.94666667 90.91       88.68333333 85.40666667 79.89\n",
      " 22.33333333]\n",
      "mica_erk_shuffle\n",
      "vgg\n",
      "[3.16227995e+00 1.00000094e+01 3.16229545e+01 9.99985334e+01\n",
      " 3.16232955e+02 1.10579503e+03 4.78280038e+04            nan\n",
      "            nan            nan            nan            nan]\n",
      "[92.58       91.71333333 89.58666667 86.74333333 81.49333333 67.03666667\n",
      " 12.33       10.         10.         10.         10.         10.        ]\n",
      "mica_igq_shuffle\n",
      "vgg\n",
      "[3.16228112e+00 1.00000640e+01 3.16231584e+01 1.00004648e+02\n",
      " 3.16543622e+02 1.16415660e+03 1.79843963e+04            nan\n",
      "            nan            nan            nan            nan]\n",
      "[92.49333333 92.02       90.95       88.59       84.32666667 72.83\n",
      " 12.61666667 10.         10.         10.         10.         10.        ]\n",
      "mica_erk_rand_effective_comp\n",
      "resnet\n",
      "[  3.16219777   9.9994832   31.64303256 104.13743011 565.43503988\n",
      "          nan          nan          nan]\n",
      "[88.11666667 85.31       79.46333333 67.13333333 45.32666667 10.\n",
      " 10.         10.        ]\n",
      "mica_igq_rand_effective_comp\n",
      "resnet\n",
      "[  3.16264079  10.00243695  31.67634046 105.59970602 546.44729325\n",
      "          nan          nan          nan]\n",
      "[87.93666667 85.41       80.13       67.28666667 48.04       12.06666667\n",
      " 10.         10.        ]\n",
      "mica_erk_shuffle\n",
      "resnet\n",
      "[  3.16219777   9.9994832   31.64303256 104.13743011 565.43503988\n",
      "          nan          nan          nan]\n",
      "[88.04       85.43666667 79.14       65.32       40.17333333 10.\n",
      " 10.         10.        ]\n",
      "mica_igq_shuffle\n",
      "resnet\n",
      "[  3.16264079  10.00243695  31.67634046 105.59970602 546.44729325\n",
      "          nan          nan          nan]\n",
      "[88.25       85.36666667 80.06       66.89       43.51333333 12.38\n",
      " 10.         10.        ]\n",
      "0, Dense\n",
      "1, $\\bf{MiCA}$-$\\bf{ERK}$\n",
      "2, $\\bf{MiCA}$-$\\bf{IGQ}$\n",
      "3, RPI-ERK\n",
      "4, RPI-IGQ\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The PostScript backend does not support transparency; partially transparent artists will be rendered opaque.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 123544 characters]",
      "text/plain": [
       "<Figure size 1152x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_pruners = [\n",
    "    'mica_erk_rand_effective_comp',\n",
    "    'mica_igq_rand_effective_comp',\n",
    "    'mica_erk_shuffle',\n",
    "    'mica_igq_shuffle',\n",
    "    ]\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(8*2, 8))\n",
    "\n",
    "for i, arch in enumerate(['vgg', 'resnet']):\n",
    "# for i, arch in enumerate(['vgg']):\n",
    "    if arch == 'vgg':\n",
    "        baseline = vgg_baseline\n",
    "        effective_x_ratios = vgg_effective_x_ratios\n",
    "        accs = vgg_accs\n",
    "        xticks = [10**(2*i - 1) for i in range(1,4)]\n",
    "    elif arch == 'resnet':\n",
    "        baseline = resnet_baseline\n",
    "        effective_x_ratios = resnet_effective_x_ratios\n",
    "        accs = resnet_accs\n",
    "        xticks = [10**(2*i - 1) for i in range(1,6)]\n",
    "    ax = axes[i]\n",
    "    ax.axhline(baseline, color = 'black', linewidth = lw, linestyle='dotted', label='Dense', zorder=0)\n",
    "\n",
    "    zorder = 1\n",
    "    for l, pruner in enumerate(plot_pruners):\n",
    "        print(pruner)\n",
    "        print(arch)\n",
    "        print(effective_x_ratios[pruner])\n",
    "        print(accs[pruner])\n",
    "        if '_shuffle' in pruner:\n",
    "            linestyle = 'dashed'  \n",
    "            ax.plot(effective_x_ratios[pruner], accs[pruner], color=colors[pruner], \n",
    "                    linestyle=linestyle, lw=lw,\n",
    "                    marker=markers[pruner], markersize=markersize, markerfacecolor= 'white', markeredgewidth=2,\n",
    "                    label=pruner_names[pruner],\n",
    "                    zorder=zorder)\n",
    "            zorder += 1\n",
    "        elif 'mica' in pruner:\n",
    "            linestyle = 'solid'\n",
    "            ax.plot(effective_x_ratios[pruner], accs[pruner], color=colors[pruner], \n",
    "                    linestyle=linestyle, lw=lw,\n",
    "                    marker=markers[pruner], markersize=markersize,  \n",
    "                    label=pruner_names[pruner],\n",
    "                    zorder=zorder+10)\n",
    "            zorder += 1\n",
    "\n",
    "    ax.set_ylim(7, 96)\n",
    "    ax.set_xscale('log')\n",
    "    ax.set_xticks(xticks)\n",
    "    ax.set_yticks(np.arange(10, 100, 20))\n",
    "    ax.minorticks_off()\n",
    "    \n",
    "    ax.tick_params(\n",
    "        which       = \"both\", bottom    = True,  top   = False,\n",
    "        labelbottom = True,   left      = True,  right = False,\n",
    "        labelleft   = True if arch == 'vgg' else False,   \n",
    "        direction   = 'out',  length    = 7,     width = 1.0,    \n",
    "        pad         = 8,      labelsize = fontsize\n",
    "        )\n",
    "    ax.set_title('VGG-16' if arch == 'vgg' else 'ResNet-20', size=fontsize)\n",
    "    ax.set_xlabel('Compression Ratio', size=fontsize)\n",
    "    if i == 0:\n",
    "        ax.set_ylabel('Top-1 Test Acc. (%)', size=fontsize)\n",
    "        ax.set_xlim(10**0.1, 10**6.5)\n",
    "    else:\n",
    "        ax.set_xlim(10**0.1, 10**4.75)\n",
    "    ax.grid(True)\n",
    "    ax.set_axisbelow(True)\n",
    "    \n",
    "\n",
    "    # boundary and zoom box\n",
    "    for dir in [\"top\", \"bottom\", \"right\", \"left\"]:\n",
    "        ax.spines[dir].set_linewidth(3)\n",
    "    ax.spines[\"right\"].set_linestyle((0,(3,3)))\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['top'].set_visible(False)\n",
    "\n",
    "fig.subplots_adjust(hspace=0.2)\n",
    "fig.tight_layout()\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "for nl, label in enumerate(labels):\n",
    "    print(f'{nl}, {label}')\n",
    "order = [1,2,3,4,0]\n",
    "ordered_handles = get_ordered_list(handles, order)\n",
    "ordered_labels = get_ordered_list(labels, order)\n",
    "plt.legend(\n",
    "    handles=ordered_handles,labels=ordered_labels, \n",
    "    fontsize=fontsize, loc='upper center', \n",
    "    bbox_to_anchor=(-.1, -.25), ncol=3)\n",
    "os.makedirs('Plots', exist_ok=True)\n",
    "plt.savefig('./Plots/cifar10_mica_rand_sparsity_distribution.eps', bbox_inches=\"tight\")\n",
    "plt.savefig('./Plots/cifar10_mica_rand_sparsity_distribution.svg', bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The PostScript backend does not support transparency; partially transparent artists will be rendered opaque.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vgg\n",
      "mica_cs_rand_effective_comp\n",
      "[92.28666667 91.33333333 89.84       87.78       83.94666667 73.29333333\n",
      " 10.         10.         10.         10.         10.         10.        ]\n",
      "[   3.17003865   10.09526711   32.87195132  115.54719817  456.67581796\n",
      " 2121.41761791           nan           nan           nan           nan\n",
      "           nan           nan]\n",
      "vgg\n",
      "cs\n",
      "[92.26       91.28666667 90.04666667 87.89666667 83.32666667 69.86666667\n",
      " 10.         10.         10.         10.         10.         10.        ]\n",
      "[   3.17003865   10.09526711   32.87195132  115.54719817  456.67581796\n",
      " 2121.41761791           nan           nan           nan           nan\n",
      "           nan           nan]\n",
      "vgg\n",
      "mica_grasp_rand_effective_comp\n",
      "[92.22333333 91.63333333 90.50333333 88.90666667 85.91666667 10.\n",
      " 10.         10.         10.         10.         10.         10.        ]\n",
      "[  3.16275952  10.01111482  31.82001484 103.7113329  358.16259615\n",
      "          nan          nan          nan          nan          nan\n",
      "          nan          nan]\n",
      "vgg\n",
      "grasp\n",
      "[90.89666667 90.18333333 89.20333333 88.24333333 85.56333333 10.\n",
      " 10.         10.         10.         10.         10.         10.        ]\n",
      "[  3.16275952  10.01111482  31.82001484 103.7113329  358.16259615\n",
      "          nan          nan          nan          nan          nan\n",
      "          nan          nan]\n",
      "vgg\n",
      "mica_sf_rand_effective_comp\n",
      "[92.54333333 92.01666667 91.13333333 88.74333333 84.96       79.76333333\n",
      " 72.36       61.44333333 46.2        26.63       10.01       10.49666667]\n",
      "[3.16235210e+00 1.00010561e+01 3.16305684e+01 1.00065173e+02\n",
      " 3.16657136e+02 1.00190095e+03 3.17033823e+03 1.00562213e+04\n",
      " 3.20373682e+04 1.00569359e+05 3.17635172e+05            nan]\n",
      "vgg\n",
      "sf\n",
      "[92.33       92.08       91.02       88.90333333 85.97333333 80.31666667\n",
      " 70.85333333 59.82333333 45.07666667 29.99333333 20.27333333 10.        ]\n",
      "[3.16235210e+00 1.00010561e+01 3.16305684e+01 1.00065173e+02\n",
      " 3.16657136e+02 1.00190095e+03 3.17033823e+03 1.00562213e+04\n",
      " 3.20373682e+04 1.00569359e+05 3.17635172e+05            nan]\n",
      "resnet\n",
      "mica_cs_rand_effective_comp\n",
      "[87.65333333 85.02       77.85333333 65.51666667 46.07333333 24.72666667\n",
      " 14.59666667 10.        ]\n",
      "[3.17093326e+00 1.04405459e+01 3.68259797e+01 1.43291052e+02\n",
      " 6.18096664e+02 6.65719342e+03            nan            nan]\n",
      "resnet\n",
      "cs\n",
      "[86.83666667 83.42666667 75.84       62.12333333 40.84333333 22.80333333\n",
      " 14.25666667 10.        ]\n",
      "[3.17093326e+00 1.04405459e+01 3.68259797e+01 1.43291052e+02\n",
      " 6.18096664e+02 6.65719342e+03            nan            nan]\n",
      "resnet\n",
      "mica_grasp_rand_effective_comp\n",
      "[87.87       85.         79.28666667 67.07666667 51.20333333 26.93\n",
      " 13.01333333 10.        ]\n",
      "[   3.16473561   10.21027478   34.61947354  123.21172478  467.24132287\n",
      " 1798.05203399           nan           nan]\n",
      "resnet\n",
      "grasp\n",
      "[84.63333333 82.57333333 75.03666667 63.65333333 42.95666667 23.67\n",
      " 11.98666667 10.        ]\n",
      "[   3.16473561   10.21027478   34.61947354  123.21172478  467.24132287\n",
      " 1798.05203399           nan           nan]\n",
      "resnet\n",
      "mica_sf_rand_effective_comp\n",
      "[88.28333333 85.22       80.1        69.25       55.06       38.08333333\n",
      " 25.82333333 19.86333333]\n",
      "[3.16282545e+00 1.00036687e+01 3.16553595e+01 1.00121877e+02\n",
      " 3.16837714e+02 1.01465849e+03 3.18730581e+03 9.67485702e+03]\n",
      "resnet\n",
      "sf\n",
      "[87.68666667 84.28333333 77.3        66.85333333 52.16333333 33.25333333\n",
      " 22.39666667 17.85333333]\n",
      "[3.16282545e+00 1.00036687e+01 3.16553595e+01 1.00121877e+02\n",
      " 3.16837714e+02 1.01465849e+03 3.18730581e+03 9.67485702e+03]\n",
      "0, Dense\n",
      "1,  \n",
      "2, $\\bf{MiCA}$-$\\bf{SNIP}$\n",
      "3, RbPI-SNIP\n",
      "4, $\\bf{MiCA}$-$\\bf{GraSP}$\n",
      "5, RbPI-GraSP\n",
      "6, $\\bf{MiCA}$-$\\bf{SynFlow}$\n",
      "7, RbPI-SynFlow\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 136792 characters]",
      "text/plain": [
       "<Figure size 1066.18x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_pruners = [\n",
    "    'mica_cs_rand_effective_comp',\n",
    "    'cs',\n",
    "    'mica_grasp_rand_effective_comp',\n",
    "    'grasp',\n",
    "    'mica_sf_rand_effective_comp',\n",
    "    'sf'\n",
    "    ]\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(8*1.851, 8))\n",
    "\n",
    "for i, arch in enumerate(['vgg', 'resnet']):\n",
    "# for i, arch in enumerate(['vgg']):\n",
    "    if arch == 'vgg':\n",
    "        baseline = vgg_baseline\n",
    "        effective_x_ratios = vgg_effective_x_ratios\n",
    "        accs = vgg_accs\n",
    "        xticks = [10**(2*i - 1) for i in range(1,4)]\n",
    "    elif arch == 'resnet':\n",
    "        baseline = resnet_baseline\n",
    "        effective_x_ratios = resnet_effective_x_ratios\n",
    "        accs = resnet_accs\n",
    "        xticks = [10**(2*i - 1) for i in range(1,6)]\n",
    "    ax = axes[i]\n",
    "    ax.axhline(baseline, color = 'black', linewidth = lw, linestyle='dotted', label='Dense', zorder=0)\n",
    "    ax.scatter(0, 0, marker=\"\", label=' ')\n",
    "    \n",
    "    zorder = 1\n",
    "    for l, pruner in enumerate(plot_pruners):\n",
    "        print(arch)\n",
    "        print(pruner)\n",
    "        print(accs[pruner])\n",
    "        print(effective_x_ratios[pruner])\n",
    "        if 'mica' in pruner:\n",
    "            linestyle = 'solid'\n",
    "            ax.plot(effective_x_ratios[pruner], accs[pruner], color=colors[pruner], \n",
    "                    linestyle=linestyle, lw=lw,\n",
    "                    marker=markers[pruner], markersize=markersize,  \n",
    "                    label=pruner_names[pruner],\n",
    "                    zorder=zorder+10)\n",
    "            zorder += 1\n",
    "        else:\n",
    "            linestyle = 'dashdot'\n",
    "            ax.plot(effective_x_ratios[pruner], accs[pruner], color=colors[pruner], \n",
    "                    linestyle=linestyle, lw=lw,\n",
    "                    marker=markers[pruner], markersize=markersize,  \n",
    "                    markerfacecolor= 'white', markeredgewidth=2,\n",
    "                    label=pruner_names[pruner],\n",
    "                    zorder=-zorder)\n",
    "            zorder += 1\n",
    "\n",
    "    ax.set_ylim(7, 96)\n",
    "    ax.set_xscale('log')\n",
    "    ax.set_xticks(xticks)\n",
    "    ax.set_yticks(np.arange(10, 100, 20))\n",
    "    ax.minorticks_off()\n",
    "    \n",
    "    ax.tick_params(\n",
    "        which       = \"both\", bottom    = True,  top   = False,\n",
    "        labelbottom = True,   left      = True,  right = False,\n",
    "        labelleft   = False,   \n",
    "        direction   = 'out',  length    = 7,     width = 1.0,    \n",
    "        pad         = 8,      labelsize = fontsize\n",
    "        )\n",
    "    ax.set_title('VGG-16' if arch == 'vgg' else 'ResNet-20', size=fontsize)\n",
    "    ax.set_xlabel('Compression Ratio', size=fontsize)\n",
    "    if i == 0:\n",
    "        ax.set_xlim(10**0.1, 10**6.5)\n",
    "    else:\n",
    "        ax.set_xlim(10**0.1, 10**4.75)\n",
    "    ax.grid(True)\n",
    "    ax.set_axisbelow(True)\n",
    "    \n",
    "\n",
    "    # boundary and zoom box\n",
    "    for dir in [\"top\", \"bottom\", \"right\", \"left\"]:\n",
    "        ax.spines[dir].set_linewidth(3)\n",
    "    ax.spines[\"right\"].set_linestyle((0,(3,3)))\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['top'].set_visible(False)\n",
    "\n",
    "fig.subplots_adjust(hspace=0.2)\n",
    "fig.tight_layout()\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "for nl, label in enumerate(labels):\n",
    "    print(f'{nl}, {label}')\n",
    "order = [2,4,6,3,5,7,0,1]\n",
    "ordered_handles = get_ordered_list(handles, order)\n",
    "ordered_labels = get_ordered_list(labels, order)\n",
    "plt.legend(\n",
    "    handles=ordered_handles,labels=ordered_labels, \n",
    "    fontsize=fontsize, loc='upper center', \n",
    "    bbox_to_anchor=(-.03, -.25), ncol=3)\n",
    "os.makedirs('Plots', exist_ok=True)\n",
    "plt.savefig('./Plots/cifar10_mica_rank_sparsity_distribution.eps', bbox_inches=\"tight\")\n",
    "plt.savefig('./Plots/cifar10_mica_rank_sparsity_distribution.svg', bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input Nodes Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expid  = 'cifar10_lottery_vgg16_bn'\n",
    "ratios = [\n",
    "    '0.5', '1.0', '1.5', '2.0', \n",
    "    '2.5', '3.0', '3.5', '4.0',\n",
    "    '4.5', '5.0', '5.5', '6.0'\n",
    "    ]\n",
    "\n",
    "# pruners = np.array(\n",
    "#     [\n",
    "#         'mica_erk_rand_effective_comp',\n",
    "#         'mica_igq_rand_effective_comp',\n",
    "#         'mica_erk_shuffle',\n",
    "#         'mica_igq_shuffle'\n",
    "#     ],\n",
    "#     dtype=object\n",
    "#     )\n",
    "\n",
    "for pruner in pruners:\n",
    "    if '_shuffle' in pruner:\n",
    "        colors[pruner]  = colors[pruner.replace('_shuffle', '')]\n",
    "        markers[pruner]  = markers[pruner.replace('_shuffle', '')]\n",
    "        pepochs[pruner] = pepochs[pruner.replace('_shuffle', '')]\n",
    "    elif '_effective_comp' in pruner:\n",
    "        colors[pruner]  = colors[pruner.replace('_effective_comp', '')]\n",
    "        markers[pruner]  = markers[pruner.replace('_effective_comp', '')]\n",
    "        pepochs[pruner] = pepochs[pruner.replace('_effective_comp', '')]\n",
    "    elif '_rand' in pruner:\n",
    "        colors[pruner]  = colors[pruner.replace('_rand', '')]\n",
    "        markers[pruner]  = markers[pruner.replace('_rand', '')]\n",
    "        pepochs[pruner] = pepochs[pruner.replace('_rand', '')]\n",
    "\n",
    "vgg_pruner_n_losses    = {}\n",
    "vgg_pruner_n_zeros     = {}\n",
    "vgg_pruner_n_params    = {}\n",
    "vgg_pruner_n_mask_V_in = {}\n",
    "vgg_effective_x_ratios = {}\n",
    "\n",
    "for l, pruner in enumerate(pruners):\n",
    "    vgg_pruner_n_losses[pruner]    = []\n",
    "    vgg_pruner_n_zeros[pruner]     = []\n",
    "    vgg_pruner_n_params[pruner]    = []\n",
    "    vgg_pruner_n_mask_V_in[pruner] = []\n",
    "    vgg_effective_x_ratios[pruner] = []\n",
    "\n",
    "for ratio in ratios:\n",
    "    for l, pruner in enumerate(pruners):\n",
    "        effective_x_ratios = []\n",
    "        n_losses_arr      = []\n",
    "        n_zeros_arr       = []\n",
    "        n_params_arr      = []\n",
    "        n_mask_V_in_arr   = []\n",
    "        n_errors = 0\n",
    "        for r, run_num in enumerate(runs):\n",
    "            experimentid      = get_experimentid(pruner, ratio, pepochs[pruner], run_num, expid)\n",
    "            model_file        = experimentid + 'model.pt'\n",
    "            compression_file  = experimentid + 'compression.pkl'\n",
    "            try:\n",
    "                compression       = pd.read_pickle(compression_file)\n",
    "            except:\n",
    "                n_errors += 1\n",
    "                continue\n",
    "            prunable          = compression['prunable']==True\n",
    "            weight            = compression['param']=='weight'\n",
    "            modules           = compression[np.logical_and(prunable, weight)]['module']\n",
    "            shape             = compression[np.logical_and(prunable,weight)]['shape']\n",
    "\n",
    "            with open(experimentid + 'args.json') as f:\n",
    "                args = json.load(f)\n",
    "            input_shape, num_classes = load.dimension(args['dataset']) \n",
    "            loss = nn.CrossEntropyLoss()\n",
    "            if l == 0 and r == 0:\n",
    "                test_loader, _, _ = load.dataloader(\n",
    "                    args['dataset'], args['test_batch_size'], args['test_batch_size'],\n",
    "                    False, args['workers'], 0, 0)\n",
    "                dataset = args['dataset']\n",
    "                test_batch_size = args['test_batch_size']\n",
    "                workers = args['workers']\n",
    "                (data, _) = next(iter(test_loader))\n",
    "                input_dim = list(data[0,:].shape)\n",
    "                input = torch.ones([1] + input_dim).to('cuda')\n",
    "            assert args['seed'] == run_num * 1000 # adhoc\n",
    "            assert args['run_number'] == str(run_num)\n",
    "            assert dataset == args['dataset']\n",
    "            assert test_batch_size == args['test_batch_size']\n",
    "            assert workers == args['workers']\n",
    "\n",
    "            model = load.model(args['model'], args['model_class'])(\n",
    "                input_shape, num_classes, args['dense_classifier'], args['pretrained']).to('cuda')\n",
    "            model.load_state_dict(torch.load(model_file))\n",
    "            print(f'Forward :')\n",
    "            model.set_effective_sparsity(is_forward=True, x=input)\n",
    "            print(f'Backward :')\n",
    "            model.set_effective_sparsity(is_forward=False, x=None)\n",
    "\n",
    "            effective_densities = []\n",
    "            set_modules = set(modules.unique())\n",
    "            print(f'Modules :')\n",
    "            n_mask_V_in = 0\n",
    "            for n, m in model.named_modules():\n",
    "                if hasattr(m, 'weight_mask') and n in set_modules:\n",
    "                    print(n)\n",
    "                    effective_densities.append((m.weight_mask.sum() / m.weight_mask.numel()).to('cpu'))\n",
    "                    reshaped_weight_mask = m.weight_mask.reshape(m.weight_mask.size()[0], -1)\n",
    "                    comp_mask            = reshaped_weight_mask[torch.any(reshaped_weight_mask != 0, dim=1),:]\n",
    "                    comp_mask            = comp_mask[:, torch.any(comp_mask != 0, dim=0)]\n",
    "                    assert comp_mask.dim() == 2\n",
    "                    n_mask_V_in += (1.0 * (comp_mask.sum(dim=0) != 0)).sum().item()\n",
    "            n_mask_V_in_arr.append(n_mask_V_in)\n",
    "\n",
    "            effective_densities = np.array(effective_densities)\n",
    "            e = np.array([np.prod(s) for s in shape]) * (effective_densities)\n",
    "            if np.all(e == 0):                \n",
    "                actual_comp = float('nan')\n",
    "            else:\n",
    "                actual_comp = np.sum(np.array([np.prod(s) for s in shape])) / np.sum(e)\n",
    "            effective_x_ratios.append(actual_comp)\n",
    "\n",
    "            n_losses, n_zeros, n_params = model.check_V_in_loss(x=input, get_sparsities=True)\n",
    "            n_losses_arr.append(sum(n_losses))\n",
    "            if isinstance(n_zeros, torch.Tensor):\n",
    "                n_zeros = n_zeros.to('cpu').numpy()\n",
    "            if isinstance(n_params, torch.Tensor):\n",
    "                n_params = n_params.to('cpu').numpy()\n",
    "            n_zeros_arr.append(n_zeros)\n",
    "            n_params_arr.append(n_params)\n",
    "            \n",
    "        if n_errors == len(runs):\n",
    "            pass\n",
    "        elif n_errors == 0:\n",
    "            vgg_pruner_n_mask_V_in[pruner].append(np.mean(n_mask_V_in_arr))\n",
    "            vgg_pruner_n_losses[pruner].append(np.mean(n_losses_arr))\n",
    "            vgg_pruner_n_zeros[pruner].append(np.mean(n_zeros_arr))\n",
    "            vgg_pruner_n_params[pruner].append(np.mean(n_params_arr))\n",
    "            vgg_effective_x_ratios[pruner].append(np.mean(effective_x_ratios)) \n",
    "        else:\n",
    "            raise ValueError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expid  = 'cifar10_lottery_resnet20'\n",
    "ratios = [\n",
    "    '0.5', '1.0', '1.5', '2.0', \n",
    "    '2.5', '3.0', '3.5', '4.0'\n",
    "    ]\n",
    "\n",
    "# pruners = np.array(\n",
    "#     [\n",
    "#         'mica_erk_rand_effective_comp',\n",
    "#         'mica_igq_rand_effective_comp',\n",
    "#         'mica_erk_shuffle',\n",
    "#         'mica_igq_shuffle'\n",
    "#     ],\n",
    "#     dtype=object\n",
    "#     )\n",
    "\n",
    "for pruner in pruners:\n",
    "    if '_shuffle' in pruner:\n",
    "        colors[pruner]  = colors[pruner.replace('_shuffle', '')]\n",
    "        markers[pruner]  = markers[pruner.replace('_shuffle', '')]\n",
    "        pepochs[pruner] = pepochs[pruner.replace('_shuffle', '')]\n",
    "    elif '_effective_comp' in pruner:\n",
    "        colors[pruner]  = colors[pruner.replace('_effective_comp', '')]\n",
    "        markers[pruner]  = markers[pruner.replace('_effective_comp', '')]\n",
    "        pepochs[pruner] = pepochs[pruner.replace('_effective_comp', '')]\n",
    "    elif '_rand' in pruner:\n",
    "        colors[pruner]  = colors[pruner.replace('_rand', '')]\n",
    "        markers[pruner]  = markers[pruner.replace('_rand', '')]\n",
    "        pepochs[pruner] = pepochs[pruner.replace('_rand', '')]\n",
    "\n",
    "resnet_pruner_n_losses    = {}\n",
    "resnet_pruner_n_zeros     = {}\n",
    "resnet_pruner_n_params    = {}\n",
    "resnet_pruner_n_mask_V_in = {}\n",
    "resnet_effective_x_ratios = {}\n",
    "\n",
    "for l, pruner in enumerate(pruners):\n",
    "    resnet_pruner_n_losses[pruner]    = []\n",
    "    resnet_pruner_n_zeros[pruner]     = []\n",
    "    resnet_pruner_n_params[pruner]    = []\n",
    "    resnet_pruner_n_mask_V_in[pruner] = []\n",
    "    resnet_effective_x_ratios[pruner] = []\n",
    "\n",
    "for ratio in ratios:\n",
    "    for l, pruner in enumerate(pruners):\n",
    "        effective_x_ratios = []\n",
    "        n_losses_arr      = []\n",
    "        n_zeros_arr       = []\n",
    "        n_params_arr      = []\n",
    "        n_mask_V_in_arr   = []\n",
    "        n_errors = 0\n",
    "        for r, run_num in enumerate(runs):\n",
    "            experimentid      = get_experimentid(pruner, ratio, pepochs[pruner], run_num, expid)\n",
    "            model_file        = experimentid + 'model.pt'\n",
    "            compression_file  = experimentid + 'compression.pkl'\n",
    "            try:\n",
    "                compression       = pd.read_pickle(compression_file)\n",
    "            except:\n",
    "                n_errors += 1\n",
    "                continue\n",
    "            prunable          = compression['prunable']==True\n",
    "            weight            = compression['param']=='weight'\n",
    "            modules           = compression[np.logical_and(prunable, weight)]['module']\n",
    "            shape             = compression[np.logical_and(prunable,weight)]['shape']\n",
    "\n",
    "            with open(experimentid + 'args.json') as f:\n",
    "                args = json.load(f)\n",
    "            input_shape, num_classes = load.dimension(args['dataset']) \n",
    "            loss = nn.CrossEntropyLoss()\n",
    "            if l == 0 and r == 0:\n",
    "                test_loader, _, _ = load.dataloader(\n",
    "                    args['dataset'], args['test_batch_size'], args['test_batch_size'],\n",
    "                    False, args['workers'], 0, 0)\n",
    "                dataset = args['dataset']\n",
    "                test_batch_size = args['test_batch_size']\n",
    "                workers = args['workers']\n",
    "                (data, _) = next(iter(test_loader))\n",
    "                input_dim = list(data[0,:].shape)\n",
    "                input = torch.ones([1] + input_dim).to('cuda')\n",
    "            assert args['seed'] == run_num * 1000 # adhoc\n",
    "            assert args['run_number'] == str(run_num)\n",
    "            assert dataset == args['dataset']\n",
    "            assert test_batch_size == args['test_batch_size']\n",
    "            assert workers == args['workers']\n",
    "\n",
    "            model = load.model(args['model'], args['model_class'])(\n",
    "                input_shape, num_classes, args['dense_classifier'], args['pretrained']).to('cuda')\n",
    "            model.load_state_dict(torch.load(model_file))\n",
    "            print(f'Forward :')\n",
    "            model.set_effective_sparsity(is_forward=True, x=input)\n",
    "            print(f'Backward :')\n",
    "            model.set_effective_sparsity(is_forward=False, x=None)\n",
    "\n",
    "            effective_densities = []\n",
    "            set_modules = set(modules.unique())\n",
    "            print(f'Modules :')\n",
    "            n_mask_V_in = 0\n",
    "            for n, m in model.named_modules():\n",
    "                if hasattr(m, 'weight_mask') and n in set_modules:\n",
    "                    print(n)\n",
    "                    effective_densities.append((m.weight_mask.sum() / m.weight_mask.numel()).to('cpu'))\n",
    "                    reshaped_weight_mask = m.weight_mask.reshape(m.weight_mask.size()[0], -1)\n",
    "                    comp_mask            = reshaped_weight_mask[torch.any(reshaped_weight_mask != 0, dim=1),:]\n",
    "                    comp_mask            = comp_mask[:, torch.any(comp_mask != 0, dim=0)]\n",
    "                    assert comp_mask.dim() == 2\n",
    "                    n_mask_V_in += (1.0 * (comp_mask.sum(dim=0) != 0)).sum().item()\n",
    "            n_mask_V_in_arr.append(n_mask_V_in)\n",
    "\n",
    "            effective_densities = np.array(effective_densities)\n",
    "            e = np.array([np.prod(s) for s in shape]) * (effective_densities)\n",
    "            if np.all(e == 0):                \n",
    "                actual_comp = float('nan')\n",
    "            else:\n",
    "                actual_comp = np.sum(np.array([np.prod(s) for s in shape])) / np.sum(e)\n",
    "            effective_x_ratios.append(actual_comp)\n",
    "\n",
    "            n_losses, n_zeros, n_params = model.check_V_in_loss(x=input, get_sparsities=True)\n",
    "            n_losses_arr.append(sum(n_losses))\n",
    "            if isinstance(n_zeros, torch.Tensor):\n",
    "                n_zeros = n_zeros.to('cpu').numpy()\n",
    "            if isinstance(n_params, torch.Tensor):\n",
    "                n_params = n_params.to('cpu').numpy()\n",
    "            n_zeros_arr.append(n_zeros)\n",
    "            n_params_arr.append(n_params)\n",
    "            \n",
    "        if n_errors == len(runs):\n",
    "            pass\n",
    "        elif n_errors == 0:\n",
    "            resnet_pruner_n_mask_V_in[pruner].append(np.mean(n_mask_V_in_arr))\n",
    "            resnet_pruner_n_losses[pruner].append(np.mean(n_losses_arr))\n",
    "            resnet_pruner_n_zeros[pruner].append(np.mean(n_zeros_arr))\n",
    "            resnet_pruner_n_params[pruner].append(np.mean(n_params_arr))\n",
    "            resnet_effective_x_ratios[pruner].append(np.mean(effective_x_ratios)) \n",
    "        else:\n",
    "            raise ValueError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The PostScript backend does not support transparency; partially transparent artists will be rendered opaque.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mica_erk_rand_effective_comp\n",
      "[3.1622799470525984, 10.000009371053153, 31.622954468524068, 99.99853342637431, 316.2329551336324, 1105.7950258688973, 47828.003793993616, nan, nan, nan, nan, nan]\n",
      "[92.46       91.78       89.47333333 87.00666667 82.58666667 76.84333333\n",
      " 39.30333333 10.         10.         10.         10.         10.        ]\n",
      "mica_igq_rand_effective_comp\n",
      "[3.162281116060736, 10.000064033209846, 31.62315836356267, 100.00464803075143, 316.5436221146768, 1164.156603260417, 17984.396294835507]\n",
      "[92.37       91.94666667 90.91       88.68333333 85.40666667 79.89\n",
      " 22.33333333]\n",
      "mica_erk_rand_effective_comp\n",
      "[3.1621977719064023, 9.99948319610604, 31.64303256400557, 104.13743011063504, 565.4350398798734, nan, nan, nan]\n",
      "[88.11666667 85.31       79.46333333 67.13333333 45.32666667 10.\n",
      " 10.         10.        ]\n",
      "mica_igq_rand_effective_comp\n",
      "[3.1626407861829393, 10.002436953070191, 31.67634045642123, 105.59970602101411, 546.4472932536529, nan, nan, nan]\n",
      "[87.93666667 85.41       80.13       67.28666667 48.04       12.06666667\n",
      " 10.         10.        ]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 131648 characters]",
      "text/plain": [
       "<Figure size 1152x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_pruners = [\n",
    "    'mica_erk_rand_effective_comp',\n",
    "    'mica_erk_shuffle',\n",
    "    'mica_igq_rand_effective_comp',\n",
    "    'mica_igq_shuffle',\n",
    "    ]\n",
    "\n",
    "lw=4\n",
    "markersize=30\n",
    "fontsize=35\n",
    "\n",
    "plt.rcParams[\"font.family\"] = \"Segoe UI\"\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(8*2, 8))\n",
    "\n",
    "for i, arch in enumerate(['vgg', 'resnet']):\n",
    "\n",
    "    if arch == 'vgg':\n",
    "        effective_x_ratios = vgg_effective_x_ratios\n",
    "        pruner_n_mask_V_in = vgg_pruner_n_mask_V_in\n",
    "        pruner_n_losses    = vgg_pruner_n_losses\n",
    "        accs               = vgg_accs\n",
    "        xticks = [10**(2*i - 1) for i in range(1,4)]\n",
    "    elif arch == 'resnet':\n",
    "        effective_x_ratios = resnet_effective_x_ratios\n",
    "        pruner_n_mask_V_in = resnet_pruner_n_mask_V_in\n",
    "        pruner_n_losses    = resnet_pruner_n_losses\n",
    "        accs               = resnet_accs\n",
    "        xticks = [10**(2*i - 1) for i in range(1,6)]\n",
    "    ax = axes[i]\n",
    "\n",
    "    zorder = 1\n",
    "    for l, pruner in enumerate(plot_pruners):\n",
    "        if '_shuffle' in pruner:\n",
    "            linestyle = 'dashed'  \n",
    "            ax.plot(effective_x_ratios[pruner], \n",
    "                    1 - np.array(pruner_n_losses[pruner]) / (np.array(pruner_n_mask_V_in[pruner])+np.array(pruner_n_losses[pruner])), \n",
    "                    color=colors[pruner], \n",
    "                    linestyle=linestyle, lw=lw,\n",
    "                    marker=markers[pruner], markersize=markersize, markerfacecolor= 'white', markeredgewidth=2,\n",
    "                    label=pruner_names[pruner],\n",
    "                    zorder=0)\n",
    "            zorder += 1\n",
    "        elif 'mica' in pruner:\n",
    "            linestyle = 'solid'\n",
    "            print(pruner)\n",
    "            print(effective_x_ratios[pruner])\n",
    "            print(accs[pruner])\n",
    "            ax.plot(effective_x_ratios[pruner], \n",
    "                    1 - np.array(pruner_n_losses[pruner]) / (np.array(pruner_n_mask_V_in[pruner])+np.array(pruner_n_losses[pruner])), \n",
    "                    color=colors[pruner], \n",
    "                    linestyle=linestyle, lw=lw,\n",
    "                    marker=markers[pruner], markersize=markersize,  \n",
    "                    label=pruner_names[pruner],\n",
    "                    zorder=10)\n",
    "            zorder += 1\n",
    "\n",
    "    ax.set_xscale('log')\n",
    "    ax.set_xticks(xticks)\n",
    "    ax.set_ylim(-0.0375, 1.075)\n",
    "    ax.set_yticks([0, 0.25, 0.5, 0.75, 1.0])\n",
    "    ax.minorticks_off()\n",
    "\n",
    "    ax.tick_params(\n",
    "        which       = \"both\", bottom    = True,  top   = False,\n",
    "        labelbottom = True,   left      = True,  right = False,\n",
    "        labelleft   = True if arch == 'vgg' else False,   \n",
    "        direction   = 'out',  length    = 7,     width = 1.0,    \n",
    "        pad         = 8,      labelsize = fontsize\n",
    "        )\n",
    "\n",
    "    ax.set_title('VGG-16' if arch == 'vgg' else 'ResNet-20', size=fontsize)\n",
    "    ax.set_xlabel('Compression Ratio', size=fontsize)\n",
    "    if i == 0:\n",
    "        ax.set_ylabel('Ratio of Functional Nodes', size=fontsize)\n",
    "        ax.set_xlim(10**0.1, 10**6.5)\n",
    "    else:\n",
    "        ax.set_xlim(10**0.1, 10**4.75)\n",
    "    ax.grid(True)\n",
    "    ax.set_axisbelow(True)\n",
    "    \n",
    "\n",
    "    # boundary and zoom box\n",
    "    for dir in [\"top\", \"bottom\", \"right\", \"left\"]:\n",
    "        ax.spines[dir].set_linewidth(3)\n",
    "    ax.spines[\"right\"].set_linestyle((0,(3,3)))\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['top'].set_visible(False)\n",
    "\n",
    "fig.subplots_adjust(hspace=0.2)\n",
    "fig.tight_layout()\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "# for nl, label in enumerate(labels):\n",
    "#     print(f'{nl}, {label}')\n",
    "order = [0, 2, 1, 3]\n",
    "ordered_handles = get_ordered_list(handles, order)\n",
    "ordered_labels = get_ordered_list(labels, order)\n",
    "plt.legend(\n",
    "    handles=ordered_handles,labels=ordered_labels, \n",
    "    fontsize=fontsize, loc='upper center', \n",
    "    bbox_to_anchor=(-.1, -.25), ncol=2)\n",
    "os.makedirs('Plots', exist_ok=True)\n",
    "plt.savefig('./Plots/cifar10_mica_rand_functional_nodes.eps', bbox_inches=\"tight\")\n",
    "plt.savefig('./Plots/cifar10_mica_rand_functional_nodes.svg', bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The PostScript backend does not support transparency; partially transparent artists will be rendered opaque.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mica_cs_rand_effective_comp\n",
      "[3.1700386515201, 10.095267108879666, 32.87195132279081, 115.54719817470324, 456.6758179628107, 2121.417617907245, nan, nan, nan, nan, nan, nan]\n",
      "[92.28666667 91.33333333 89.84       87.78       83.94666667 73.29333333\n",
      " 10.         10.         10.         10.         10.         10.        ]\n",
      "cs\n",
      "[3.1700386515201, 10.095267108879666, 32.87195132279081, 115.54719817470324, 456.6758179628107, 2121.417617907245, nan, nan, nan, nan, nan, nan]\n",
      "[92.26       91.28666667 90.04666667 87.89666667 83.32666667 69.86666667\n",
      " 10.         10.         10.         10.         10.         10.        ]\n",
      "mica_grasp_rand_effective_comp\n",
      "[3.16275951540686, 10.011114824379083, 31.820014839878667, 103.71133290442737, 358.162596148767, nan, nan, nan, nan, nan, nan, nan]\n",
      "[92.22333333 91.63333333 90.50333333 88.90666667 85.91666667 10.\n",
      " 10.         10.         10.         10.         10.         10.        ]\n",
      "grasp\n",
      "[3.16275951540686, 10.011114824379083, 31.820014839878667, 103.71133290442737, 358.162596148767, nan, nan, nan, nan, nan, nan, nan]\n",
      "[90.89666667 90.18333333 89.20333333 88.24333333 85.56333333 10.\n",
      " 10.         10.         10.         10.         10.         10.        ]\n",
      "mica_sf_rand_effective_comp\n",
      "[3.162352096810315, 10.001056066885516, 31.630568374025547, 100.06517347374415, 316.6571356095927, 1001.9009468412396, 3170.3382286001724, 10056.221293233904, 32037.368156818775, 100569.35910511257, 317635.1718609704, nan]\n",
      "[92.54333333 92.01666667 91.13333333 88.74333333 84.96       79.76333333\n",
      " 72.36       61.44333333 46.2        26.63       10.01       10.49666667]\n",
      "sf\n",
      "[3.162352096810315, 10.001056066885516, 31.630568374025547, 100.06517347374415, 316.6571356095927, 1001.9009468412396, 3170.3382286001724, 10056.221293233904, 32037.368156818775, 100569.35910511257, 317635.1718609704, nan]\n",
      "[92.33       92.08       91.02       88.90333333 85.97333333 80.31666667\n",
      " 70.85333333 59.82333333 45.07666667 29.99333333 20.27333333 10.        ]\n",
      "mica_cs_rand_effective_comp\n",
      "[3.1709332619129804, 10.440545928128016, 36.825979748590214, 143.2910521377031, 618.096664452679, 6657.1934194004725, nan, nan]\n",
      "[87.65333333 85.02       77.85333333 65.51666667 46.07333333 24.72666667\n",
      " 14.59666667 10.        ]\n",
      "cs\n",
      "[3.1709332619129804, 10.440545928128016, 36.825979748590214, 143.2910521377031, 618.096664452679, 6657.1934194004725, nan, nan]\n",
      "[86.83666667 83.42666667 75.84       62.12333333 40.84333333 22.80333333\n",
      " 14.25666667 10.        ]\n",
      "mica_grasp_rand_effective_comp\n",
      "[3.1647356110242684, 10.210274776741366, 34.61947353622923, 123.2117247751687, 467.2413228706337, 1798.0520339940795, nan, nan]\n",
      "[87.87       85.         79.28666667 67.07666667 51.20333333 26.93\n",
      " 13.01333333 10.        ]\n",
      "grasp\n",
      "[3.1647356110242684, 10.210274776741366, 34.61947353622923, 123.2117247751687, 467.2413228706337, 1798.0520339940795, nan, nan]\n",
      "[84.63333333 82.57333333 75.03666667 63.65333333 42.95666667 23.67\n",
      " 11.98666667 10.        ]\n",
      "mica_sf_rand_effective_comp\n",
      "[3.162825453933992, 10.003668664621488, 31.655359519905456, 100.12187693025108, 316.83771351777125, 1014.6584943366676, 3187.305813817906, 9674.857018427632]\n",
      "[88.28333333 85.22       80.1        69.25       55.06       38.08333333\n",
      " 25.82333333 19.86333333]\n",
      "sf\n",
      "[3.162825453933992, 10.003668664621488, 31.655359519905456, 100.12187693025108, 316.83771351777125, 1014.6584943366676, 3187.305813817906, 9674.857018427632]\n",
      "[87.68666667 84.28333333 77.3        66.85333333 52.16333333 33.25333333\n",
      " 22.39666667 17.85333333]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 161996 characters]",
      "text/plain": [
       "<Figure size 1066.18x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_pruners = [\n",
    "    'mica_cs_rand_effective_comp',\n",
    "    'cs',\n",
    "    'mica_grasp_rand_effective_comp',\n",
    "    'grasp',\n",
    "    'mica_sf_rand_effective_comp',\n",
    "    'sf'\n",
    "    ]\n",
    "\n",
    "lw=4\n",
    "markersize=30\n",
    "fontsize=35\n",
    "\n",
    "plt.rcParams[\"font.family\"] = \"Segoe UI\"\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(8*1.851, 8))\n",
    "\n",
    "for i, arch in enumerate(['vgg', 'resnet']):\n",
    "\n",
    "    if arch == 'vgg':\n",
    "        effective_x_ratios = vgg_effective_x_ratios\n",
    "        pruner_n_mask_V_in = vgg_pruner_n_mask_V_in\n",
    "        pruner_n_losses    = vgg_pruner_n_losses\n",
    "        accs               = vgg_accs\n",
    "        xticks = [10**(2*i - 1) for i in range(1,4)]\n",
    "    elif arch == 'resnet':\n",
    "        effective_x_ratios = resnet_effective_x_ratios\n",
    "        pruner_n_mask_V_in = resnet_pruner_n_mask_V_in\n",
    "        pruner_n_losses    = resnet_pruner_n_losses\n",
    "        accs               = resnet_accs\n",
    "        xticks = [10**(2*i - 1) for i in range(1,4)]\n",
    "    ax = axes[i]\n",
    "\n",
    "    zorder = 1\n",
    "    for l, pruner in enumerate(plot_pruners):\n",
    "        if 'mica' in pruner:\n",
    "            linestyle = 'solid'\n",
    "            print(pruner)\n",
    "            print(effective_x_ratios[pruner])\n",
    "            print(accs[pruner])\n",
    "            ax.plot(effective_x_ratios[pruner], \n",
    "                    1 - np.array(pruner_n_losses[pruner]) / (np.array(pruner_n_mask_V_in[pruner])+np.array(pruner_n_losses[pruner])), \n",
    "                    color=colors[pruner], \n",
    "                    linestyle=linestyle, lw=lw,\n",
    "                    marker=markers[pruner], markersize=markersize,  \n",
    "                    label=pruner_names[pruner],\n",
    "                    zorder=10)\n",
    "            zorder += 1\n",
    "        else:\n",
    "            linestyle = 'dashdot'\n",
    "            print(pruner)\n",
    "            print(effective_x_ratios[pruner])\n",
    "            print(accs[pruner])\n",
    "            ax.plot(effective_x_ratios[pruner], \n",
    "                    1 - np.array(pruner_n_losses[pruner]) / (np.array(pruner_n_mask_V_in[pruner])+np.array(pruner_n_losses[pruner])), \n",
    "                    color=colors[pruner], \n",
    "                    linestyle=linestyle, lw=lw,\n",
    "                    marker=markers[pruner], markersize=markersize,  \n",
    "                    markerfacecolor= 'white', markeredgewidth=2,\n",
    "                    label=pruner_names[pruner],\n",
    "                    zorder=-10)\n",
    "            zorder += 1\n",
    "\n",
    "    ax.set_xscale('log')\n",
    "    ax.set_xticks(xticks)\n",
    "    ax.set_ylim(-0.0375, 1.075)\n",
    "    ax.set_yticks([0, 0.25, 0.5, 0.75, 1.0])\n",
    "    ax.minorticks_off()\n",
    "\n",
    "    ax.tick_params(\n",
    "        which       = \"both\", bottom    = True,  top   = False,\n",
    "        labelbottom = True,   left      = True,  right = False,\n",
    "        labelleft   =  False,   \n",
    "        direction   = 'out',  length    = 7,     width = 1.0,    \n",
    "        pad         = 8,      labelsize = fontsize\n",
    "        )\n",
    "\n",
    "    ax.set_title('VGG-16' if arch == 'vgg' else 'ResNet-20', size=fontsize)\n",
    "    ax.set_xlabel('Compression Ratio', size=fontsize)\n",
    "    if i == 0:\n",
    "        ax.set_xlim(10**0.1, 10**6.5)\n",
    "    else:\n",
    "        ax.set_xlim(10**0.1, 10**4.75)\n",
    "    ax.grid(True)\n",
    "    ax.set_axisbelow(True)\n",
    "    \n",
    "\n",
    "    # boundary and zoom box\n",
    "    for dir in [\"top\", \"bottom\", \"right\", \"left\"]:\n",
    "        ax.spines[dir].set_linewidth(3)\n",
    "    ax.spines[\"right\"].set_linestyle((0,(3,3)))\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['top'].set_visible(False)\n",
    "\n",
    "fig.subplots_adjust(hspace=0.2)\n",
    "fig.tight_layout()\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "# for nl, label in enumerate(labels):\n",
    "#     print(f'{nl}, {label}')\n",
    "order = [0, 2, 4, 1, 3, 5]\n",
    "ordered_handles = get_ordered_list(handles, order)\n",
    "ordered_labels = get_ordered_list(labels, order)\n",
    "plt.legend(\n",
    "    handles=ordered_handles,labels=ordered_labels, \n",
    "    fontsize=fontsize, loc='upper center', \n",
    "    bbox_to_anchor=(-.1, -.25), ncol=2)\n",
    "os.makedirs('Plots', exist_ok=True)\n",
    "plt.savefig('./Plots/cifar10_mica_rank_functional_nodes.eps', bbox_inches=\"tight\")\n",
    "plt.savefig('./Plots/cifar10_mica_rank_functional_nodes.svg', bbox_inches=\"tight\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "connected_edge",
   "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.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
