{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6b6cbac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# notebook with basic loading tests for dataset classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7515de4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset_base\n",
    "# comparison: permutation of checkpoint -> vectorization -> slicing\n",
    "import logging\n",
    "logging.basicConfig(level=logging.WARNING)\n",
    "\n",
    "from pathlib import Path\n",
    "import json\n",
    "from shrp.datasets.dataset_base import ModelDatasetBase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5c87caab",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-01-31 18:03:53,510\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading checkpoints from [PosixPath('/netscratch2/kschuerholt/code/versai/model_zoos/zoos/CIFAR10/resnet19/kaiming_uniform/tune_zoo_cifar10_resnet18_kaiming_uniform')]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████| 14/14 [00:07<00:00,  1.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data loaded. found 14 usable samples out of potential 14 samples.\n",
      "Load properties for samples from paths.\n",
      "### load data for dict_keys(['test_acc', 'training_iteration', 'ggap'])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "0it [00:00, ?it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training iteration 4 and epoch 5 don't match.\n",
      "### load data for dict_keys(['test_acc', 'training_iteration', 'ggap'])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "14it [00:00, 37.23it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Properties loaded.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "config_key_list = []\n",
    "result_key_list = [\n",
    "    \"test_acc\",\n",
    "    \"training_iteration\",\n",
    "    \"ggap\",\n",
    "#     \"sparsity_ratio\",\n",
    "]\n",
    "property_keys = {\n",
    "    \"result_keys\": result_key_list,\n",
    "    \"config_keys\": config_key_list,\n",
    "}\n",
    "\n",
    "# path_root = Path('/netscratch2/dtaskiran/zoos/MNIST/tune_zoo_mnist_uniform/')\n",
    "path_root = Path('/netscratch2/kschuerholt/code/versai/model_zoos/zoos/CIFAR10/resnet19/kaiming_uniform/tune_zoo_cifar10_resnet18_kaiming_uniform')\n",
    "\n",
    "dataset_test = ModelDatasetBase(\n",
    "    root=path_root,\n",
    "    epoch_lst=[5,10],\n",
    "    mode=\"checkpoint\",\n",
    "    train_val_test=\"train\",  # determines whcih dataset split to use\n",
    "    ds_split=[0.7, 0.15,0.15],  #\n",
    "    max_samples=10,\n",
    "    weight_threshold=float(\"inf\"),\n",
    "    filter_function=None,  # gets sample path as argument and returns True if model needs to be filtered out\n",
    "    property_keys=property_keys,\n",
    "    num_threads=4,\n",
    "    shuffle_path=True,\n",
    "    verbosity=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "922f7ac5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "14"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dataset_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2d94ba92",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████| 14/14 [00:11<00:00,  1.21it/s]\n"
     ]
    }
   ],
   "source": [
    "weights_per_channel = dataset_test.__get_weights_per_channel__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "edcab048",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "14"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(weights_per_channel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "71c88a19",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "18"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idx = 4\n",
    "len(weights_per_channel[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3638a147",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['layer_0-conv1.weight',\n",
       " 'layer_6-layer1.0.conv1.weight',\n",
       " 'layer_12-layer1.0.conv2.weight',\n",
       " 'layer_18-layer1.1.conv1.weight',\n",
       " 'layer_24-layer1.1.conv2.weight',\n",
       " 'layer_30-layer2.0.conv1.weight',\n",
       " 'layer_36-layer2.0.conv2.weight',\n",
       " 'layer_48-layer2.1.conv1.weight',\n",
       " 'layer_54-layer2.1.conv2.weight',\n",
       " 'layer_60-layer3.0.conv1.weight',\n",
       " 'layer_66-layer3.0.conv2.weight',\n",
       " 'layer_78-layer3.1.conv1.weight',\n",
       " 'layer_84-layer3.1.conv2.weight',\n",
       " 'layer_90-layer4.0.conv1.weight',\n",
       " 'layer_96-layer4.0.conv2.weight',\n",
       " 'layer_108-layer4.1.conv1.weight',\n",
       " 'layer_114-layer4.1.conv2.weight',\n",
       " 'layer_120-fc.weight']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "dataset_test.channel_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2756b51f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 27])\n",
      "torch.Size([64, 576])\n",
      "torch.Size([64, 576])\n",
      "torch.Size([64, 576])\n",
      "torch.Size([64, 576])\n",
      "torch.Size([128, 576])\n",
      "torch.Size([128, 1152])\n",
      "torch.Size([128, 1152])\n",
      "torch.Size([128, 1152])\n",
      "torch.Size([256, 1152])\n",
      "torch.Size([256, 2304])\n",
      "torch.Size([256, 2304])\n",
      "torch.Size([256, 2304])\n",
      "torch.Size([512, 2304])\n",
      "torch.Size([512, 4608])\n",
      "torch.Size([512, 4608])\n",
      "torch.Size([512, 4608])\n",
      "torch.Size([10, 513])\n",
      "10992330\n"
     ]
    }
   ],
   "source": [
    "count = 0\n",
    "for jdx in range(len(dataset_test.weights[idx])):\n",
    "    print(dataset_test.weights[idx][jdx].shape)\n",
    "    count += dataset_test.weights[idx][jdx].numel()\n",
    "print(count) # ought to be 2464 for normal MNIST-Seed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98bd5e01",
   "metadata": {},
   "source": [
    "# test new model dataset class with epoch resolution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c0821227",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset_base\n",
    "from pathlib import Path\n",
    "import json\n",
    "from shrp.datasets.dataset_epochs import ModelDatasetBaseEpochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae9f933a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-12 20:39:09,403\tINFO worker.py:1538 -- Started a local Ray instance.\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 35700/35700 [04:12<00:00, 141.54it/s]\n",
      "0it [00:00, ?it/s]\n",
      "ERROR:root:Exception occurred: training iteration 0 and epoch 1 don't match.\n",
      "Traceback (most recent call last):\n",
      "  File \"/netscratch2/kschuerholt/code/shrp/src/shrp/datasets/dataset_epochs.py\", line 204, in __init__\n",
      "    self.read_properties(\n",
      "  File \"/netscratch2/kschuerholt/code/shrp/src/shrp/datasets/dataset_epochs.py\", line 293, in read_properties\n",
      "    assert (\n",
      "AssertionError: training iteration 0 and epoch 1 don't match.\n",
      "692it [02:44,  4.39it/s]"
     ]
    }
   ],
   "source": [
    "config_key_list = []\n",
    "result_key_list = [\n",
    "    \"test_acc\",\n",
    "    \"training_iteration\",\n",
    "    \"ggap\",\n",
    "#     \"sparsity_ratio\",\n",
    "]\n",
    "property_keys = {\n",
    "    \"result_keys\": result_key_list,\n",
    "    \"config_keys\": config_key_list,\n",
    "}\n",
    "\n",
    "# path_root = Path('/netscratch2/dtaskiran/zoos/MNIST/tune_zoo_mnist_uniform/')\n",
    "path_root = Path('/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/')\n",
    "# path_root = Path('/netscratch2/kschuerholt/code/versai/model_zoos/zoos/CIFAR10/resnet19/kaiming_uniform/tune_zoo_cifar10_resnet18_kaiming_uniform')\n",
    "\n",
    "dataset_ep_test = ModelDatasetBaseEpochs(\n",
    "    root=path_root,\n",
    "    epoch_lst=list(range(0,51)),\n",
    "    mode=\"checkpoint\",\n",
    "    train_val_test=\"train\",  # determines whcih dataset split to use\n",
    "    ds_split=[0.7, 0.15,0.15],  #\n",
    "#     max_samples=10,\n",
    "    weight_threshold=float(\"inf\"),\n",
    "    filter_function=None,  # gets sample path as argument and returns True if model needs to be filtered out\n",
    "    property_keys=property_keys,\n",
    "    num_threads=4,\n",
    "    shuffle_path=True,\n",
    "    verbosity=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b805579c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "700\n",
      "700\n",
      "700\n",
      "700\n",
      "700\n",
      "700\n",
      "700\n"
     ]
    }
   ],
   "source": [
    "print(len(dataset_ep_test.data))\n",
    "print(len(dataset_ep_test.paths))\n",
    "print(len(dataset_ep_test.epochs))\n",
    "print(len(dataset_ep_test.labels))\n",
    "\n",
    "for key in dataset_ep_test.properties.keys():\n",
    "    print(len(dataset_ep_test.properties[key]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f85363bf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 127248 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "for idx in range(25):\n",
    "    sns.lineplot(x=dataset_ep_test.properties['training_iteration'][idx],y=dataset_ep_test.properties['test_acc'][idx],ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cc8b638",
   "metadata": {},
   "outputs": [],
   "source": [
    "lst = [\n",
    "    [1, 2, None],\n",
    "    [None, 2, None],\n",
    "    [3,1,2],\n",
    "    [1,None,2],\n",
    "]\n",
    "lst2 = [1,2,1,None,None,1,None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "085ce73e",
   "metadata": {},
   "outputs": [],
   "source": [
    "lst2clean = [ddx for ddx in lst2 if ddx]\n",
    "lst2clean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c63cad83",
   "metadata": {},
   "outputs": [],
   "source": [
    "lst = [[ddx for ddx in lst[idx] if ddx] for idx in range(len(lst))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fdf5943",
   "metadata": {},
   "outputs": [],
   "source": [
    "lst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9ce39589",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "module_list.0.weight\n",
      "torch.Size([8, 1, 5, 5])\n",
      "module_list.0.bias\n",
      "torch.Size([8])\n",
      "module_list.3.weight\n",
      "torch.Size([6, 8, 5, 5])\n",
      "module_list.3.bias\n",
      "torch.Size([6])\n",
      "module_list.6.weight\n",
      "torch.Size([4, 6, 2, 2])\n",
      "module_list.6.bias\n",
      "torch.Size([4])\n",
      "module_list.9.weight\n",
      "torch.Size([20, 36])\n",
      "module_list.9.bias\n",
      "torch.Size([20])\n",
      "module_list.11.weight\n",
      "torch.Size([10, 20])\n",
      "module_list.11.bias\n",
      "torch.Size([10])\n"
     ]
    }
   ],
   "source": [
    "for key in dataset_ep_test.data[0][2].keys():\n",
    "    print(key)\n",
    "    print(dataset_ep_test.data[0][2][key].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5a5907f2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 16])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "w = torch.randn([5,15])\n",
    "b = torch.randn([5])\n",
    "\n",
    "w2 = torch.cat([w,b.unsqueeze(dim=1)],dim=1)\n",
    "w2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "db80bb1e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "75"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w.shape\n",
    "w.numel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "37d0c23e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "75"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w3 = torch.randn([5,5,3])\n",
    "w3.numel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "275cf991",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 5, 3])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w4 = w.view(w3.shape)\n",
    "w4.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdb77164",
   "metadata": {},
   "source": [
    "# Test SimCLRDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f087bda7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset_base\n",
    "\n",
    "import logging\n",
    "logging.basicConfig(level=logging.WARNING)\n",
    "\n",
    "from pathlib import Path\n",
    "import json\n",
    "from shrp.datasets.dataset_simclr import SimCLRDataset\n",
    "\n",
    "from shrp.git_re_basin.git_re_basin import (\n",
    "    PermutationSpec,\n",
    "    zoo_cnn_permutation_spec,\n",
    "    weight_matching,\n",
    "    apply_permutation,\n",
    ")\n",
    "\n",
    "# import logging\n",
    "# logging.basicConfig(level=logging.DEBUG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "9119ae3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-15 00:07:43,744\tINFO worker.py:1538 -- Started a local Ray instance.\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [00:01<00:00,  5.15it/s]\n",
      "0it [00:00, ?it/s]\n",
      "ERROR:root:Exception occurred: training iteration 4 and epoch 5 don't match.\n",
      "Traceback (most recent call last):\n",
      "  File \"/netscratch2/kschuerholt/code/shrp/src/shrp/datasets/dataset_epochs.py\", line 204, in __init__\n",
      "    self.read_properties(\n",
      "  File \"/netscratch2/kschuerholt/code/shrp/src/shrp/datasets/dataset_epochs.py\", line 293, in read_properties\n",
      "    assert (\n",
      "AssertionError: training iteration 4 and epoch 5 don't match.\n",
      "10it [00:00, 220.84it/s]\n"
     ]
    }
   ],
   "source": [
    "config_key_list = []\n",
    "result_key_list = [\n",
    "    \"test_acc\",\n",
    "    \"training_iteration\",\n",
    "    \"ggap\",\n",
    "#     \"sparsity_ratio\",\n",
    "]\n",
    "property_keys = {\n",
    "    \"result_keys\": result_key_list,\n",
    "    \"config_keys\": config_key_list,\n",
    "}\n",
    "\n",
    "# path_root = Path('/netscratch2/dtaskiran/zoos/MNIST/tune_zoo_mnist_uniform/')\n",
    "path_root = Path('/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/')\n",
    "# path_root = Path('/netscratch2/kschuerholt/code/versai/model_zoos/zoos/CIFAR10/resnet19/kaiming_uniform/tune_zoo_cifar10_resnet18_kaiming_uniform')\n",
    "\n",
    "dataset_ep_test = SimCLRDataset(\n",
    "    root=path_root,\n",
    "    epoch_lst=[5],\n",
    "    mode=\"checkpoint\",\n",
    "#     mode=\"vector\",\n",
    "    permutations_number=0,\n",
    "    permutation_spec = zoo_cnn_permutation_spec(),\n",
    "    view_1_canonical = False,\n",
    "    view_2_canonical = False,\n",
    "    add_noise_view_1 = 0.0,  # [('input', 0.15), ('output', 0.013)]\n",
    "    add_noise_view_2 = 0.0,  # [('input', 0.15), ('output', 0.013)]\n",
    "    erase_augment=None,  # {\"p\": 0.5,\"scale\":(0.02,0.33),\"value\":0,\"mode\":\"block\"}\n",
    "    windowsize = 5,\n",
    "    train_val_test=\"train\",  # determines whcih dataset split to use\n",
    "    ds_split=[0.7, 0.15,0.15],  #\n",
    "    max_samples=15,\n",
    "    weight_threshold=float(\"inf\"),\n",
    "    filter_function=None,  # gets sample path as argument and returns True if model needs to be filtered out\n",
    "    property_keys=property_keys,\n",
    "    num_threads=4,\n",
    "    shuffle_path=True,\n",
    "    verbosity=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5331cd52",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_ep_test.set_module_window(windowsize = 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cc544f49",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'collections.OrderedDict'>\n",
      "10\n"
     ]
    }
   ],
   "source": [
    "print(type(dataset_ep_test.data[0][-1]))\n",
    "print(len(dataset_ep_test.data[0][-1]))\n",
    "# print(len(dataset_ep_test.data[0][-1][0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caaa031e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# idea:\n",
    "have index tensor 1: within components\n",
    "index tensor 2: numbers per component\n",
    "# permute both\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "63a5b112",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import torch\n",
    "import ray\n",
    "if ray.is_initialized():\n",
    "    ray.shutdown()\n",
    "from shrp.datasets.progress_bar import ProgressBar\n",
    "\n",
    "def precompute_permutations(ref_checkpoint,permutation_number, perm_spec, num_threads=6):\n",
    "        logging.info(\"start precomputing permutations\")\n",
    "        model_curr = ref_checkpoint\n",
    "        # find permutation of model to itself as reference\n",
    "        reference_permutation = weight_matching(\n",
    "            ps=perm_spec, params_a=model_curr, params_b=model_curr\n",
    "        )\n",
    "\n",
    "        logging.info(\"get random permutation dicts\")\n",
    "        # compute random permutations\n",
    "        permutation_dicts = []\n",
    "        for ndx in range(permutation_number):\n",
    "            perm = copy.deepcopy(reference_permutation)\n",
    "            for key in perm.keys():\n",
    "                # get permuted indecs for current layer\n",
    "                perm[key] = torch.randperm(perm[key].shape[0]).float()\n",
    "            # append to list of permutation dicts\n",
    "            permutation_dicts.append(perm)\n",
    "\n",
    "\n",
    "        logging.info(\"get permutation indices\")\n",
    "        # get permutation data\n",
    "        ## get reference checkpoint\n",
    "        ref_checkpoint = copy.deepcopy(model_curr)\n",
    "        ## vectoirze\n",
    "        ref_vec_global = vectorize_checkpoint(ref_checkpoint)\n",
    "        ref_vec_kernel = copy.deepcopy(ref_vec_global)\n",
    "        ## get reference index vec\n",
    "        for idx, module in enumerate(ref_vec_global):\n",
    "            # get global index of permutation between kernels\n",
    "            index_global = torch.ones(module.numel())*idx\n",
    "            index_global = index_global.view(module.shape)\n",
    "            ref_vec_global[idx] = index_global\n",
    "            # got local index of permutation within kernels\n",
    "            index_kernel = torch.tensor(list(range(module.numel())))\n",
    "            index_kernel = index_kernel.view(module.shape)\n",
    "            ref_vec_kernel[idx] = index_kernel\n",
    "        ## map to checkpoint\n",
    "        ref_checkpoint_global = vector_to_checkpoint(\n",
    "            vector=ref_vec_global, reference_checkpoint=ref_checkpoint\n",
    "        )\n",
    "        ref_checkpoint_kernel = vector_to_checkpoint(\n",
    "            vector=ref_vec_kernel, reference_checkpoint=ref_checkpoint\n",
    "        )\n",
    "                                             \n",
    "\n",
    "        ## init multiprocessing environment ############\n",
    "        ray.init(num_cpus=num_threads)\n",
    "        pb = ProgressBar(total=permutation_number)\n",
    "        pb_actor = pb.actor\n",
    "        # get permutations\n",
    "        permutations_global = []\n",
    "        permutations_kernel = []\n",
    "        for perm_dict in permutation_dicts:\n",
    "            perm_curr_global = compute_single_perm.remote(\n",
    "                reference_checkpoint=ref_checkpoint_global,\n",
    "                permutation_dict=perm_dict,\n",
    "                perm_spec=perm_spec,\n",
    "                pba=pb_actor,\n",
    "            )\n",
    "            \n",
    "            perm_curr_kernel = compute_single_perm.remote(\n",
    "                reference_checkpoint=ref_checkpoint_kernel,\n",
    "                permutation_dict=perm_dict,\n",
    "                perm_spec=perm_spec,\n",
    "                pba=pb_actor,\n",
    "            )\n",
    "            \n",
    "            permutations_global.append(perm_curr_global)\n",
    "            permutations_kernel.append(perm_curr_kernel)\n",
    "\n",
    "        permutations_global = ray.get(permutations_global)\n",
    "        permutations_kernel = ray.get(permutations_kernel)\n",
    "                \n",
    "        permutations_global = [torch.tensor([perm[0].item() for perm in perm_g]).int() for perm_g in permutations_global]\n",
    "\n",
    "        permutations = [(perm_g, perm_k) for (perm_g, perm_k) in zip(permutations_global,permutations_kernel)]\n",
    "                \n",
    "        ray.shutdown()\n",
    "        \n",
    "        return permutation_dicts, permutations\n",
    "    \n",
    "@ray.remote(num_returns=1)\n",
    "def compute_single_perm(reference_checkpoint, permutation_dict, perm_spec, pba):\n",
    "    # copy reference checkpoint\n",
    "    index_check = copy.deepcopy(reference_checkpoint)\n",
    "    # apply permutation on checkpoint\n",
    "    index_check_perm = apply_permutation(\n",
    "        ps=perm_spec, perm=permutation_dict, params=index_check\n",
    "    )\n",
    "    # vectorize\n",
    "    index_perm = vectorize_checkpoint(index_check_perm)\n",
    "    # update counter\n",
    "    pba.update.remote(1)\n",
    "    # return list\n",
    "    return index_perm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1bd8e917",
   "metadata": {},
   "outputs": [],
   "source": [
    "def vectorize_checkpoint(checkpoint):\n",
    "    out = []\n",
    "    # use only weights and biases\n",
    "    for key in checkpoint.keys():\n",
    "        if \"weight\" in key:\n",
    "            w = checkpoint[key]\n",
    "            # flatten to out_channels x n\n",
    "            w = w.view(w.shape[0], -1)\n",
    "            # cat biases to channels if they exist in checkpoint\n",
    "            if key.replace(\"weight\", \"bias\") in checkpoint:\n",
    "                b = checkpoint[key.replace(\"weight\", \"bias\")]\n",
    "                w = torch.cat([w, b.unsqueeze(dim=1)], dim=1)\n",
    "            # split weights in slices along output channel dims\n",
    "            w = torch.split(w, w.shape[0])\n",
    "            # extend out with new tokens, zero's (and only entry) is a list\n",
    "            out.extend(w[0])\n",
    "\n",
    "    return out\n",
    "\n",
    "\n",
    "def vector_to_checkpoint(vector, reference_checkpoint):\n",
    "    # make copy to prevent memory management issues\n",
    "    checkpoint = copy.deepcopy(reference_checkpoint)\n",
    "    # use only weights and biases\n",
    "    idx_start = 0\n",
    "    for key in checkpoint.keys():\n",
    "        if \"weight\" in key:\n",
    "            # get correct slice of modules out of vec sequence\n",
    "            out_channels = checkpoint[key].shape[0]\n",
    "            idx_end = idx_start + out_channels\n",
    "            w = vector[idx_start:idx_end]\n",
    "            # update start\n",
    "            idx_start = idx_end\n",
    "            # get weight matrix from list of vectors\n",
    "            w = torch.stack(w, dim=0)\n",
    "            # extract bias\n",
    "            if key.replace(\"weight\", \"bias\") in checkpoint:\n",
    "                b = w[:, -1]\n",
    "                checkpoint[key.replace(\"weight\", \"bias\")] = b\n",
    "                w = w[:, :-1]\n",
    "            # reshape weight vector\n",
    "            w = w.view(checkpoint[key].shape)\n",
    "            logging.debug(\n",
    "                f\"vector_to_checkpoint: created weight {w.shape} for checkpoint weight {checkpoint[key].shape}\"\n",
    "            )\n",
    "            checkpoint[key] = w\n",
    "    return checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4cbb720e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-14 23:54:58,595\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "permdicts, perms = precompute_permutations(ref_checkpoint=dataset_ep_test.data[0][-1],permutation_number=3, perm_spec=zoo_cnn_permutation_spec())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0f3abe09",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'P_bg0': tensor([5., 1., 2., 7., 6., 3., 4., 0.]),\n",
       "  'P_bg1': tensor([1., 3., 2., 4., 5., 0.]),\n",
       "  'P_bg2': tensor([2., 1., 3., 0.]),\n",
       "  'P_bg3': tensor([19.,  0.,  8., 12.,  9.,  6.,  7.,  5., 10., 13.,  1.,  3., 17., 11.,\n",
       "           4., 16., 18., 14.,  2., 15.])},\n",
       " {'P_bg0': tensor([2., 0., 4., 5., 1., 6., 7., 3.]),\n",
       "  'P_bg1': tensor([0., 2., 1., 4., 3., 5.]),\n",
       "  'P_bg2': tensor([2., 0., 3., 1.]),\n",
       "  'P_bg3': tensor([12., 13., 18., 15.,  3., 16.,  7.,  0., 19.,  6., 11.,  1., 17., 10.,\n",
       "           4.,  5.,  9., 14.,  2.,  8.])},\n",
       " {'P_bg0': tensor([6., 2., 7., 1., 4., 0., 5., 3.]),\n",
       "  'P_bg1': tensor([3., 1., 0., 2., 5., 4.]),\n",
       "  'P_bg2': tensor([0., 2., 1., 3.]),\n",
       "  'P_bg3': tensor([ 8., 16.,  3.,  2., 14., 13.,  7.,  4., 15., 17., 19., 11., 12.,  1.,\n",
       "           6., 18.,  5.,  0.,  9., 10.])}]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "permdicts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6af939f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3\n",
      "2\n"
     ]
    }
   ],
   "source": [
    "print(len(perms))\n",
    "print(len(perms[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dbb477c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'tuple'>\n",
      "2\n",
      "tensor([ 5,  1,  2,  7,  6,  3,  4,  0,  9, 11, 10, 12, 13,  8, 16, 15, 17, 14,\n",
      "        37, 18, 26, 30, 27, 24, 25, 23, 28, 31, 19, 21, 35, 29, 22, 34, 36, 32,\n",
      "        20, 33, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47], dtype=torch.int32)\n"
     ]
    }
   ],
   "source": [
    "print(type(perms[0]))\n",
    "print(len(perms[0]))\n",
    "print(perms[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f9fe74e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def permute_model_vector(vec,idx_start,window,perm):\n",
    "    # create index vector of tokens\n",
    "    index = list(range(len(vec)))\n",
    "    \n",
    "    # apply global permutation on index\n",
    "    # using (slices of) the permuted index to access tokens equals permuting all tokens and slicing after\n",
    "    perm_glob = perm[0]\n",
    "    index = [index[idx] for idx in perm_glob]\n",
    "    \n",
    "    # slice index\n",
    "    idx_end = idx_start + window\n",
    "    index = index[idx_start:idx_end]\n",
    "    \n",
    "    # slice token sequence\n",
    "    vec = [vec[idx] for idx in index] \n",
    "        \n",
    "    # slice permutations\n",
    "    perm_loc = perm[1]\n",
    "    perm_loc = [perm_loc[idx] for idx in index]\n",
    "\n",
    "    # apply token permutation\n",
    "    vec = [vecdx[permdx] for (vecdx,permdx) in zip(vec,perm_loc)]\n",
    "    \n",
    "    # return tokens\n",
    "    return vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "32033871",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'collections.OrderedDict'>\n",
      "<class 'list'>\n",
      "48\n"
     ]
    }
   ],
   "source": [
    "check1 = dataset_ep_test.data[0][-1]\n",
    "print(type(check1))\n",
    "vec1 = vectorize_checkpoint(check1)\n",
    "print(type(vec1))\n",
    "print(len(vec1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96d1dae9",
   "metadata": {},
   "source": [
    "# Test 1: computational load - compare permutations on checkpoint with permutation on vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a40adcdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6573c702",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.57 µs ± 30.8 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# baseline - how long does it take to draw a random integer\n",
    "idx = random.randint(0,len(dataset_ep_test.data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ee5ce77e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare a list of vecs\n",
    "vecs = []\n",
    "for idx in range(len(dataset_ep_test.data)):\n",
    "    checkdx = dataset_ep_test.data[idx][-1]\n",
    "    vecdx = vectorize_checkpoint(checkdx)\n",
    "    vecs.append(vecdx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "992864ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "154 µs ± 312 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# compute random sliced permutations\n",
    "idx = random.randint(0,len(vecs)-1)\n",
    "perm1 = perms[0]\n",
    "vec1 = vecs[idx]\n",
    "vec2 = permute_model_vector(vec=vec1,idx_start=2,window=13,perm=perm1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "id": "93ddebe9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7"
      ]
     },
     "execution_count": 143,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# comparison: permute checkpoints -> vectorize -> slice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "30d1258a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "616 µs ± 1.21 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "idx = random.randint(0,len(dataset_ep_test.data)-1)\n",
    "\n",
    "check1 = dataset_ep_test.data[idx][-1]\n",
    "\n",
    "permd = permdicts[0]\n",
    "\n",
    "check_perm = apply_permutation(ps=zoo_cnn_permutation_spec(), perm=permd, params=check1)\n",
    "\n",
    "vec3 = vectorize_checkpoint(check_perm )\n",
    "\n",
    "vec4 = vec3[2:15]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c4a5eca",
   "metadata": {},
   "source": [
    "permutation on vectorized form appears <25% of the runtime of permutation on the checkpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e91ff96",
   "metadata": {},
   "source": [
    "# test 2: check slice of permutation ar the same to sliced permutations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "id": "daf50ccd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "# check that slice of permuted model is the same as permute/slice at the same time\n",
    "\n",
    "# since all true -> that works :)\n",
    "perm1 = perms[0]\n",
    "sstart = 7\n",
    "slen = 34\n",
    "vec2 = permute_model_vector(vec=vec1,idx_start=0,window=len(vec1),perm=perm1)\n",
    "vec3 = permute_model_vector(vec=vec1,idx_start=sstart,window=slen,perm=perm1)\n",
    "\n",
    "vec23 = vec2[sstart:sstart+slen]\n",
    "for m1,m2 in zip(vec23,vec3):\n",
    "    print(torch.allclose(m1,m2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "070a58a1",
   "metadata": {},
   "source": [
    "# test 3: test that permuted models have the same mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2414a0b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# permute model with full window\n",
    "# cast back to checkpoint\n",
    "# load both checkpoint\n",
    "# verify closeness of predictions\n",
    "\n",
    "# repeat for 10 models and 15 iterations\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "45ff0f69",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load config\n",
    "from shrp.models.def_net import NNmodule\n",
    "config_path = dataset_ep_test.paths[0][0].joinpath('params.json')\n",
    "config = json.load(config_path.open('r'))\n",
    "config['training::batchsize'] = 1000\n",
    "model = NNmodule(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "196d9923",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load dataset\n",
    "imgdata = torch.load(config['dataset::dump'].replace('netscratch','netscratch2'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "172e1c1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from shrp.datasets.def_FastTensorDataLoader import FastTensorDataLoader\n",
    "n_samples = 1000\n",
    "\n",
    "testset_1, testset_2 = torch.utils.data.random_split(\n",
    "            imgdata['testset'], [n_samples, len(imgdata['testset'])-n_samples], generator=torch.Generator().manual_seed(42))\n",
    "\n",
    "testloader_raw = torch.utils.data.DataLoader(\n",
    "    dataset=testset_1, batch_size=len(testset_1), shuffle=True\n",
    "        )\n",
    "assert testloader_raw.__len__() == 1, \"temp testloader has more than one batch\"\n",
    "for test_data, test_labels in testloader_raw:\n",
    "    pass\n",
    "\n",
    "testset = torch.utils.data.TensorDataset(test_data, test_labels)\n",
    "testloader = FastTensorDataLoader(testset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "a846fb4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-15 00:06:50,151\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "permdicts, perms = precompute_permutations(ref_checkpoint=dataset_ep_test.data[0][-1],permutation_number=15, perm_spec=zoo_cnn_permutation_spec())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "9b70f1e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load first checkpoint\n",
      "compute predictions of first checkpoint\n",
      "load permuted checkpoint\n",
      "compute predictions of permuted checkpoint\n",
      "check equality\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "check1 = dataset_ep_test.data[0][-1]\n",
    "permd = permdicts[0]\n",
    "check_perm = apply_permutation(ps=zoo_cnn_permutation_spec(), perm=permd, params=check1)\n",
    "\n",
    "print('load first checkpoint')\n",
    "model.model.load_state_dict(check1)\n",
    "print('compute predictions of first checkpoint')\n",
    "with torch.no_grad():\n",
    "    pred1 = model(test_data)\n",
    "    \n",
    "print('load permuted checkpoint')\n",
    "model.model.load_state_dict(check_perm)\n",
    "print('compute predictions of permuted checkpoint')\n",
    "with torch.no_grad():\n",
    "    pred2 = model(test_data)\n",
    "    \n",
    "print('check equality')\n",
    "torch.allclose(pred1,pred2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "e9204c7a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███████████████████████████████████████████████▍                                                                                                              | 3/10 [00:52<02:02, 17.52s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [03:01<00:00, 18.12s/it]\n"
     ]
    }
   ],
   "source": [
    "import tqdm\n",
    "\n",
    "n_models = 10\n",
    "n_perms = 15\n",
    "\n",
    "res = []\n",
    "\n",
    "for idx in tqdm.tqdm(range(n_models)):\n",
    "    for jdx in range(n_perms):\n",
    "\n",
    "        check1 = dataset_ep_test.data[idx][-1]\n",
    "        permd = permdicts[jdx]\n",
    "        check_perm = apply_permutation(ps=zoo_cnn_permutation_spec(), perm=permd, params=check1)\n",
    "\n",
    "#         print('load first checkpoint')\n",
    "        model.model.load_state_dict(check1)\n",
    "#         print('compute predictions of first checkpoint')\n",
    "        with torch.no_grad():\n",
    "            pred1 = model(test_data)\n",
    "\n",
    "#         print('load permuted checkpoint')\n",
    "        model.model.load_state_dict(check_perm)\n",
    "#         print('compute predictions of permuted checkpoint')\n",
    "        with torch.no_grad():\n",
    "            pred2 = model(test_data)\n",
    "\n",
    "#         print('check equality')\n",
    "        restmp = torch.allclose(pred1,pred2)\n",
    "        if not restmp:\n",
    "            print(torch.allclose(pred1,pred2,atol=1e-05, rtol=1e-03))\n",
    "        res.append(restmp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "e864ceb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.9866666666666667\n"
     ]
    }
   ],
   "source": [
    "print(sum(res)/len(res))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b33d89ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test the same on vector level"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "6f633df7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare a list of vecs\n",
    "vecs = []\n",
    "for idx in range(len(dataset_ep_test.data)):\n",
    "    checkdx = dataset_ep_test.data[idx][-1]\n",
    "    vecdx = vectorize_checkpoint(checkdx)\n",
    "    vecs.append(vecdx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "b9f07845",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load first checkpoint\n",
      "compute predictions of first checkpoint\n",
      "load permuted checkpoint\n",
      "compute predictions of permuted checkpoint\n",
      "check equality\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vec1 = vecs[0]\n",
    "perm1 = perms[0]\n",
    "\n",
    "vecperm = permute_model_vector(vec=vec1,idx_start=0,window=len(vec1),perm=perm1)\n",
    "\n",
    "check1 = vector_to_checkpoint(vector=vec1, reference_checkpoint=check1)\n",
    "check_perm = vector_to_checkpoint(vector=vecperm, reference_checkpoint=check1)\n",
    "\n",
    "print('load first checkpoint')\n",
    "model.model.load_state_dict(check1)\n",
    "print('compute predictions of first checkpoint')\n",
    "with torch.no_grad():\n",
    "    pred1 = model(test_data)\n",
    "    \n",
    "print('load permuted checkpoint')\n",
    "model.model.load_state_dict(check_perm)\n",
    "print('compute predictions of permuted checkpoint')\n",
    "with torch.no_grad():\n",
    "    pred2 = model(test_data)\n",
    "    \n",
    "print('check equality')\n",
    "torch.allclose(pred1,pred2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "44257999",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███████████████████████████████████████████████▍                                                                                                              | 3/10 [00:55<02:10, 18.58s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [03:04<00:00, 18.49s/it]\n"
     ]
    }
   ],
   "source": [
    "import tqdm\n",
    "\n",
    "n_models = 10\n",
    "n_perms = 15\n",
    "\n",
    "res = []\n",
    "\n",
    "for idx in tqdm.tqdm(range(n_models)):\n",
    "    for jdx in range(n_perms):\n",
    "\n",
    "        vec1 = vecs[idx]\n",
    "        perm1 = perms[jdx]\n",
    "\n",
    "        vecperm = permute_model_vector(vec=vec1,idx_start=0,window=len(vec1),perm=perm1)\n",
    "\n",
    "        check1 = vector_to_checkpoint(vector=vec1, reference_checkpoint=check1)\n",
    "        check_perm = vector_to_checkpoint(vector=vecperm, reference_checkpoint=check1)\n",
    "\n",
    "#         print('load first checkpoint')\n",
    "        model.model.load_state_dict(check1)\n",
    "#         print('compute predictions of first checkpoint')\n",
    "        with torch.no_grad():\n",
    "            pred1 = model(test_data)\n",
    "\n",
    "#         print('load permuted checkpoint')\n",
    "        model.model.load_state_dict(check_perm)\n",
    "#         print('compute predictions of permuted checkpoint')\n",
    "        with torch.no_grad():\n",
    "            pred2 = model(test_data)\n",
    "\n",
    "#         print('check equality')\n",
    "        restmp = torch.allclose(pred1,pred2)\n",
    "        if not restmp:\n",
    "            print(torch.allclose(pred1,pred2,atol=1e-05, rtol=1e-03))\n",
    "        res.append(restmp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "7136fdd8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.,  2.,  3.],\n",
       "        [21., 22., 23.],\n",
       "        [31., 32., 33.]])"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lst = [\n",
    "    torch.tensor([1,2,3]),\n",
    "    torch.tensor([21,22,23]),\n",
    "    torch.tensor([31,32,33]),\n",
    "]\n",
    "\n",
    "tst = torch.stack(lst,dim=0).float()\n",
    "tst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "d9024b1b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(18.6667)"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.mean(tst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "26ddadd0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(13.2571)"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.std(tst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25850cca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test dataset with standardization, permutation, map to canonical\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "726c6b56",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset_base\n",
    "\n",
    "import logging\n",
    "# logging.basicConfig(level=logging.INFO)\n",
    "logging.basicConfig(level=logging.DEBUG)\n",
    "\n",
    "from pathlib import Path\n",
    "import json\n",
    "from shrp.datasets.dataset_simclr import SimCLRDataset\n",
    "\n",
    "from shrp.git_re_basin.git_re_basin import (\n",
    "    PermutationSpec,\n",
    "    zoo_cnn_permutation_spec,\n",
    "    weight_matching,\n",
    "    apply_permutation,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "72b9b91a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-18 00:13:11,412\tINFO worker.py:1538 -- Started a local Ray instance.\n",
      "INFO:root:loading checkpoints from [PosixPath('/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform')]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30/30 [00:03<00:00,  9.11it/s]\n",
      "INFO:root:Data loaded. found 10 usable samples out of potential 30 samples.\n",
      "INFO:root:Load properties for samples from paths.\n",
      "INFO:root:### load data for dict_keys(['test_acc', 'training_iteration', 'ggap'])\n",
      "10it [00:00, 13.75it/s]\n",
      "INFO:root:Properties loaded.\n",
      "INFO:root:prepare canonical form\n",
      "2023-03-18 00:13:22,034\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "preparing computing canon form...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:init dataset length\n",
      "INFO:root:Get Positions\n",
      "INFO:root:vectorize data\n",
      "INFO:root:Get layer mapping\n",
      "DEBUG:root:layer mapping: {'0': {'start_idx': 0, 'end_idx': 7}, '1': {'start_idx': 8, 'end_idx': 13}, '2': {'start_idx': 14, 'end_idx': 17}, '3': {'start_idx': 18, 'end_idx': 37}, '4': {'start_idx': 38, 'end_idx': 47}}\n",
      "INFO:root:Get layer-wise mean and std\n",
      "INFO:root:Apply standardization\n",
      "INFO:root:Discover tokensize\n"
     ]
    }
   ],
   "source": [
    "config_key_list = []\n",
    "result_key_list = [\n",
    "    \"test_acc\",\n",
    "    \"training_iteration\",\n",
    "    \"ggap\",\n",
    "#     \"sparsity_ratio\",\n",
    "]\n",
    "property_keys = {\n",
    "    \"result_keys\": result_key_list,\n",
    "    \"config_keys\": config_key_list,\n",
    "}\n",
    "\n",
    "# path_root = Path('/netscratch2/dtaskiran/zoos/MNIST/tune_zoo_mnist_uniform/')\n",
    "path_root = Path('/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/')\n",
    "# path_root = Path('/netscratch2/kschuerholt/code/versai/model_zoos/zoos/CIFAR10/resnet19/kaiming_uniform/tune_zoo_cifar10_resnet18_kaiming_uniform')\n",
    "\n",
    "dataset_ep_test = SimCLRDataset(\n",
    "    root=path_root,\n",
    "    epoch_lst=[1,5,25],\n",
    "#     mode=\"checkpoint\",\n",
    "    mode=\"vector\",\n",
    "    permutations_number=0,\n",
    "    permutation_spec = zoo_cnn_permutation_spec(),\n",
    "#     view_1_canonical = False,\n",
    "    view_1_canonical = False,\n",
    "    view_2_canonical = True,\n",
    "#     view_2_canonical = False,\n",
    "    add_noise_view_1 = 0.1,  # [('input', 0.15), ('output', 0.013)]\n",
    "    add_noise_view_2 = 0.0,  # [('input', 0.15), ('output', 0.013)]\n",
    "    noise_multiplicative=True,\n",
    "    erase_augment=None,  # {\"p\": 0.5,\"scale\":(0.02,0.33),\"value\":0,\"mode\":\"block\"}\n",
    "    windowsize = 12,\n",
    "    standardize=True,\n",
    "    train_val_test=\"train\",  # determines whcih dataset split to use\n",
    "    ds_split=[0.7, 0.15,0.15],  #\n",
    "    max_samples=15,\n",
    "    weight_threshold=float(\"inf\"),\n",
    "    filter_function=None,  # gets sample path as argument and returns True if model needs to be filtered out\n",
    "    property_keys=property_keys,\n",
    "    num_threads=4,\n",
    "    shuffle_path=True,\n",
    "    verbosity=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d46bc894",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Checks\n",
    "# -[x] All options run without error\n",
    "# -[x] NAN / INFS\n",
    "# -[x] run permutation checks from above again\n",
    "# -[x] check lenghts of windows\n",
    "# -[x] check standardization worsk\n",
    "# -[x] check noising works\n",
    "# -[x] check augmentations are applied correctly\n",
    "# -[x] check relation of all datapoints is preserved\n",
    "# -[x] check positions\n",
    "\n",
    "# =[]comment all code \n",
    "# =[]code review while commenting\n",
    "# =[]load models, compute test_acc, compare to recorded one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a4791954",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "30"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_ep_test.__len__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d41cd090",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dataset_ep_test.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "74696f9e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dataset_ep_test.data[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fbd1df22",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "a947cbe7",
   "metadata": {},
   "source": [
    "### test noise "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "54bdb2fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.6178,  1.0873,  0.0407,  0.0093,  0.5361, -0.5169,  0.6908, -0.2309,\n",
      "        -0.3285,  0.1974,  0.4596,  0.4653, -0.0077,  0.2002,  0.4147, -0.5666,\n",
      "         1.7816,  1.4433,  0.5817,  1.0785,  0.2645,  0.6473, -0.2172,  0.7955,\n",
      "        -0.1908])\n",
      "tensor([ 0.8240,  0.9750,  0.0380,  0.0066,  0.4346, -0.5247,  0.6876, -0.2180,\n",
      "        -0.2840,  0.1956,  0.4612,  0.4618, -0.0086,  0.2470,  0.3720, -0.5622,\n",
      "         1.3629,  1.6588,  0.5280,  1.1187,  0.2334,  0.5823, -0.2271,  0.7855,\n",
      "        -0.1828])\n",
      "tensor([ 0.7212,  1.0149,  0.0400,  0.0066,  0.4646, -0.4545,  0.6207, -0.2456,\n",
      "        -0.3473,  0.1681,  0.4642,  0.3805, -0.0081,  0.2214,  0.4619, -0.5562,\n",
      "         1.5033,  1.5660,  0.6301,  1.1723,  0.2358,  0.6135, -0.2521,  0.8279,\n",
      "        -0.1709])\n",
      "tensor([ 0.7916,  1.1427,  0.0344,  0.0080,  0.4719, -0.4345,  0.7106, -0.2521,\n",
      "        -0.2859,  0.1705,  0.4819,  0.3783, -0.0058,  0.1908,  0.3508, -0.3702,\n",
      "         1.7695,  1.4537,  0.6129,  1.0355,  0.2838,  0.6718, -0.2450,  0.8243,\n",
      "        -0.1934])\n",
      "tensor([ 0.7389,  1.1845,  0.0446,  0.0081,  0.4063, -0.4910,  0.7072, -0.2420,\n",
      "        -0.2235,  0.2006,  0.4358,  0.4096, -0.0072,  0.2246,  0.4255, -0.4792,\n",
      "         1.5815,  1.7055,  0.5801,  1.1538,  0.2635,  0.5545, -0.2333,  0.7678,\n",
      "        -0.1939])\n"
     ]
    }
   ],
   "source": [
    "dataset_ep_test.set_module_window(len(dataset_ep_test.data[0][-1]))\n",
    "index = 21\n",
    "for _ in range(5):\n",
    "    tw,_,_,_,pw = dataset_ep_test.__getitem__(index)\n",
    "    print(tw[15])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db8b8bc0",
   "metadata": {},
   "source": [
    "seems to work as expected "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "f5c98392",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00548_548_seed=549_2021-07-27_15-51-15#_#epoch_1\n",
      "0.18335125998770743 - 1\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00935_935_seed=936_2021-07-28_08-01-46#_#epoch_5\n",
      "0.23044714197910265 - 5\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00257_257_seed=258_2021-07-27_03-43-04#_#epoch_1\n",
      "0.19817916410571604 - 1\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00786_786_seed=787_2021-07-28_01-50-08#_#epoch_5\n",
      "0.21838506453595574 - 5\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00321_321_seed=322_2021-07-27_06-12-40#_#epoch_5\n",
      "0.21842347879532883 - 5\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00935_935_seed=936_2021-07-28_08-01-46#_#epoch_25\n",
      "0.7004456054087277 - 25\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00935_935_seed=936_2021-07-28_08-01-46#_#epoch_5\n",
      "0.23044714197910265 - 5\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00599_599_seed=600_2021-07-27_18-04-45#_#epoch_1\n",
      "0.1958358942839582 - 1\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00827_827_seed=828_2021-07-28_03-34-52#_#epoch_1\n",
      "0.19564382298709282 - 1\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00689_689_seed=690_2021-07-27_21-54-40#_#epoch_5\n",
      "0.1958743085433313 - 5\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00548_548_seed=549_2021-07-27_15-51-15#_#epoch_1\n",
      "0.18335125998770743 - 1\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00548_548_seed=549_2021-07-27_15-51-15#_#epoch_1\n",
      "0.18335125998770743 - 1\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00689_689_seed=690_2021-07-27_21-54-40#_#epoch_1\n",
      "0.1977181929932391 - 1\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00321_321_seed=322_2021-07-27_06-12-40#_#epoch_1\n",
      "0.19652735095267362 - 1\n",
      "/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00599_599_seed=600_2021-07-27_18-04-45#_#epoch_5\n",
      "0.21231561155500922 - 5\n"
     ]
    }
   ],
   "source": [
    "### check relation - checkpoint - properties\n",
    "for _ in range(15):\n",
    "    index = random.randint(0,len(dataset_ep_test)-1)\n",
    "    _,l,_,_,pw = dataset_ep_test.__getitem__(index)\n",
    "    mdx,edx = dataset_ep_test._index[index]\n",
    "    props = dataset_ep_test.properties\n",
    "    print(l)\n",
    "    print(f\"{props['test_acc'][mdx][edx]} - {props['training_iteration'][mdx][edx]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0cda91a",
   "metadata": {},
   "source": [
    "### check positions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e02930a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(0, 0),\n",
       " (0, 1),\n",
       " (0, 2),\n",
       " (0, 3),\n",
       " (0, 4),\n",
       " (0, 5),\n",
       " (0, 6),\n",
       " (0, 7),\n",
       " (1, 0),\n",
       " (1, 1),\n",
       " (1, 2),\n",
       " (1, 3),\n",
       " (1, 4),\n",
       " (1, 5),\n",
       " (2, 0),\n",
       " (2, 1),\n",
       " (2, 2),\n",
       " (2, 3),\n",
       " (3, 0),\n",
       " (3, 1),\n",
       " (3, 2),\n",
       " (3, 3),\n",
       " (3, 4),\n",
       " (3, 5),\n",
       " (3, 6),\n",
       " (3, 7),\n",
       " (3, 8),\n",
       " (3, 9),\n",
       " (3, 10),\n",
       " (3, 11),\n",
       " (3, 12),\n",
       " (3, 13),\n",
       " (3, 14),\n",
       " (3, 15),\n",
       " (3, 16),\n",
       " (3, 17),\n",
       " (3, 18),\n",
       " (3, 19),\n",
       " (4, 0),\n",
       " (4, 1),\n",
       " (4, 2),\n",
       " (4, 3),\n",
       " (4, 4),\n",
       " (4, 5),\n",
       " (4, 6),\n",
       " (4, 7),\n",
       " (4, 8),\n",
       " (4, 9)]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_ep_test.positions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62c0c3e3",
   "metadata": {},
   "source": [
    "### check standardization\n",
    "seems to work well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "172b8261",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'0': {'start_idx': 0,\n",
       "  'end_idx': 7,\n",
       "  'mean': tensor(0.3796),\n",
       "  'std': tensor(0.7530)},\n",
       " '1': {'start_idx': 8,\n",
       "  'end_idx': 13,\n",
       "  'mean': tensor(0.3528),\n",
       "  'std': tensor(0.7909)},\n",
       " '2': {'start_idx': 14,\n",
       "  'end_idx': 17,\n",
       "  'mean': tensor(0.3028),\n",
       "  'std': tensor(0.5307)},\n",
       " '3': {'start_idx': 18,\n",
       "  'end_idx': 37,\n",
       "  'mean': tensor(0.3382),\n",
       "  'std': tensor(0.5479)},\n",
       " '4': {'start_idx': 38,\n",
       "  'end_idx': 47,\n",
       "  'mean': tensor(0.4690),\n",
       "  'std': tensor(0.5501)}}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_ep_test.layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "9021ec16",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10\n",
      "30\n",
      "mean: 1.6513841405441099e-09| std: 0.999972939491272\n",
      "nans: False | infs: False\n"
     ]
    }
   ],
   "source": [
    "# stack all components\n",
    "\n",
    "data = dataset_ep_test.data\n",
    "print(len(data))\n",
    "data = [ddx for idx in range(len(data)) for ddx in data[idx]]\n",
    "print(len(data))\n",
    "data = [torch.cat(ddx,dim=0)for ddx in data]\n",
    "data = torch.stack(data,dim=0)\n",
    "data.shape\n",
    "\n",
    "print(f'mean: {torch.mean(data)}| std: {torch.std(data)}')\n",
    "print(f'nans: {torch.isnan(data).any()} | infs: {torch.isinf(data).any()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a7c55be",
   "metadata": {},
   "source": [
    "### check permutations (again)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9ea077d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset_base\n",
    "\n",
    "import logging\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "\n",
    "from pathlib import Path\n",
    "import json\n",
    "from shrp.datasets.dataset_simclr import SimCLRDataset\n",
    "\n",
    "from shrp.git_re_basin.git_re_basin import (\n",
    "    PermutationSpec,\n",
    "    zoo_cnn_permutation_spec,\n",
    "    weight_matching,\n",
    "    apply_permutation,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "58aceb76",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-03-16 22:11:47,569\tINFO worker.py:1538 -- Started a local Ray instance.\n",
      "INFO:root:loading checkpoints from [PosixPath('/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform')]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 105/105 [00:02<00:00, 43.87it/s]\n",
      "INFO:root:Data loaded. found 35 usable samples out of potential 105 samples.\n",
      "INFO:root:Load properties for samples from paths.\n",
      "INFO:root:### load data for dict_keys(['test_acc', 'training_iteration', 'ggap'])\n",
      "35it [00:00, 80.65it/s]\n",
      "INFO:root:Properties loaded.\n",
      "INFO:root:prepare canonical form\n",
      "2023-03-16 22:11:56,780\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "preparing computing canon form...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:init dataset length\n",
      "INFO:root:Get Positions\n",
      "INFO:root:vectorize data\n",
      "INFO:root:init permutations\n",
      "INFO:root:start precomputing permutations\n",
      "INFO:root:0/P_bg3: 0.0\n",
      "INFO:root:0/P_bg0: 0.0\n",
      "INFO:root:0/P_bg2: 0.0\n",
      "INFO:root:0/P_bg1: 0.0\n",
      "INFO:root:get random permutation dicts\n",
      "INFO:root:get permutation indices\n",
      "2023-03-16 22:12:06,684\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "config_key_list = []\n",
    "result_key_list = [\n",
    "    \"test_acc\",\n",
    "    \"training_iteration\",\n",
    "    \"ggap\",\n",
    "#     \"sparsity_ratio\",\n",
    "]\n",
    "property_keys = {\n",
    "    \"result_keys\": result_key_list,\n",
    "    \"config_keys\": config_key_list,\n",
    "}\n",
    "\n",
    "# path_root = Path('/netscratch2/dtaskiran/zoos/MNIST/tune_zoo_mnist_uniform/')\n",
    "path_root = Path('/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/')\n",
    "# path_root = Path('/netscratch2/kschuerholt/code/versai/model_zoos/zoos/CIFAR10/resnet19/kaiming_uniform/tune_zoo_cifar10_resnet18_kaiming_uniform')\n",
    "\n",
    "dataset_ep_test = SimCLRDataset(\n",
    "    root=path_root,\n",
    "    epoch_lst=[1,5,25],\n",
    "#     mode=\"checkpoint\",\n",
    "    mode=\"vector\",\n",
    "    permutations_number=20,\n",
    "    permutation_spec = zoo_cnn_permutation_spec(),\n",
    "    view_1_canonical = False,\n",
    "    view_2_canonical = True,\n",
    "#     view_2_canonical = False,\n",
    "    add_noise_view_1 = 0.0,  # [('input', 0.15), ('output', 0.013)]\n",
    "    add_noise_view_2 = 0.0,  # [('input', 0.15), ('output', 0.013)]\n",
    "    erase_augment=None,  # {\"p\": 0.5,\"scale\":(0.02,0.33),\"value\":0,\"mode\":\"block\"}\n",
    "#     windowsize = 12,\n",
    "    standardize=False,\n",
    "    train_val_test=\"train\",  # determines whcih dataset split to use\n",
    "    ds_split=[0.7, 0.15,0.15],  #\n",
    "    max_samples=50,\n",
    "    weight_threshold=float(\"inf\"),\n",
    "    filter_function=None,  # gets sample path as argument and returns True if model needs to be filtered out\n",
    "    property_keys=property_keys,\n",
    "    num_threads=4,\n",
    "    shuffle_path=True,\n",
    "    verbosity=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "70e3c0ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_ep_test.set_module_window(len(dataset_ep_test.data[0][-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c66336eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:cuda unavailable:: train model on cpu\n",
      "INFO:root:=> creating model CNN\n",
      "INFO:root:initialze model\n"
     ]
    }
   ],
   "source": [
    "# load config\n",
    "import torch\n",
    "from shrp.models.def_net import NNmodule\n",
    "config_path = dataset_ep_test.paths[0][0].joinpath('params.json')\n",
    "config = json.load(config_path.open('r'))\n",
    "config['training::batchsize'] = 1000\n",
    "model = NNmodule(config)\n",
    "# load dataset\n",
    "imgdata = torch.load(config['dataset::dump'].replace('netscratch','netscratch2'))\n",
    "from shrp.datasets.def_FastTensorDataLoader import FastTensorDataLoader\n",
    "n_samples = 1000\n",
    "\n",
    "testset_1, testset_2 = torch.utils.data.random_split(\n",
    "            imgdata['testset'], [n_samples, len(imgdata['testset'])-n_samples], generator=torch.Generator().manual_seed(42))\n",
    "\n",
    "testloader_raw = torch.utils.data.DataLoader(\n",
    "    dataset=testset_1, batch_size=len(testset_1), shuffle=True\n",
    "        )\n",
    "assert testloader_raw.__len__() == 1, \"temp testloader has more than one batch\"\n",
    "for test_data, test_labels in testloader_raw:\n",
    "    pass\n",
    "\n",
    "testset = torch.utils.data.TensorDataset(test_data, test_labels)\n",
    "testloader = FastTensorDataLoader(testset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "03cb6cc0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                                                                                                      | 0/10 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|███████████████████████████████▌                                                                                                                              | 2/10 [01:35<06:18, 47.32s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███████████████████████████████████████████████▍                                                                                                              | 3/10 [02:22<05:32, 47.53s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|███████████████████████████████████████████████████████████████▏                                                                                              | 4/10 [03:11<04:47, 47.99s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|███████████████████████████████████████████████████████████████████████████████                                                                               | 5/10 [04:03<04:07, 49.44s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▍                               | 8/10 [06:24<01:35, 47.89s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 90%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏               | 9/10 [07:17<00:49, 49.68s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [08:10<00:00, 49.08s/it]\n"
     ]
    }
   ],
   "source": [
    "import tqdm\n",
    "import random\n",
    "from shrp.datasets.dataset_simclr import vector_to_checkpoint\n",
    "\n",
    "n_models = 10\n",
    "n_perms = 15\n",
    "\n",
    "res = []\n",
    "\n",
    "check = dataset_ep_test.reference_checkpoint\n",
    "for idx in tqdm.tqdm(range(n_models)):\n",
    "        \n",
    "    for jdx in range(n_perms):\n",
    "        \n",
    "        edx = random.randint(0,2)\n",
    "        jdx = idx*3+edx\n",
    "        vec1,lab1,_,_,_ = dataset_ep_test.__getitem__(jdx)\n",
    "        vec2,lab2,_,_,_ = dataset_ep_test.__getitem__(jdx)\n",
    "\n",
    "        check1 = vector_to_checkpoint(vector=vec1, reference_checkpoint=check)\n",
    "        check2 = vector_to_checkpoint(vector=vec2, reference_checkpoint=check)\n",
    "\n",
    "#         print('load first checkpoint')\n",
    "        model.model.load_state_dict(check1)\n",
    "#         print('compute predictions of first checkpoint')\n",
    "        with torch.no_grad():\n",
    "            pred1 = model(test_data)\n",
    "\n",
    "#         print('load permuted checkpoint')\n",
    "        model.model.load_state_dict(check2)\n",
    "#         print('compute predictions of permuted checkpoint')\n",
    "        with torch.no_grad():\n",
    "            pred2 = model(test_data)\n",
    "\n",
    "#         print('check equality')\n",
    "        restmp = torch.allclose(pred1,pred2)\n",
    "        if not restmp:\n",
    "            print(torch.allclose(pred1,pred2,atol=1e-05, rtol=1e-03))\n",
    "        res.append(restmp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3819829f",
   "metadata": {},
   "source": [
    "No errors, all close enough with relaxed tolerances. Permutations still work "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f20f47d2",
   "metadata": {},
   "source": [
    "### compare recorded to computed accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "85a9ecbd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                                                                                                      | 0/10 [00:00<?, ?it/s]INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.180163366317749; accuracy: 0.213\n",
      "INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.180163356781006; accuracy: 0.213\n",
      " 10%|███████████████▊                                                                                                                                              | 1/10 [01:25<12:50, 85.60s/it]INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 1.1837517328262328; accuracy: 0.614\n",
      "INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 1.1837517595291138; accuracy: 0.614\n",
      " 20%|███████████████████████████████▌                                                                                                                              | 2/10 [02:29<09:44, 73.01s/it]INFO:root:validate at epoch 0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc_1: 0.614 | acc_2: 0.614 | acc_rec: 0.6262676705593117\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:test ::: loss: 2.2183909492492675; accuracy: 0.197\n",
      "INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.2183907890319823; accuracy: 0.197\n",
      " 30%|███████████████████████████████████████████████▍                                                                                                              | 3/10 [03:47<08:46, 75.24s/it]INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.0843050088882444; accuracy: 0.276\n",
      "INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.0843049564361573; accuracy: 0.276\n",
      " 40%|███████████████████████████████████████████████████████████████▏                                                                                              | 4/10 [05:23<08:20, 83.48s/it]INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.2148335647583006; accuracy: 0.198\n",
      "INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.2148336029052733; accuracy: 0.198\n",
      " 50%|███████████████████████████████████████████████████████████████████████████████                                                                               | 5/10 [06:21<06:11, 74.33s/it]INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.1776208724975588; accuracy: 0.223\n",
      "INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.177620843887329; accuracy: 0.223\n",
      " 60%|██████████████████████████████████████████████████████████████████████████████████████████████▊                                                               | 6/10 [07:26<04:43, 70.98s/it]INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 1.1041694135665894; accuracy: 0.653\n",
      "INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 1.1041693668365478; accuracy: 0.653\n",
      " 70%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████▌                                               | 7/10 [09:09<04:04, 81.48s/it]INFO:root:validate at epoch 0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acc_1: 0.653 | acc_2: 0.653 | acc_rec: 0.6694068838352797\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:test ::: loss: 2.191834867477417; accuracy: 0.219\n",
      "INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.1918348865509034; accuracy: 0.219\n",
      " 80%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▍                               | 8/10 [10:24<02:38, 79.29s/it]INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 1.1670266799926758; accuracy: 0.649\n",
      "INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 1.167026635169983; accuracy: 0.649\n",
      " 90%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏               | 9/10 [11:24<01:13, 73.54s/it]INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.109017961502075; accuracy: 0.292\n",
      "INFO:root:validate at epoch 0\n",
      "INFO:root:test ::: loss: 2.1090178623199463; accuracy: 0.292\n",
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [12:52<00:00, 77.21s/it]\n"
     ]
    }
   ],
   "source": [
    "import tqdm\n",
    "import random\n",
    "from shrp.datasets.dataset_simclr import vector_to_checkpoint\n",
    "\n",
    "repetitions = 10\n",
    "\n",
    "res = []\n",
    "\n",
    "check = dataset_ep_test.reference_checkpoint\n",
    "for idx in tqdm.tqdm(range(repetitions)):\n",
    "        \n",
    "    kdx = random.randint(0,len(dataset_ep_test)-1)\n",
    "    mdx,edx = dataset_ep_test._index[kdx]\n",
    "    vec1,lab1,vec2,_,_ = dataset_ep_test.__getitem__(kdx)\n",
    "\n",
    "    check1 = vector_to_checkpoint(vector=vec1, reference_checkpoint=check)\n",
    "    check2 = vector_to_checkpoint(vector=vec2, reference_checkpoint=check)\n",
    "\n",
    "#         print('load first checkpoint')\n",
    "    model.model.load_state_dict(check1)\n",
    "#         print('compute predictions of first checkpoint')\n",
    "    with torch.no_grad():\n",
    "        loss_1,acc_1 = model.test_epoch(testloader,epoch=0)\n",
    "\n",
    "#         print('load permuted checkpoint')\n",
    "    model.model.load_state_dict(check2)\n",
    "#         print('compute predictions of permuted checkpoint')\n",
    "    with torch.no_grad():\n",
    "        loss_2,acc_2 = model.test_epoch(testloader,epoch=0)\n",
    "\n",
    "#         print('check equality')\n",
    "\n",
    "    acc_recorded = dataset_ep_test.properties['test_acc'][mdx][edx]\n",
    "    restmp_pair = torch.allclose(torch.tensor(acc_1),torch.tensor(acc_2),atol=1e-02, rtol=1e-03)\n",
    "    restmp_base = torch.allclose(torch.tensor(acc_1),torch.tensor(acc_recorded),atol=1e-02, rtol=1e-03)\n",
    "    if not (restmp_pair and restmp_base):\n",
    "        print(f'acc_1: {acc_1} | acc_2: {acc_2} | acc_rec: {acc_recorded}')\n",
    "    res.append(restmp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "aafefb6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "t1 = torch.randn(3)\n",
    "t2 = torch.randn(5)\n",
    "t3 = torch.randn(4)\n",
    "\n",
    "t = [t1,t2,t3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "8d7268b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([0.3595, 0.7960, 0.8118]),\n",
       " tensor([0.8026, 0.0560, 0.8345, 1.0049, 1.4833]),\n",
       " tensor([-0.0634,  2.2114, -0.4224, -0.2023])]"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "b9de9b77",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 0.],\n",
       "        [0., 0., 0., 0., 0.]])"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = torch.zeros(len(t),5)\n",
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "ccf1bd11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "\n",
    "def tokenize(t:List[torch.tensor],tokensize:int)->torch.tensor:\n",
    "    \"\"\"\n",
    "    transforms list of tokens of differen lenght to tensor\n",
    "    input t: List[torch.tensor]: list of 1d input tokens of different lenghts\n",
    "    input tokensize: int output dimension of each token\n",
    "    return out: torch.tensor\n",
    "    \"\"\"\n",
    "    # init output with zeros\n",
    "    tokens = torch.zeros(len(t),tokensize)\n",
    "    mask = torch.zeros(len(t),tokensize)\n",
    "    # iterate over inputs\n",
    "    for idx,tdx in enumerate(t):\n",
    "        # get end of token, either the length of the input or them maximum length\n",
    "        tdx_end = min(tdx.shape[0],tokensize)\n",
    "        # put at position idx\n",
    "        tokens[idx,:tdx_end] = tdx[:tdx_end]\n",
    "        mask[idx,:tdx_end] = torch.ones(tdx_end)\n",
    "        \n",
    "    # return\n",
    "    return tokens,mask\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "75a0d4df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.3595,  0.7960,  0.8118,  0.0000,  0.0000,  0.0000],\n",
      "        [ 0.8026,  0.0560,  0.8345,  1.0049,  1.4833,  0.0000],\n",
      "        [-0.0634,  2.2114, -0.4224, -0.2023,  0.0000,  0.0000]])\n",
      "tensor([[1., 1., 1., 0., 0., 0.],\n",
      "        [1., 1., 1., 1., 1., 0.],\n",
      "        [1., 1., 1., 1., 0., 0.]])\n"
     ]
    }
   ],
   "source": [
    "tt,mt = tokenize(t,tokensize=6)\n",
    "print(tt)\n",
    "print(mt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "3d1600fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.3595,  0.7960,  0.8118,  0.0000,  0.0000],\n",
       "        [ 0.8026,  0.0560,  0.8345,  1.0049,  1.4833],\n",
       "        [-0.0634,  2.2114, -0.4224, -0.2023,  0.0000]])"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e42fc0e8",
   "metadata": {},
   "source": [
    "# test position embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1dc4ee24",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "testloader = torch.utils.data.DataLoader(dataset_ep_test,batch_size=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8d6d2f31",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([5, 12, 201])\n",
      "torch.Size([5, 12, 201])\n",
      "('/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00321_321_seed=322_2021-07-27_06-12-40#_#epoch_1', '/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00321_321_seed=322_2021-07-27_06-12-40#_#epoch_5', '/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00321_321_seed=322_2021-07-27_06-12-40#_#epoch_25', '/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00689_689_seed=690_2021-07-27_21-54-40#_#epoch_1', '/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00689_689_seed=690_2021-07-27_21-54-40#_#epoch_5')\n",
      "torch.Size([5, 12, 201])\n",
      "torch.Size([5, 12, 201])\n",
      "('/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00321_321_seed=322_2021-07-27_06-12-40#_#epoch_1#_#canon', '/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00321_321_seed=322_2021-07-27_06-12-40#_#epoch_5#_#canon', '/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00321_321_seed=322_2021-07-27_06-12-40#_#epoch_25#_#canon', '/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00689_689_seed=690_2021-07-27_21-54-40#_#epoch_1#_#canon', '/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00689_689_seed=690_2021-07-27_21-54-40#_#epoch_5#_#canon')\n",
      "torch.Size([5, 12, 2])\n"
     ]
    }
   ],
   "source": [
    "w1,m1,l1,w2,m2,l2,p = next(iter(testloader))\n",
    "print(w1.shape)\n",
    "print(m1.shape)\n",
    "print(l1)\n",
    "print(w2.shape)\n",
    "print(m2.shape)\n",
    "print(l2)\n",
    "print(p.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "c9cf2d1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "class PositionEmbs(nn.Module):\n",
    "    \"\"\"Adds learned positional embeddings to the inputs.\n",
    "    Attributes:\n",
    "        posemb_init: positional embedding initializer.\n",
    "        max_positions: maximum number of positions to embed.\n",
    "        embedding_dim: dimension of the input embeddings.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, max_positions=[48,256], embedding_dim=128):\n",
    "        super().__init__()\n",
    "        self.max_positions = max_positions\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.pe1 = nn.Embedding(max_positions[0], embedding_dim//2)\n",
    "        self.pe2 = nn.Embedding(max_positions[1], embedding_dim//2)\n",
    "\n",
    "    def forward(self, inputs, pos):\n",
    "        \"\"\"Applies the AddPositionEmbs module.\n",
    "        Args:\n",
    "            inputs: Inputs to the layer, shape `(batch_size, seq_len, emb_dim)`.\n",
    "            pos: Position of the first token in each sequence, shape `(batch_size,seq_len,2)`.\n",
    "        Returns:\n",
    "            Output tensor with shape `(batch_size, seq_len, emb_dim + 2)`.\n",
    "        \"\"\"\n",
    "        assert (\n",
    "            inputs.ndim == 3\n",
    "        ), f\"Number of dimensions should be 3, but it is {inputs.ndim}\"\n",
    "        assert (\n",
    "            pos.shape[2] == 2\n",
    "        ), \"Position tensors should have two dimensions (layer, module in layer)\"\n",
    "        assert (\n",
    "            pos.shape[0] == inputs.shape[0]\n",
    "        ), \"Position tensors should have the same batch size as inputs\"\n",
    "        assert (\n",
    "            pos.shape[1] == inputs.shape[1]\n",
    "        ), \"Position tensors should have the same seq length as inputs\"\n",
    "\n",
    "        # TODO: fix embeding code\n",
    "        pos_emb1 = self.pe1(pos[:,:,0])\n",
    "        pos_emb2 = self.pe2(pos[:,:,1])\n",
    "        pos_emb = torch.cat([pos_emb1,pos_emb2],dim=2)\n",
    "        \n",
    "        out = torch.cat([inputs,pos_emb],dim=2)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "e011ead5",
   "metadata": {},
   "outputs": [],
   "source": [
    "PosEmbs = PositionEmbs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "0901f54f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([5, 12, 2])\n",
      "torch.Size([5, 12, 329])\n"
     ]
    }
   ],
   "source": [
    "print(p.shape)\n",
    "pse = PosEmbs(w1,p)\n",
    "print(pse.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6eb4ccf",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "03310363",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "34e7beca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1, 2],\n",
      "        [3, 4]])\n",
      "tensor([[5, 6],\n",
      "        [7, 7]])\n",
      "tensor([[ 5, 12],\n",
      "        [21, 28]])\n"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "0b475669",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "yes\n"
     ]
    }
   ],
   "source": [
    "for idx in range(5):\n",
    "    if idx==3:\n",
    "        break\n",
    "print('yes')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "7c1d058d",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = [[None,1,2,4],[3,None,1],[None,None]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "615e99f6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[1, 2, 4], [3, 1]]"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = [[ddx for ddx in data[idx] if ddx] for idx in range(len(data))]\n",
    "data = [ddx for ddx in data if len(ddx)>0]\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "cf042d31",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = Path('/netscratch2/kschuerholt/code/shrp/experiments/02_representation_learning/01_test/tune/ae_resnet_test_subset/dataset.pt')\n",
    "dataset = torch.load(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "46ad1451",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "15\n",
      "3\n",
      "9610\n",
      "27\n"
     ]
    }
   ],
   "source": [
    "testset = dataset['testset']\n",
    "print(len(testset.data))\n",
    "print(len(testset.data[0]))\n",
    "print(len(testset.data[0][0]))\n",
    "print(len(testset.data[0][0][0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "bb068d76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4608"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "testset.tokensize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1383b60e",
   "metadata": {},
   "source": [
    "# benchmark dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dca42577",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:torch.distributed.nn.jit.instantiator:Created a temporary directory at /tmp/tmpljm_5ej0\n",
      "INFO:torch.distributed.nn.jit.instantiator:Writing /tmp/tmpljm_5ej0/_remote_module_non_scriptable.py\n"
     ]
    }
   ],
   "source": [
    "#\n",
    "import logging\n",
    "\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "\n",
    "from pathlib import Path\n",
    "import json\n",
    "from shrp.datasets.dataset_simclr import SimCLRDataset\n",
    "\n",
    "from shrp.git_re_basin.git_re_basin import (\n",
    "    PermutationSpec,\n",
    "    zoo_cnn_permutation_spec,\n",
    "    weight_matching,\n",
    "    apply_permutation,\n",
    ")\n",
    "\n",
    "from shrp.models.def_AE_module import AEModule\n",
    "\n",
    "# from lightning.fabric import Fabric\n",
    "# from lightning.fabric.strategies import SingleDeviceStrategy\n",
    "\n",
    "import torch\n",
    "\n",
    "# from torch.profiler import profile, record_function, ProfilerActivity\n",
    "\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "991cbf4e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a230afdb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0186d788",
   "metadata": {},
   "outputs": [],
   "source": [
    "# without permutations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "37edcec2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:reference checkpoint found at /netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/NN_tune_trainable_97ebe_00321_321_seed=322_2021-07-27_06-12-40\n",
      "2023-04-12 03:54:18,053\tINFO worker.py:1553 -- Started a local Ray instance.\n",
      "INFO:root:loading checkpoints from [PosixPath('/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform')]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2800/2800 [00:08<00:00, 347.48it/s]\n",
      "INFO:root:Data loaded. found 700 usable samples out of potential 2800 samples.\n",
      "INFO:root:Load properties for samples from paths.\n",
      "INFO:root:### load data for dict_keys(['test_acc', 'training_iteration', 'ggap'])\n",
      "700it [00:12, 57.73it/s]\n",
      "INFO:root:Properties loaded.\n",
      "INFO:root:both view 1 and view 2 are set to canonical. number of permutations is set to 0\n",
      "INFO:root:prepare canonical form\n",
      "2023-04-12 03:54:54,356\tINFO worker.py:1553 -- Started a local Ray instance.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "preparing computing canon form...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 700/700 [00:02<00:00, 251.84it/s]\n",
      "INFO:root:init dataset length\n",
      "INFO:root:Get Positions\n",
      "INFO:root:vectorize data\n",
      "INFO:root:Get layer mapping\n",
      "INFO:root:Get layer-wise mean and std\n",
      "INFO:root:Apply standardization\n",
      "INFO:root:Discover tokensize\n"
     ]
    }
   ],
   "source": [
    "config_key_list = []\n",
    "result_key_list = [\n",
    "    \"test_acc\",\n",
    "    \"training_iteration\",\n",
    "    \"ggap\",\n",
    "    #     \"sparsity_ratio\",\n",
    "]\n",
    "property_keys = {\n",
    "    \"result_keys\": result_key_list,\n",
    "    \"config_keys\": config_key_list,\n",
    "}\n",
    "\n",
    "path_root = Path(\"/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/\")\n",
    "\n",
    "dataset_train = SimCLRDataset(\n",
    "    root=path_root,\n",
    "    epoch_lst=[1, 5, 20, 25],\n",
    "#     epoch_lst=list(range(1, 26)),\n",
    "    #     mode=\"checkpoint\",\n",
    "    mode=\"vector\",\n",
    "    permutations_number=0,\n",
    "    permutation_spec=zoo_cnn_permutation_spec(),\n",
    "    #     view_1_canonical = False,\n",
    "    view_1_canonical=True,\n",
    "    view_2_canonical=True,\n",
    "    #     view_2_canonical = False,\n",
    "    add_noise_view_1=0.1,  # [('input', 0.15), ('output', 0.013)]\n",
    "    add_noise_view_2=0.0,  # [('input', 0.15), ('output', 0.013)]\n",
    "    noise_multiplicative=True,\n",
    "    erase_augment=None,  # {\"p\": 0.5,\"scale\":(0.02,0.33),\"value\":0,\"mode\":\"block\"}\n",
    "    windowsize=12,\n",
    "    standardize=True,\n",
    "    train_val_test=\"train\",  # determines whcih dataset split to use\n",
    "    ds_split=[0.7, 0.15, 0.15],  #\n",
    "    max_samples=1000,\n",
    "    # weight_threshold=float(\"inf\"),\n",
    "    weight_threshold=15,\n",
    "    filter_function=None,  # gets sample path as argument and returns True if model needs to be filtered out\n",
    "    property_keys=property_keys,\n",
    "    num_threads=6,\n",
    "    shuffle_path=True,\n",
    "    verbosity=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88570ee9",
   "metadata": {},
   "source": [
    "### single dataset access"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ab21a4dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_int = len(dataset_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "481a632c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "331 ms ± 7.15 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "for idx in range(64):\n",
    "    kdx = random.randint(0, max_int)\n",
    "    (x_i, m_i, x_j, m_j, p) = dataset_train.__getitem__(kdx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a970162e",
   "metadata": {},
   "source": [
    "fetching one batch of 64 samples takes ~400ms (~0.4s), sequentially. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8564c9e2",
   "metadata": {},
   "source": [
    "### dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "20f5cf97",
   "metadata": {},
   "outputs": [],
   "source": [
    "dloader = torch.utils.data.DataLoader(\n",
    "    dataset_train,\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    "    prefetch_factor=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9661590f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.06 s ± 233 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "for idx, (x_i, m_i, x_j, m_j, p) in enumerate(dloader):\n",
    "    continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "48e65fee",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "43"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "77ad08cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.1755813953488372"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "7.55/len(dloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0d876851",
   "metadata": {},
   "outputs": [],
   "source": [
    "dloader = torch.utils.data.DataLoader(\n",
    "    dataset_train,\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    num_workers=8,\n",
    "    pin_memory=True,\n",
    "    prefetch_factor=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "464f978f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9.72 s ± 246 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "for idx, (x_i, m_i, l_i, x_j, m_j, _, p) in enumerate(dloader):\n",
    "    continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f6c1bfb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "dloader = torch.utils.data.DataLoader(\n",
    "    dataset_train,\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    "    prefetch_factor=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "3b2a6099",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.68 s ± 764 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "for idx, (x_i, m_i, l_i, x_j, m_j, _, p) in enumerate(dloader):\n",
    "    continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "02a3ef33",
   "metadata": {},
   "outputs": [],
   "source": [
    "dloader = torch.utils.data.DataLoader(\n",
    "    dataset_train,\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    num_workers=2,\n",
    "    pin_memory=True,\n",
    "    prefetch_factor=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b9a4c432",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11 s ± 137 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "for idx, (x_i, m_i, l_i, x_j, m_j, _, p) in enumerate(dloader):\n",
    "    continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "8aa5ccd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "dloader = torch.utils.data.DataLoader(\n",
    "    dataset_train,\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    num_workers=2,\n",
    "    pin_memory=True,\n",
    "    prefetch_factor=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "5423b59a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11.4 s ± 315 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "for idx, (x_i, m_i, l_i, x_j, m_j, _, p) in enumerate(dloader):\n",
    "    continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "98f7ff04",
   "metadata": {},
   "outputs": [],
   "source": [
    "dloader = torch.utils.data.DataLoader(\n",
    "    dataset_train,\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    num_workers=3,\n",
    "    pin_memory=True,\n",
    "    prefetch_factor=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "62c6beaf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8.64 s ± 365 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "for idx, (x_i, m_i, l_i, x_j, m_j, _, p) in enumerate(dloader):\n",
    "    continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62502cad",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c7f6d97",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "995f5a84",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "195681c8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7bea768",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "202128b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with permutations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d7c8659",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fb5a60b3",
   "metadata": {},
   "source": [
    "# baseline: random data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "2defcd7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "class RandomDataset(Dataset):\n",
    "    \"\"\"\n",
    "    This dataset produces random tensors of data\n",
    "    \"\"\"\n",
    "\n",
    "    # init\n",
    "    def __init__(\n",
    "        self,\n",
    "        windowsize=12,\n",
    "        tokensize=201,\n",
    "        samples=1000,\n",
    "        ):\n",
    "        super(RandomDataset, self).__init__()\n",
    "        self.windowsize=windowsize\n",
    "        self.tokensize=tokensize\n",
    "        self.samples = samples\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.samples\n",
    "    \n",
    "    def __getitem__(self,idx):\n",
    "        x_i = torch.randn(self.windowsize,self.tokensize)\n",
    "        m_i = torch.randn(self.windowsize,self.tokensize)\n",
    "        l_i = \"123\"\n",
    "        x_j = torch.randn(self.windowsize,self.tokensize)\n",
    "        m_j = torch.randn(self.windowsize,self.tokensize)\n",
    "        l_j = \"123\"\n",
    "        p = [(1,6) for _ in range(self.windowsize)]\n",
    "        return x_i, m_i, l_i, x_j, m_j, l_j, p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "6ad7a35d",
   "metadata": {},
   "outputs": [],
   "source": [
    "rdataset = RandomDataset(\n",
    "    windowsize=dataset_train.window,\n",
    "    tokensize=dataset_train.tokensize,\n",
    "    samples = dataset_train.__len__(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "f59193e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.24 ms ± 140 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "for idx in range(64):\n",
    "    kdx = random.randint(0, max_int)\n",
    "    (x_i, m_i, l_i, x_j, m_j, _, p) = rdataset.__getitem__(kdx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "373918ce",
   "metadata": {},
   "source": [
    "random data access is about 1% of original dataset class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "dcb6d58f",
   "metadata": {},
   "outputs": [],
   "source": [
    "dloader = torch.utils.data.DataLoader(\n",
    "    rdataset,\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    num_workers=2,\n",
    "    pin_memory=True,\n",
    "    prefetch_factor=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "22be3604",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.48 s ± 169 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "for idx, (x_i, m_i, l_i, x_j, m_j, _, p) in enumerate(dloader):\n",
    "    continue"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc5057bb",
   "metadata": {},
   "source": [
    "# test ffcv dataset type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "fdc9ba2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ffcv.writer import DatasetWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a7ceb26",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "#     RGBImageField: Handles images including (optional) compression and resizing. Pass in a PyTorch Tensor.\n",
    "\n",
    "#     IntField and FloatField: Handle simple scalar fields. Pass in int or float.\n",
    "\n",
    "#     BytesField: Stores byte arrays of variable length. Pass in numpy byte array.\n",
    "\n",
    "#     JSONField: Encodes a JSON document. Pass in dict that can be JSON-encoded.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3f8ba4b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.float32\n",
      "torch.float32\n",
      "torch.float32\n",
      "torch.float32\n",
      "torch.int64\n"
     ]
    }
   ],
   "source": [
    "(x_i, m_i, x_j, m_j, p) = next(iter(dloader))\n",
    "print(x_i.dtype)\n",
    "print(m_i.dtype)\n",
    "print(x_j.dtype)\n",
    "print(m_j.dtype)\n",
    "print(p.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d21995f0",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "cannot create 'torch.dtype' instances",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [20], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdtype\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mfloat32\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mTypeError\u001b[0m: cannot create 'torch.dtype' instances"
     ]
    }
   ],
   "source": [
    "torch.dtype('float32')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "579119b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ffcv.fields import NDArrayField, FloatField, TorchTensorField\n",
    "#  class ffcv.fields.TorchTensorField(dtype: torch.dtype, shape: Tuple[int, ...])[source]\n",
    "# Each field corresponds to an element of the data tuple returned by our dataset, \n",
    "# and specifies how the element should be written to (and later, read from) the FFCV dataset file. \n",
    "# In our case, the dataset has two fields, one for the (vector) input and the other for the corresponding \n",
    "# (scalar) label. Both of these fields already have default implementations in FFCV, which we use below:\n",
    "# (x_i, m_i, x_j, m_j, p)\n",
    "w= dataset_train.window\n",
    "t = dataset_train.tokensize\n",
    "write_path = Path('./ffcv_test_dataset.beton')\n",
    "writer = DatasetWriter(write_path, {\n",
    "    'x_i': TorchTensorField(shape=(w,t), dtype=torch.float32),\n",
    "    'm_i': TorchTensorField(shape=(w,t), dtype=torch.float32),\n",
    "    'x_j': TorchTensorField(shape=(w,t), dtype=torch.float32),\n",
    "    'm_j': TorchTensorField(shape=(w,t), dtype=torch.float32),\n",
    "    'p': TorchTensorField(shape=(w,2), dtype=torch.int64),\n",
    "\n",
    "\n",
    "}, num_workers=16)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e9e9e4ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2800/2800 [00:03<00:00, 735.12it/s]\n"
     ]
    }
   ],
   "source": [
    "writer.from_indexed_dataset(dataset_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae197bbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "3ef64d6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ffcv.loader import Loader, OrderOption\n",
    "from ffcv.fields.decoders import NDArrayDecoder, FloatDecoder\n",
    "\n",
    "# Our first step is instantiating the Loader class:\n",
    "\n",
    "BATCH_SIZE = 64\n",
    "NUM_WORKERS = 4\n",
    "\n",
    "# The order option in the loader initialization is similar to PyTorch DataLoader’s shuffle option, with some additional options. This argument takes an enum provided by ffcv.loader.OrderOption:\n",
    "\n",
    "from ffcv.loader import OrderOption\n",
    "\n",
    "# Truly random shuffling (shuffle=True in PyTorch)\n",
    "# ORDERING = OrderOption.RANDOM\n",
    "\n",
    "# Unshuffled (i.e., served in the order the dataset was written)\n",
    "# ORDERING = OrderOption.SEQUENTIAL\n",
    "\n",
    "# Memory-efficient but not truly random loading\n",
    "# Speeds up loading over RANDOM when the whole dataset does not fit in RAM!\n",
    "ORDERING = OrderOption.QUASI_RANDOM\n",
    "\n",
    "\n",
    "# In order to create a loader, we need to specify a path to the FFCV dataset, batch size, number of workers, as well as two less standard arguments, order and pipelines, which we discuss below:\n",
    "# Dataset ordering\n",
    "loader = Loader(write_path,\n",
    "                batch_size=BATCH_SIZE,\n",
    "                num_workers=NUM_WORKERS,\n",
    "                order=ORDERING,\n",
    "#                 pipelines=PIPELINES\n",
    "               )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "49c78b4c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "72.3 ms ± 1.77 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "for idx, (x_i, m_i, x_j, m_j, p) in enumerate(loader):\n",
    "    continue"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7045c87",
   "metadata": {},
   "source": [
    "gave up, opencv installation is a pain in the ... "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a64d5ed8",
   "metadata": {},
   "source": [
    "# improvement: tensorize data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc11d817",
   "metadata": {},
   "outputs": [],
   "source": [
    "# idee:\n",
    "# call tokenize at init\n",
    "# -> list of samples [n_samples,n_epochs]\n",
    "# -> each samlpe is [n_tokens,tokensize] (not list of tensors)\n",
    "# -> applying all operations on that one tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "072fec30",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "from pathlib import Path\n",
    "import random\n",
    "import copy\n",
    "\n",
    "import itertools\n",
    "from math import factorial\n",
    "\n",
    "from shrp.datasets.dataset_epochs import ModelDatasetBaseEpochs\n",
    "from shrp.git_re_basin.git_re_basin import (\n",
    "    PermutationSpec,\n",
    "    zoo_cnn_permutation_spec,\n",
    "    weight_matching,\n",
    "    apply_permutation,\n",
    ")\n",
    "\n",
    "import logging\n",
    "\n",
    "from typing import List\n",
    "\n",
    "from shrp.datasets.random_erasing import RandomErasingVector\n",
    "\n",
    "import ray\n",
    "from shrp.datasets.progress_bar import ProgressBar\n",
    "\n",
    "#####################################################################\n",
    "# Define Dataset class\n",
    "#####################################################################\n",
    "class SimCLRDatasetTensor(ModelDatasetBaseEpochs):\n",
    "    \"\"\"\n",
    "    This class inherits from the base ModelDatasetBaseEpochs class.\n",
    "    It extends it by permutations of the dataset in the init function.\n",
    "    \"\"\"\n",
    "\n",
    "    # init\n",
    "    def __init__(\n",
    "        self,\n",
    "        root,\n",
    "        epoch_lst=10,\n",
    "        mode=\"vector\",  # \"vector\", \"checkpoint\"\n",
    "        permutations_number=10,\n",
    "        permutation_spec: PermutationSpec = zoo_cnn_permutation_spec,\n",
    "        view_1_canonical: bool = False,\n",
    "        view_2_canonical: bool = False,\n",
    "        add_noise_view_1: float = 0.0,  # [('input', 0.15), ('output', 0.013)]\n",
    "        add_noise_view_2: float = 0.0,  # [('input', 0.15), ('output', 0.013)]\n",
    "        noise_multiplicative: bool = True,\n",
    "        erase_augment=None,  # {\"p\": 0.5,\"scale\":(0.02,0.33),\"value\":0,\"mode\":\"block\"}\n",
    "        windowsize: int = 5,\n",
    "        standardize: bool = True,  # wether or not to standardize the data\n",
    "        tokensize: int = 0,\n",
    "        train_val_test=\"train\",\n",
    "        ds_split=[0.7, 0.3],\n",
    "        weight_threshold: float = float(\"inf\"),\n",
    "        max_samples: int = 0,  # limit the number of models to integer number (full model trajectory, all epochs)\n",
    "        filter_function=None,  # gets sample path as argument and returns True if model needs to be filtered out\n",
    "        property_keys=None,\n",
    "        shuffle_path: bool = True,\n",
    "        num_threads=4,\n",
    "        verbosity=0,\n",
    "    ):\n",
    "        # call init of base class\n",
    "        super().__init__(\n",
    "            root=root,\n",
    "            epoch_lst=epoch_lst,\n",
    "            mode=\"checkpoint\",\n",
    "            train_val_test=train_val_test,\n",
    "            ds_split=ds_split,\n",
    "            weight_threshold=weight_threshold,\n",
    "            max_samples=max_samples,\n",
    "            filter_function=filter_function,\n",
    "            property_keys=property_keys,\n",
    "            num_threads=num_threads,\n",
    "            verbosity=verbosity,\n",
    "            shuffle_path=shuffle_path,\n",
    "        )\n",
    "        self.mode = mode\n",
    "        self.permutations_number = permutations_number\n",
    "        self.permutation_spec = permutation_spec\n",
    "        self.standardize = standardize\n",
    "        self.tokensize = tokensize\n",
    "\n",
    "        self.add_noise_view_1 = add_noise_view_1\n",
    "        assert isinstance(self.add_noise_view_1, float)\n",
    "        self.add_noise_view_2 = add_noise_view_2\n",
    "        assert isinstance(self.add_noise_view_1, float)\n",
    "        self.use_multiplicative_noise = noise_multiplicative\n",
    "\n",
    "        self.num_threads = num_threads\n",
    "\n",
    "        self.view_1_canonical = view_1_canonical\n",
    "        self.view_2_canonical = view_2_canonical\n",
    "        if view_1_canonical and view_2_canonical:\n",
    "            logging.info(\n",
    "                f\"both view 1 and view 2 are set to canonical. number of permutations is set to 0\"\n",
    "            )\n",
    "            self.permutations_number = 0\n",
    "\n",
    "        # set erase augmnet\n",
    "        self.set_erase(erase_augment)\n",
    "\n",
    "        if self.view_1_canonical or self.view_2_canonical:\n",
    "            logging.info(\"prepare canonical form\")\n",
    "            # TODO: load reference checkpoint in base class, share between train/val/test\n",
    "            self.map_models_to_canonical()\n",
    "\n",
    "        ### init len ###########################################################################################\n",
    "        logging.info(\"init dataset length\")\n",
    "        self.init_len()\n",
    "\n",
    "        ### get positions ###########################################################################################\n",
    "        logging.info(\"Get Positions\")\n",
    "        reference_checkpoint = self.reference_checkpoint\n",
    "        self.positions = get_position_mapping_from_checkpoint(reference_checkpoint)\n",
    "\n",
    "        ### vectorize data ###########################################################################################\n",
    "        if self.mode == \"vector\":\n",
    "            logging.info(\"vectorize data\")\n",
    "            # keep reference checkpoint\n",
    "            self.vectorize_data()\n",
    "\n",
    "        ### initialize permutations ##########################################################################################################################################\n",
    "        # list of permutations (list of list with indexes)\n",
    "        if self.permutations_number > 0:\n",
    "            logging.info(\"init permutations\")\n",
    "            self.precompute_permutations(\n",
    "                permutation_number=self.permutations_number,\n",
    "                perm_spec=self.permutation_spec,\n",
    "                num_threads=num_threads,\n",
    "            )\n",
    "\n",
    "        if self.standardize:\n",
    "            self.standardize_data()\n",
    "\n",
    "        ### set module window ##################################################################################################################################################################\n",
    "        self.set_module_window(windowsize=windowsize)\n",
    "\n",
    "        ### set tokensize ###################################################################################################################################################################\n",
    "        logging.info('Tokenize data')\n",
    "        self.set_token_size(tokensize=self.tokensize)\n",
    "        \n",
    "        self.tokenize_data()        \n",
    "\n",
    "    def set_module_window(self, windowsize: int = 5):\n",
    "        # check that window is within range 1,no-tokens\n",
    "        assert 0 < windowsize <= len(self.data[0][-1])\n",
    "        self.window = windowsize\n",
    "\n",
    "    def set_token_size(self, tokensize: int = 0):\n",
    "        \"\"\"\n",
    "        tokens are zero-padded to all have the same size at __getitem__\n",
    "        this function sets the size of the token either to a specific length (rest is cut off)\n",
    "        or discovers the maximum size\n",
    "        \"\"\"\n",
    "        if tokensize > 0:\n",
    "            self.tokensize = tokensize\n",
    "        else:\n",
    "            logging.info(\"Discover tokensize\")\n",
    "            # assumes data is already vectorized\n",
    "            max_len = 0\n",
    "            for tdx in self.data[0][-1]:\n",
    "                if tdx.shape[0] > max_len:\n",
    "                    max_len = tdx.shape[0]\n",
    "            self.tokensize = max_len\n",
    "\n",
    "            \n",
    "    def tokenize_data(self):\n",
    "        \"\"\"\n",
    "        cast samples as list of tokens to tensors to speed up processing\n",
    "        \"\"\"\n",
    "        # iterate over all samlpes\n",
    "        for idx in range(len(self.data)):\n",
    "            for jdx in range(len(self.data[idx])):\n",
    "                self.data[idx][jdx], mask = tokenize(\n",
    "                    self.data[idx][jdx],\n",
    "                    tokensize=self.tokensize,\n",
    "                    return_mask=True,\n",
    "                )\n",
    "        self.mask = mask\n",
    "        \n",
    "    ## get_weights ####################################################################################################################################################################\n",
    "    def __get_weights__(\n",
    "        self,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Returns:\n",
    "            torch.Tensor with full dataset as sequence of components [n_samples,n_tokens_per_sample,token_dim]\n",
    "        \"\"\"\n",
    "        #todo -> probably needs flattening or something\n",
    "        if not self.mode == \"vector\":\n",
    "            data_tmp = copy.deepcopy(self.data)\n",
    "            self.vectorize_data()\n",
    "            data_out = [\n",
    "                tokenize(\n",
    "                    self.data[idx][jdx],\n",
    "                    tokensize=self.tokensize,\n",
    "                    return_mask=True,\n",
    "                )[0]\n",
    "                for idx in range(len(self.data))\n",
    "                for jdx in range(len(self.data[idx]))\n",
    "            ]\n",
    "            mask_out = [\n",
    "                tokenize(\n",
    "                    self.data[idx][jdx],\n",
    "                    tokensize=self.tokensize,\n",
    "                    return_mask=True,\n",
    "                )[1]\n",
    "                for idx in range(len(self.data))\n",
    "                for jdx in range(len(self.data[idx]))\n",
    "            ]\n",
    "            data_out = torch.stack(data_out)\n",
    "            mask_out = torch.stack(mask_out)\n",
    "            self.data = data_tmp\n",
    "            return data_out, mask_out\n",
    "        data_out = [\n",
    "            tokenize(\n",
    "                self.data[idx][jdx],\n",
    "                tokensize=self.tokensize,\n",
    "                return_mask=True,\n",
    "            )[0]\n",
    "            for idx in range(len(self.data))\n",
    "            for jdx in range(len(self.data[idx]))\n",
    "        ]\n",
    "        mask_out = [\n",
    "            tokenize(\n",
    "                self.data[idx][jdx],\n",
    "                tokensize=self.tokensize,\n",
    "                return_mask=True,\n",
    "            )[1]\n",
    "            for idx in range(len(self.data))\n",
    "            for jdx in range(len(self.data[idx]))\n",
    "        ]\n",
    "        data_out = torch.stack(data_out)\n",
    "        mask_out = torch.stack(mask_out)\n",
    "        logging.debug(f\"shape of weight tensor: {data_out.shape}\")\n",
    "        return data_out, mask_out\n",
    "\n",
    "    ## getitem ####################################################################################################################################################################\n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            index (int): Index of the sample to be retrieved\n",
    "        Returns:\n",
    "            ddx_idx: torch.Tensor of neuron tokens with shape [n_tokens_per_sample/windowsize,token_dim] in view 1\n",
    "            mask_idx: torch.Tensor of the same shape as ddx_idx indicating the nonzero elements\n",
    "            label_idx: label for sample in view 1\n",
    "            ddx_jdx: torch.Tensor of neuron tokens with shape [n_tokens_per_sample/windowsize,token_dim] in view 2\n",
    "            mask_jdx: torch.Tensor of the same shape as ddx_jdx indicating the nonzero elements\n",
    "            label_jdx: label for sample in view 2\n",
    "            pos: positions [layer,token_in_layer] of sample in view 1 / view 2\n",
    "        \"\"\"\n",
    "        # get model and epoch index\n",
    "        mdx, edx = self._index[index]\n",
    "\n",
    "        # get permutation index -> pick random number from available perms\n",
    "        if self.permutations_number > 0:\n",
    "            perm_idx, perm_jdx = random.choices(\n",
    "                list(range(self.permutations_number)), k=2\n",
    "            )\n",
    "\n",
    "        ## mode \"vector has different workflow\"\n",
    "        if self.mode == \"vector\":\n",
    "            # get raw data, assume view 1 and view 2 are the same\n",
    "            ddx = self.data[mdx][edx]\n",
    "            label = self.labels[mdx][edx]\n",
    "\n",
    "            # slice window of size windowsize\n",
    "            idx_start = random.randint(0, len(ddx) - self.window)\n",
    "\n",
    "            # get position\n",
    "            pos = self.positions[idx_start : idx_start + self.window]\n",
    "\n",
    "            # get permutation index -> pick random number from available perms\n",
    "            if self.permutations_number > 0:\n",
    "                perm_idx, perm_jdx = random.choices(\n",
    "                    list(range(self.permutations_number)), k=2\n",
    "                )\n",
    "\n",
    "            # permutation\n",
    "            # permute data idx\n",
    "            if self.view_1_canonical:\n",
    "                ddx_idx = copy.deepcopy(ddx)\n",
    "                ddx_idx = ddx_idx[idx_start : idx_start + self.window]\n",
    "                label_idx = f\"{label}#_#canon\"\n",
    "            elif self.permutations_number > 0:\n",
    "                # get current permutation for slice of window\n",
    "                ddx_idx = copy.deepcopy(ddx)\n",
    "                permdx = self.permutations[perm_idx]\n",
    "                ddx_idx = permute_model_vector(\n",
    "                    vec=ddx_idx, perm=permdx, idx_start=idx_start, window=self.window\n",
    "                )\n",
    "                # update label\n",
    "                label_idx = f\"{label}#_#per_{perm_idx}\"\n",
    "            else:\n",
    "                ddx_idx = copy.deepcopy(ddx)\n",
    "                ddx_idx = ddx_idx[idx_start : idx_start + self.window]\n",
    "                label_idx = copy.deepcopy(label)\n",
    "\n",
    "            # permute data jdx\n",
    "            if self.view_2_canonical:\n",
    "                ddx_jdx = copy.deepcopy(ddx)\n",
    "                ddx_jdx = ddx_jdx[idx_start : idx_start + self.window]\n",
    "                label_jdx = f\"{label}#_#canon\"\n",
    "            elif self.permutations_number > 0:\n",
    "                # get current permutation for slice of window\n",
    "                ddx_jdx = copy.deepcopy(ddx)\n",
    "                permdx = self.permutations[perm_jdx]\n",
    "                ddx_jdx = permute_model_vector(\n",
    "                    vec=ddx_jdx, perm=permdx, idx_start=idx_start, window=self.window\n",
    "                )\n",
    "                # update label\n",
    "                label_jdx = f\"{label}#_#per_{perm_jdx}\"\n",
    "            else:\n",
    "                ddx_jdx = copy.deepcopy(ddx)\n",
    "                ddx_jdx = ddx_jdx[idx_start : idx_start + self.window]\n",
    "                label_jdx = copy.deepcopy(label)\n",
    "\n",
    "            # noise\n",
    "            if not self.add_noise_view_1 == False:\n",
    "                # add noise to input\n",
    "                # check sigma is larger than 0\n",
    "                if self.add_noise_view_1 > 0:\n",
    "                    if self.use_multiplicative_noise:\n",
    "                        # multiply each token with noise (uniform distribution around 1)\n",
    "                        # noise idx\n",
    "                        ddx_idx = ddx_idx * (1.0 + self.add_noise_view_1 * torch.randn(ddx_idx.shape))\n",
    "                    else:\n",
    "                        # add to each token noise (uniform distribution around 0)\n",
    "                        # noise idx\n",
    "                        ddx_idx = ddx_idx + self.add_noise_view_1 * torch.randn(ddx_idx.shape)\n",
    "\n",
    "            if not self.add_noise_view_2 == False:\n",
    "                # check sigma is number\n",
    "                assert isinstance(self.add_noise_view_2, float)\n",
    "                # add noise to input\n",
    "                # check sigma is larger than 0\n",
    "                if self.add_noise_view_2 > 0:\n",
    "                    if self.use_multiplicative_noise:\n",
    "                        # multiply each token with noise (uniform distribution around 1)\n",
    "                        # noise jdx\n",
    "                        ddx_jdx = ddx_jdx * (1.0 + self.add_noise_view_2 * torch.randn(ddx_jdx.shape))\n",
    "                    else:\n",
    "                        # add to each token noise (uniform distribution around 0)\n",
    "                        # noise jdx\n",
    "                        ddx_jdx = ddx_jdx + self.add_noise_view_2 * torch.randn(ddx_jdx.shape)\n",
    "\n",
    "            # erase_input/output augmentation\n",
    "            if self.erase_augment is not None:\n",
    "                # apply erase_augment on each token\n",
    "                ddx_idx = [self.erase_augment(iddx) for iddx in ddx_idx]\n",
    "                ddx_jdx = [self.erase_augment(iddx) for iddx in ddx_jdx]\n",
    "\n",
    "            # tokenize\n",
    "            ddx_idx, mask_idx = tokenize(t=ddx_idx, tokensize=self.tokensize)\n",
    "            ddx_jdx, mask_jdx = tokenize(t=ddx_jdx, tokensize=self.tokensize)\n",
    "\n",
    "            return ddx_idx, mask_idx, label_idx, ddx_jdx, mask_jdx, label_jdx, pos\n",
    "        ### end mode==\"vector\"\n",
    "\n",
    "        # TODO implement version on checkpoints, relatively straight-forward\n",
    "        raise NotImplementedError(\n",
    "            \"simclr dataset is not yet implemented for checkpoints. hang in there..\"\n",
    "        )\n",
    "\n",
    "    ### len ##################################################################################################################################################################\n",
    "    def init_len(self):\n",
    "        index = []\n",
    "        for idx, ddx in enumerate(self.data):\n",
    "            idx_tmp = [(idx, jdx) for jdx in range(len(ddx))]\n",
    "            index.extend(idx_tmp)\n",
    "        self._len = len(index)\n",
    "        self._index = index\n",
    "\n",
    "    def __len__(self):\n",
    "        return self._len\n",
    "\n",
    "    ### set erase ############################################################\n",
    "    def set_erase(self, erase=None):\n",
    "        if erase is not None:\n",
    "            assert (\n",
    "                self.mode == \"vectorize\" or self.mode == \"vector\"\n",
    "            ), \"erasing is only for vectorized mode implemenetd\"\n",
    "            erase = RandomErasingVector(\n",
    "                p=erase[\"p\"],\n",
    "                scale=erase[\"scale\"],\n",
    "                value=erase[\"value\"],\n",
    "                mode=erase[\"mode\"],\n",
    "            )\n",
    "        else:\n",
    "            erase = None\n",
    "        self.erase_augment = erase\n",
    "\n",
    "    ### vectorize_data #########################################################################################################################################################\n",
    "    def vectorize_data(self):\n",
    "        # iterate over models\n",
    "        for idx in range(len(self.data)):\n",
    "            # iterate over epochs\n",
    "            for jdx in range(len(self.data[idx])):\n",
    "                # get checkpoint\n",
    "                checkpoint = copy.deepcopy(self.data[idx][jdx])\n",
    "                # get vectorize\n",
    "                ddx = vectorize_checkpoint(checkpoint)\n",
    "                # overwrite data\n",
    "                self.data[idx][jdx] = ddx\n",
    "\n",
    "    ### standardize data #########################################################################################################################################################\n",
    "    def standardize_data(self):\n",
    "        \"\"\"\n",
    "        standardize data to zero-mean / unit std. per layer\n",
    "        store per-layer mean / std\n",
    "        \"\"\"\n",
    "        logging.info(\"Get layer mapping\")\n",
    "        # step 1: get token-layer index relation\n",
    "        layers = {}\n",
    "        # init vals\n",
    "        cur_layer = 0\n",
    "        cur_layer_start = 0\n",
    "        # iterate over models\n",
    "        for idx in range(self.positions.shape[0]):\n",
    "            ldx, kdx = self.positions[idx, 0].item(), self.positions[idx, 1].item()\n",
    "            # if new layer\n",
    "            if not ldx == cur_layer:\n",
    "                # add previous layer to dict\n",
    "                layer = str(cur_layer)\n",
    "                cur_layer_end = idx - 1\n",
    "                layers[layer] = {\n",
    "                    \"start_idx\": cur_layer_start,\n",
    "                    \"end_idx\": cur_layer_end,\n",
    "                }\n",
    "                # start new layer\n",
    "                cur_layer = ldx\n",
    "                cur_layer_start = idx\n",
    "        # last layer\n",
    "        layer = str(ldx)\n",
    "        cur_layer_end = idx\n",
    "        layers[layer] = {\n",
    "            \"start_idx\": cur_layer_start,\n",
    "            \"end_idx\": cur_layer_end,\n",
    "        }\n",
    "        logging.debug(f\"layer mapping: {layers}\")\n",
    "\n",
    "        logging.info(\"Get layer-wise mean and std\")\n",
    "\n",
    "        # iterate over layers\n",
    "        for layer in layers:\n",
    "            idx_start = layers[layer][\"start_idx\"]\n",
    "            idx_end = layers[layer][\"end_idx\"]\n",
    "\n",
    "            # collect all tokens within the layer for all models\n",
    "            tmp = []\n",
    "            # iterate over models\n",
    "            for idx in range(len(self.data)):\n",
    "                for jdx in range(len(self.data[idx])):\n",
    "                    tmp2 = [\n",
    "                        self.data[idx][jdx][ldx]\n",
    "                        for ldx in range(idx_start, idx_end + 1)\n",
    "                    ]\n",
    "                    tmp.extend(tmp2)\n",
    "\n",
    "            # stack / cat\n",
    "            tmp = torch.stack(tmp, dim=0)\n",
    "\n",
    "            # compute mean / std\n",
    "            mu = torch.mean(tmp)\n",
    "            sigma = torch.std(tmp)\n",
    "\n",
    "            # store in layer\n",
    "            layers[layer][\"mean\"] = mu\n",
    "            layers[layer][\"std\"] = sigma\n",
    "\n",
    "            # free memory\n",
    "            del tmp\n",
    "\n",
    "        self.layers = layers\n",
    "\n",
    "        logging.info(\"Apply standardization\")\n",
    "        # TODO: make more efficient, cut out the first two for loops with list expression?\n",
    "        # standardize:\n",
    "        # # iterate over models\n",
    "        for idx in range(len(self.data)):\n",
    "            for jdx in range(len(self.data[idx])):\n",
    "                # # iterate over tokens of that layer\n",
    "                for tdx in range(len(self.data[idx][jdx])):\n",
    "                    # get position\n",
    "                    pos = self.positions[tdx]\n",
    "                    ldx = pos[0].item()\n",
    "                    # get mu/sigma for that token\n",
    "                    mu = layers[str(ldx)][\"mean\"]\n",
    "                    std = layers[str(ldx)][\"std\"]\n",
    "                    # # standardize with mean / std\n",
    "                    self.data[idx][jdx][tdx] = (self.data[idx][jdx][tdx] - mu) / std\n",
    "\n",
    "    ### precompute_permutation_index #########################################################################################################################################################\n",
    "    def precompute_permutation_index(self):\n",
    "        # ASSUMES THAT DATA IS ALREADY VECTORIZED\n",
    "        permutation_index_list = []\n",
    "        # create index vector\n",
    "        # print(f\"vector shape: {self.data_in[0].shape}\")\n",
    "        index_vector = torch.tensor(list(range(self.data_in[0].shape[0])))\n",
    "        # cast index vector to double\n",
    "        index_vector = index_vector.double()\n",
    "        # print(f\"index vector: {index_vector}\")\n",
    "        # reference checkpoint\n",
    "        reference_checkpoint = copy.deepcopy(self.reference_checkpoint)\n",
    "        # cast index vector to checkpoint\n",
    "        index_checkpoint = vector_to_checkpoint(\n",
    "            checkpoint=copy.deepcopy(reference_checkpoint),\n",
    "            vector=copy.deepcopy(index_vector),\n",
    "            layer_lst=self.layer_lst,\n",
    "            use_bias=self.use_bias,\n",
    "        )\n",
    "\n",
    "        ## init multiprocessing environment ############\n",
    "        ray.init(num_cpus=self.num_threads)\n",
    "\n",
    "        ### gather data #############################################################################################\n",
    "        logging.info(f\"preparing permutation indices from {self.root}\")\n",
    "        pb = ProgressBar(total=self.permutations_number)\n",
    "        pb_actor = pb.actor\n",
    "\n",
    "        # loop over all permutations in self.permutations_number\n",
    "        for pdx in range(self.permutations_number):\n",
    "            # get perm dict\n",
    "            prmt_dct = self.permutations_dct_lst[pdx]\n",
    "            #\n",
    "            index_p = compute_single_index_vector_remote.remote(\n",
    "                index_checkpoint=copy.deepcopy(index_checkpoint),\n",
    "                prmt_dct=prmt_dct,\n",
    "                layer_lst=self.layer_lst,\n",
    "                permute_layers=self.permute_layers,\n",
    "                use_bias=self.use_bias,\n",
    "                pba=pb_actor,\n",
    "            )\n",
    "            # append to permutation_index_list\n",
    "            permutation_index_list.append(index_p)\n",
    "\n",
    "        # update progress bar\n",
    "        pb.print_until_done()\n",
    "\n",
    "        # collect actual data\n",
    "        permutation_index_list = ray.get(permutation_index_list)\n",
    "\n",
    "        ray.shutdown()\n",
    "\n",
    "        self.permutation_index_list = permutation_index_list\n",
    "\n",
    "    ### map data to canoncial #############################################################################################\n",
    "    def map_models_to_canonical(self):\n",
    "        \"\"\"\n",
    "        define reference model\n",
    "        iterate over all models\n",
    "        get permutation w.r.t last epoch (best convergence)\n",
    "        apply same permutation on all epochs (on raw data)\n",
    "        \"\"\"\n",
    "        # use first model / last epoch as reference model (might be sub-optimal)\n",
    "        reference_model = self.reference_checkpoint\n",
    "\n",
    "        ## init multiprocessing environment ############\n",
    "        ray.init(num_cpus=self.num_threads)\n",
    "\n",
    "        ### gather data #############################################################################################\n",
    "        print(f\"preparing computing canon form...\")\n",
    "        pb = ProgressBar(total=len(self.data))\n",
    "        pb_actor = pb.actor\n",
    "\n",
    "        for idx in range(len(self.data)):\n",
    "            # align models using git-re-basin\n",
    "            perm_spec = self.permutation_spec\n",
    "            # get second\n",
    "            model_curr = self.data[idx]\n",
    "            model_curr = compute_single_canon_form.remote(\n",
    "                reference_model=reference_model,\n",
    "                data_curr=model_curr,\n",
    "                perm_spec=perm_spec,\n",
    "                pba=pb_actor,\n",
    "            )\n",
    "            self.data[idx] = model_curr\n",
    "\n",
    "        # update progress bar\n",
    "        pb.print_until_done()\n",
    "\n",
    "        self.data = [ray.get(self.data[idx]) for idx in range(len(self.data))]\n",
    "        ray.shutdown()\n",
    "\n",
    "    ### precompute_permutations #############################################################################################\n",
    "    def precompute_permutations(self, permutation_number, perm_spec, num_threads=6):\n",
    "        \"\"\"\n",
    "        - get permutation_dict as template\n",
    "        - generate random permutations\n",
    "        - generate index checkpoint\n",
    "            - copy actual checkpoint\n",
    "            - flatten tensor\n",
    "            - generate list with indices of same shape\n",
    "            - put on tensor\n",
    "            - reshape to original view\n",
    "        - apply permutation on checkpoint: store checkpoint as permutation dict.\n",
    "\n",
    "        - permuting modules translates to:\n",
    "            - get right index for current module\n",
    "            - .flatten()\n",
    "            - apply permutation /slice\n",
    "            - .view()\n",
    "        \"\"\"\n",
    "        logging.info(\"start precomputing permutations\")\n",
    "        # model_curr = self.data[0][-1]\n",
    "        model_curr = self.reference_checkpoint\n",
    "        # find permutation of model to itself as reference\n",
    "        reference_permutation = weight_matching(\n",
    "            ps=perm_spec, params_a=model_curr, params_b=model_curr\n",
    "        )\n",
    "\n",
    "        logging.info(\"get random permutation dicts\")\n",
    "        # compute random permutations\n",
    "        permutation_dicts = []\n",
    "        for ndx in range(permutation_number):\n",
    "            perm = copy.deepcopy(reference_permutation)\n",
    "            for key in perm.keys():\n",
    "                # get permuted indecs for current layer\n",
    "                perm[key] = torch.randperm(perm[key].shape[0]).float()\n",
    "            # append to list of permutation dicts\n",
    "            permutation_dicts.append(perm)\n",
    "\n",
    "        self.permutation_dicts = permutation_dicts\n",
    "\n",
    "        if self.mode == \"vector\":\n",
    "            logging.info(\"get permutation indices\")\n",
    "            # get permutation data\n",
    "            ## get reference checkpoint\n",
    "            ref_checkpoint = copy.deepcopy(model_curr)\n",
    "            ## vectoirze\n",
    "            ref_vec_global = vectorize_checkpoint(ref_checkpoint)\n",
    "            ref_vec_kernel = copy.deepcopy(ref_vec_global)\n",
    "            ## get reference index vec\n",
    "            for idx, module in enumerate(ref_vec_global):\n",
    "                # get global index of permutation between kernels\n",
    "                index_global = torch.ones(module.numel()) * idx\n",
    "                index_global = index_global.view(module.shape)\n",
    "                ref_vec_global[idx] = index_global\n",
    "                # got local index of permutation within kernels\n",
    "                index_kernel = torch.tensor(list(range(module.numel())))\n",
    "                index_kernel = index_kernel.view(module.shape)\n",
    "                ref_vec_kernel[idx] = index_kernel\n",
    "            ## map to checkpoint\n",
    "            ref_checkpoint_global = vector_to_checkpoint(\n",
    "                vector=ref_vec_global, reference_checkpoint=ref_checkpoint\n",
    "            )\n",
    "            ref_checkpoint_kernel = vector_to_checkpoint(\n",
    "                vector=ref_vec_kernel, reference_checkpoint=ref_checkpoint\n",
    "            )\n",
    "\n",
    "            ## init multiprocessing environment ############\n",
    "            ray.init(num_cpus=num_threads)\n",
    "            pb = ProgressBar(total=permutation_number)\n",
    "            pb_actor = pb.actor\n",
    "            # get permutations\n",
    "            permutations_global = []\n",
    "            permutations_kernel = []\n",
    "            for perm_dict in permutation_dicts:\n",
    "                perm_curr_global = compute_single_perm.remote(\n",
    "                    reference_checkpoint=ref_checkpoint_global,\n",
    "                    permutation_dict=perm_dict,\n",
    "                    perm_spec=perm_spec,\n",
    "                    pba=pb_actor,\n",
    "                )\n",
    "\n",
    "                perm_curr_kernel = compute_single_perm.remote(\n",
    "                    reference_checkpoint=ref_checkpoint_kernel,\n",
    "                    permutation_dict=perm_dict,\n",
    "                    perm_spec=perm_spec,\n",
    "                    pba=pb_actor,\n",
    "                )\n",
    "\n",
    "                permutations_global.append(perm_curr_global)\n",
    "                permutations_kernel.append(perm_curr_kernel)\n",
    "\n",
    "            permutations_global = ray.get(permutations_global)\n",
    "            permutations_kernel = ray.get(permutations_kernel)\n",
    "\n",
    "            permutations_global = [\n",
    "                torch.tensor([perm[0].item() for perm in perm_g]).int()\n",
    "                for perm_g in permutations_global\n",
    "            ]\n",
    "\n",
    "            permutations = [\n",
    "                (perm_g, perm_k)\n",
    "                for (perm_g, perm_k) in zip(permutations_global, permutations_kernel)\n",
    "            ]\n",
    "\n",
    "            ray.shutdown()\n",
    "\n",
    "            self.permutations = permutations\n",
    "\n",
    "\n",
    "def permute_model_vector(\n",
    "    vec: List[torch.Tensor], perm: List[tuple], idx_start: int, window: int\n",
    "):\n",
    "    \"\"\"\n",
    "    performs permutation on vectorized model and returns slice of tokens\n",
    "    input vec: list(torch.tensor) with the weights per output channel of the full model\n",
    "    input perm: contains two pieces of information for permutation. perm[0] is the global permutation of the tokens, perm[1] contains the permutaion mappings within tokens per token\n",
    "    input idx_start: int marks the start of a slice of tokens to keep\n",
    "    input window: int marks the size of the slice to keep\n",
    "    return vec: list(torch.tensor) of permuted and sliced tokens\n",
    "    \"\"\"\n",
    "    # create index vector of tokens\n",
    "    index = list(range(len(vec)))\n",
    "\n",
    "    # apply global permutation on index\n",
    "    # using (slices of) the permuted index to access tokens equals permuting all tokens and slicing after\n",
    "    perm_glob = perm[0]\n",
    "    index = [index[idx] for idx in perm_glob]\n",
    "\n",
    "    # slice index\n",
    "    idx_end = idx_start + window\n",
    "    index = index[idx_start:idx_end]\n",
    "\n",
    "    # slice token sequence\n",
    "    vec = [vec[idx] for idx in index]\n",
    "\n",
    "    # slice permutations\n",
    "    perm_loc = perm[1]\n",
    "    perm_loc = [perm_loc[idx] for idx in index]\n",
    "\n",
    "    # apply token permutation\n",
    "    vec = [vecdx[permdx] for (vecdx, permdx) in zip(vec, perm_loc)]\n",
    "\n",
    "    # return tokens\n",
    "    return vec\n",
    "\n",
    "\n",
    "### helper parallel function #############################################################################################\n",
    "@ray.remote(num_returns=1)\n",
    "def compute_single_perm(reference_checkpoint, permutation_dict, perm_spec, pba):\n",
    "    # copy reference checkpoint\n",
    "    index_check = copy.deepcopy(reference_checkpoint)\n",
    "    # apply permutation on checkpoint\n",
    "    index_check_perm = apply_permutation(\n",
    "        ps=perm_spec, perm=permutation_dict, params=index_check\n",
    "    )\n",
    "    # vectorize\n",
    "    index_perm = vectorize_checkpoint(index_check_perm)\n",
    "    # update counter\n",
    "    pba.update.remote(1)\n",
    "    # return list\n",
    "    return index_perm\n",
    "\n",
    "\n",
    "### helper parallel function #############################################################################################\n",
    "@ray.remote(num_returns=1)\n",
    "def compute_single_canon_form(reference_model, data_curr, perm_spec, pba):\n",
    "    # get second\n",
    "    model_curr = data_curr[-1]\n",
    "    # find permutation to match params_b to params_a\n",
    "    logging.debug(\n",
    "        f\"compute canonical form: params a {type(reference_model)} params b {type(model_curr)}\"\n",
    "    )\n",
    "    match_permutation = weight_matching(\n",
    "        ps=perm_spec, params_a=reference_model, params_b=model_curr\n",
    "    )\n",
    "    # apply permutation on all epochs\n",
    "    for jdx in range(len(data_curr)):\n",
    "        model_curr = data_curr[jdx]\n",
    "        model_curr_perm = apply_permutation(\n",
    "            ps=perm_spec, perm=match_permutation, params=model_curr\n",
    "        )\n",
    "        # put back in data\n",
    "        data_curr[jdx] = model_curr_perm\n",
    "\n",
    "    # update counter\n",
    "    pba.update.remote(1)\n",
    "    # return list\n",
    "    return data_curr\n",
    "\n",
    "\n",
    "### helper parallel function #############################################################################################\n",
    "@ray.remote(num_returns=1)\n",
    "def compute_single_index_vector_remote(\n",
    "    index_checkpoint, prmt_dct, layer_lst, permute_layers, use_bias, pba\n",
    "):\n",
    "    # apply permutation on copy of unit checkpoint\n",
    "    # TODO: fix function,\n",
    "    chkpt_p = permute_checkpoint(\n",
    "        index_checkpoint,\n",
    "        layer_lst,\n",
    "        permute_layers,\n",
    "        prmt_dct,\n",
    "    )\n",
    "\n",
    "    # cast back to vector\n",
    "    vector_p = vectorize_checkpoint(copy.deepcopy(chkpt_p), layer_lst, use_bias)\n",
    "    # cast vector back to int\n",
    "    vector_p = vector_p.int()\n",
    "    # we specifically don't check for uniqueness of indices. we'd rather let this run into index errors to catch the issue\n",
    "    index_p = copy.deepcopy(vector_p.tolist())\n",
    "    # update counter\n",
    "    pba.update.remote(1)\n",
    "    # return list\n",
    "    return index_p\n",
    "\n",
    "\n",
    "def vectorize_checkpoint(checkpoint):\n",
    "    out = []\n",
    "    # use only weights and biases\n",
    "    for key in checkpoint.keys():\n",
    "        if \"weight\" in key:\n",
    "            w = checkpoint[key]\n",
    "            # flatten to out_channels x n\n",
    "            w = w.view(w.shape[0], -1)\n",
    "            # cat biases to channels if they exist in checkpoint\n",
    "            if key.replace(\"weight\", \"bias\") in checkpoint:\n",
    "                b = checkpoint[key.replace(\"weight\", \"bias\")]\n",
    "                w = torch.cat([w, b.unsqueeze(dim=1)], dim=1)\n",
    "            # split weights in slices along output channel dims\n",
    "            w = torch.split(w, w.shape[0])\n",
    "            # extend out with new tokens, zero's (and only entry) is a list\n",
    "            out.extend(w[0])\n",
    "\n",
    "    return out\n",
    "\n",
    "\n",
    "def vector_to_checkpoint(vector, reference_checkpoint):\n",
    "    # make copy to prevent memory management issues\n",
    "    checkpoint = copy.deepcopy(reference_checkpoint)\n",
    "    # use only weights and biases\n",
    "    idx_start = 0\n",
    "    for key in checkpoint.keys():\n",
    "        if \"weight\" in key:\n",
    "            # get correct slice of modules out of vec sequence\n",
    "            out_channels = checkpoint[key].shape[0]\n",
    "            idx_end = idx_start + out_channels\n",
    "            w = vector[idx_start:idx_end]\n",
    "            # get weight matrix from list of vectors\n",
    "            try:\n",
    "                w = torch.stack(w, dim=0)\n",
    "            except Exception as e:\n",
    "                logging.error(\n",
    "                    f\"vector_to_checkpoint: layer {key}, idx: [{idx_start},{idx_end}] created weight {w.shape} for checkpoint weight {checkpoint[key].shape}\"\n",
    "                )\n",
    "                logging.error(e)\n",
    "            # extract bias\n",
    "            if key.replace(\"weight\", \"bias\") in checkpoint:\n",
    "                b = w[:, -1]\n",
    "                checkpoint[key.replace(\"weight\", \"bias\")] = b\n",
    "                w = w[:, :-1]\n",
    "            # reshape weight vector\n",
    "            w = w.view(checkpoint[key].shape)\n",
    "            logging.debug(\n",
    "                f\"vector_to_checkpoint: layer {key}, idx: [{idx_start},{idx_end}] tried to create weights from {[wdx.shape] for wdx in w} for checkpoint weight {checkpoint[key].shape}\"\n",
    "            )\n",
    "            checkpoint[key] = w\n",
    "            # update start\n",
    "            idx_start = idx_end\n",
    "    return checkpoint\n",
    "\n",
    "\n",
    "def get_position_mapping_from_checkpoint(checkpoint) -> list:\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        checkpoint: Collections.OrderedDict model checkpoint\n",
    "    Returns:\n",
    "        output tensor with 2d positions for every token in the vectorized model sequence\n",
    "    \"\"\"\n",
    "    out = []\n",
    "    # start layer index counter at 0\n",
    "    idx = 0\n",
    "    # iterate over modules\n",
    "    for key in checkpoint.keys():\n",
    "        # if module with weights is found -> add index\n",
    "        if \"weight\" in key:\n",
    "            w = checkpoint[key]\n",
    "            # create tuple (layer_idx, channel_idx) for every channel\n",
    "            idx_layer = [torch.tensor([idx, jdx]) for jdx in range(w.shape[0])]\n",
    "            # add to overall position\n",
    "            out.extend(idx_layer)\n",
    "            # increase layer counter\n",
    "            idx += 1\n",
    "    out = torch.stack(out, dim=0)\n",
    "    return out\n",
    "\n",
    "\n",
    "def tokenize(t: List[torch.tensor], tokensize: int, return_mask: bool = True):\n",
    "    \"\"\"\n",
    "    transforms list of tokens of differen lenght to tensor\n",
    "    Args:\n",
    "        t: List[torch.tensor]: list of 1d input tokens of different lenghts\n",
    "        tokensize: int output dimension of each token\n",
    "        return_mask: bool wether to return the mask of nonzero values\n",
    "    Returns:\n",
    "        tokens: torch.tensor with tokens stacked along dim=0\n",
    "        mask: torch.tensor indicating the shape of the original tokens\n",
    "    \"\"\"\n",
    "    # init output with zeros\n",
    "    tokens = torch.zeros(len(t), tokensize)\n",
    "    mask = torch.zeros(len(t), tokensize)\n",
    "    # iterate over inputs\n",
    "    for idx, tdx in enumerate(t):\n",
    "        # get end of token, either the length of the input or them maximum length\n",
    "        tdx_end = min(tdx.shape[0], tokensize)\n",
    "        # put at position idx\n",
    "        tokens[idx, :tdx_end] = tdx[:tdx_end]\n",
    "        mask[idx, :tdx_end] = torch.ones(tdx_end)\n",
    "\n",
    "    # return\n",
    "    if return_mask:\n",
    "        return tokens, mask\n",
    "    else:\n",
    "        return tokens\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8188133",
   "metadata": {},
   "outputs": [],
   "source": [
    "def permute_model_tensor(\n",
    "    vec: List[torch.Tensor], perm: List[tuple], idx_start: int, window: int\n",
    "):\n",
    "    \"\"\"\n",
    "    performs permutation on vectorized model and returns slice of tokens\n",
    "    input vec: list(torch.tensor) with the weights per output channel of the full model\n",
    "    input perm: contains two pieces of information for permutation. perm[0] is the global permutation of the tokens, perm[1] contains the permutaion mappings within tokens per token\n",
    "    input idx_start: int marks the start of a slice of tokens to keep\n",
    "    input window: int marks the size of the slice to keep\n",
    "    return vec: list(torch.tensor) of permuted and sliced tokens\n",
    "    \"\"\"\n",
    "    # create index vector of tokens\n",
    "    index = list(range(len(vec)))\n",
    "\n",
    "    # apply global permutation on index\n",
    "    # using (slices of) the permuted index to access tokens equals permuting all tokens and slicing after\n",
    "    perm_glob = perm[0]\n",
    "    index = [index[idx] for idx in perm_glob]\n",
    "\n",
    "    # slice index\n",
    "    idx_end = idx_start + window\n",
    "    index = index[idx_start:idx_end]\n",
    "\n",
    "    # slice token sequence\n",
    "    vec = [vec[idx] for idx in index]\n",
    "\n",
    "    # slice permutations\n",
    "    perm_loc = perm[1]\n",
    "    perm_loc = [perm_loc[idx] for idx in index]\n",
    "\n",
    "    # apply token permutation\n",
    "    vec = [vecdx[permdx] for (vecdx, permdx) in zip(vec, perm_loc)]\n",
    "\n",
    "    # return tokens\n",
    "    return vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0e595c70",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-04-12 03:46:52,774\tINFO worker.py:1553 -- Started a local Ray instance.\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2800/2800 [00:08<00:00, 337.88it/s]\n",
      "700it [00:12, 57.53it/s]\n",
      "2023-04-12 03:47:29,804\tINFO worker.py:1553 -- Started a local Ray instance.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "preparing computing canon form...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 700/700 [00:02<00:00, 272.80it/s]\n"
     ]
    }
   ],
   "source": [
    "config_key_list = []\n",
    "result_key_list = [\n",
    "    \"test_acc\",\n",
    "    \"training_iteration\",\n",
    "    \"ggap\",\n",
    "    #     \"sparsity_ratio\",\n",
    "]\n",
    "property_keys = {\n",
    "    \"result_keys\": result_key_list,\n",
    "    \"config_keys\": config_key_list,\n",
    "}\n",
    "\n",
    "path_root = Path(\"/netscratch2/dtaskiran/zoos/SVHN/tune_zoo_svhn_uniform/\")\n",
    "\n",
    "dataset_tensor = SimCLRDatasetTensor(\n",
    "    root=path_root,\n",
    "    epoch_lst=[1, 5, 20, 25],\n",
    "#     epoch_lst=list(range(1, 26)),\n",
    "    #     mode=\"checkpoint\",\n",
    "    mode=\"vector\",\n",
    "    permutations_number=0,\n",
    "    permutation_spec=zoo_cnn_permutation_spec(),\n",
    "    #     view_1_canonical = False,\n",
    "    view_1_canonical=True,\n",
    "    view_2_canonical=True,\n",
    "    #     view_2_canonical = False,\n",
    "    add_noise_view_1=0.1,  # [('input', 0.15), ('output', 0.013)]\n",
    "    add_noise_view_2=0.0,  # [('input', 0.15), ('output', 0.013)]\n",
    "    noise_multiplicative=True,\n",
    "    erase_augment=None,  # {\"p\": 0.5,\"scale\":(0.02,0.33),\"value\":0,\"mode\":\"block\"}\n",
    "    windowsize=12,\n",
    "    standardize=True,\n",
    "    train_val_test=\"train\",  # determines whcih dataset split to use\n",
    "    ds_split=[0.7, 0.15, 0.15],  #\n",
    "    max_samples=1000,\n",
    "    # weight_threshold=float(\"inf\"),\n",
    "    weight_threshold=15,\n",
    "    filter_function=None,  # gets sample path as argument and returns True if model needs to be filtered out\n",
    "    property_keys=property_keys,\n",
    "    num_threads=6,\n",
    "    shuffle_path=True,\n",
    "    verbosity=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e0d0577a",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_int = len(dataset_tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "578ed3b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "91 µs ± 645 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "for idx in range(64):\n",
    "    kdx = random.randint(0, max_int)\n",
    "#     (x_i, m_i, l_i, x_j, m_j, _, p) = dataset_tensor.__getitem__(kdx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "5f072f44",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.41 s ± 531 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "for idx in range(64):\n",
    "    kdx = random.randint(0, max_int-1)\n",
    "    (x_i, m_i, l_i, x_j, m_j, _, p) = dataset_tensor.__getitem__(kdx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "559dcf0b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
