test_multimodel_ckpt.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "id": "bb5f6d9d",
   "metadata": {},
   "source": [
    "Possible adds\n",
    "* Other ways to ensemble (e.g. Ensemble probabilities from experts instead of features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b7d19f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import random\n",
    "from PIL import Image\n",
    "import os\n",
    "from torchvision.datasets import CIFAR10, CIFAR100, CocoCaptions, ImageNet\n",
    "import slip_models\n",
    "from tokenizer import SimpleTokenizer\n",
    "\n",
    "import torchvision.transforms as T\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import numpy as np\n",
    "%pip install ipywidgets\n",
    "%pip install update tqdm%matplotlib inline\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a7e613a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !/export/home/google-cloud-sdk/bin/gsutil -m cp gs://sfr-tpu-us-east1-research/bwallace/ckpts/models/multimodel_clip/parallel_vision_standard_text/viz_parallel_text_standard_head_dim_64_per_expert_norm_bs_1024_epoch_30.ckpt ckpts/viz_parallel_text_standard_head_dim_64_per_expert_norm_bs_1024_epoch_30.ckpt\n",
    "!/export/home/google-cloud-sdk/bin/gsutil -m cp gs://sfr-tpu-us-east1-research/bwallace/ckpts/models/multimodel_clip/parallel_vision_standard_text/viz_parallel_text_standard_head_dim_64_top1_expert_loss_no_entroy_bs_1024_epoch_6.ckpt ckpts/viz_parallel_text_standard_head_dim_64_top1_expert_loss_no_entroy_bs_1024_epoch_6.ckpt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "383b607d",
   "metadata": {},
   "outputs": [],
   "source": [
    "preprocess = T.Compose(\n",
    "                    [\n",
    "                        T.Resize(256),\n",
    "                        T.CenterCrop(224),\n",
    "                        T.ToTensor(),\n",
    "                        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "                    ]\n",
    "                        )\n",
    "\n",
    "grayscale_preprocess = T.Compose(\n",
    "                    [\n",
    "                        T.Resize(256),\n",
    "                        T.CenterCrop(224),\n",
    "                        T.ToTensor(),\n",
    "                        T.Normalize(mean=[np.average([0.485, 0.456, 0.406])], std=[np.average([0.229, 0.224, 0.225])]),\n",
    "                        T.Lambda(lambda x: x.repeat(3, 1, 1) ),\n",
    "                    ]\n",
    "                        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2552cd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "grayscale_preprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "578d6aa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = slip_models.CLIP_VITB16()\n",
    "# ckpt = torch.load('ckpts/clip_vit_epoch_35.ckpt', map_location='cpu')\n",
    "\n",
    "model = slip_models.VisionParallelTextStandard(64, 64)\n",
    "ckpt = torch.load('ckpts/viz_parallel_text_standard_head_dim_64_top1_expert_loss_no_entroy_bs_1024_epoch_6.ckpt', map_location='cpu')\n",
    "per_group_norm = True\n",
    "# ckpt = torch.load('tmpp_save/test_epoch_4.ckpt', map_location='cpu')\n",
    "# model = slip_models.VisionParallelTextStandard(64, 8)\n",
    "# ckpt = torch.load('ckpts/viz_parallel_text_standard_bs_1024_epoch_35.ckpt', map_location='cpu')\n",
    "# per_group_norm = False\n",
    "\n",
    "# model = slip_models.CLIP_VITB16(num_prompt_tokens=64, num_text_outputs=1000)\n",
    "# ckpt = torch.load('ckpts/epoch_30_prompted_clip_may_27.ckpt', map_location='cpu')\n",
    "\n",
    "\n",
    "# model.load_state_dict({k.replace('module.',''):v for k,v in ckpt[\"model\"].items()})\n",
    "model.load_state_dict(ckpt[\"model\"])\n",
    "\n",
    "model.eval()\n",
    "model = model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "433d4dd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cc_model = slip_models.CLIP_VITB16() # ssl_mlp_dim=4096, ssl_emb_dim=256)\n",
    "# cc_model =  slip_models.CLIP_VITB16()\n",
    "# ckpt = torch.load('ckpts/clip_vit_epoch_35.ckpt', map_location='cpu')\n",
    "# ckpt = torch.load('ckpts/retest_baseline_jul_27_epoch_5.ckpt', map_location='cpu')\n",
    "\n",
    "\n",
    "\n",
    "cc_model =  slip_models.CLIP_ResNet18()\n",
    "# ckpt = torch.load('ckpts/clip_vit_epoch_35.ckpt', map_location='cpu')\n",
    "ckpt = torch.load('ckpts/resnet18_baseline_jul_27_epoch_35.ckpt', map_location='cpu')\n",
    "\n",
    "# ckpt = torch.load('ckpts/slip_base_100ep.pt', map_location='cpu')\n",
    "# cc_model.load_state_dict({k.replace('module.',''):v for k,v in ckpt[\"state_dict\"].items()})\n",
    "cc_model.load_state_dict({k.replace('module.',''):v for k,v in ckpt[\"model\"].items()})\n",
    "cc_model = cc_model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "152dd04a",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = SimpleTokenizer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fff1d32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# image = preprocess(Image.open(\"pics/golden-retriever.png\")).unsqueeze(0).to(device)\n",
    "# image = preprocess(Image.open(\"pics/CLIP.png\")).unsqueeze(0).to(device)\n",
    "base_text = [\"a diagram\", \"a dog\", \"a cat\"]\n",
    "model.eval()\n",
    "dog_image = preprocess(Image.open(\"pics/golden-retriever.png\"))\n",
    "diagram_image = preprocess(Image.open(\"pics/CLIP.png\").convert(\"RGB\"))\n",
    "cat_image = preprocess(Image.open(\"pics/cat.jpg\"))\n",
    "images = torch.stack([diagram_image, dog_image, cat_image]).cuda()\n",
    "\n",
    "text = tokenizer([f\"a picture of {s}\" for s in base_text]).cuda()\n",
    "\n",
    "with torch.no_grad():\n",
    "    image_features = model.encode_image(images)\n",
    "    text_features = model.encode_text(text)\n",
    "    \n",
    "    if per_group_norm:\n",
    "        image_features = image_features.view(image_features.shape[0], 64, -1)\n",
    "        text_features = text_features.view(image_features.shape[0], 64, -1)\n",
    "    image_features = image_features / image_features.norm(dim=-1, keepdim=True)\n",
    "    text_features = text_features / text_features.norm(dim=-1, keepdim=True)\n",
    "    \n",
    "    image_features = image_features.flatten(start_dim=1)\n",
    "    text_features = text_features.flatten(start_dim=1)\n",
    "    \n",
    "    logits_per_image = model.logit_scale.exp() * image_features @ text_features.t()\n",
    "    logits_per_text = logits_per_image.t()\n",
    "    # logits_per_image, logits_per_text = model(images, text)\n",
    "    per_image_probs = logits_per_image.softmax(dim=-1).cpu().numpy()\n",
    "    per_text_probs = logits_per_text.softmax(dim=-1).cpu().numpy()\n",
    "\n",
    "print(\"(Per Image) Label probs:\\n\", per_image_probs)  # prints: [[0.9927937  0.00421068 0.00299572]]\n",
    "print(\"(Per Text) Label probs:\\n\", per_text_probs)  # prints: [[0.9927937  0.00421068 0.00299572]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3376a2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_features @ text_features.t()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c8b8939",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_features @ image_features.t()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6970496c",
   "metadata": {},
   "outputs": [],
   "source": [
    "grouped_if = image_features.view(3, 64, 64)\n",
    "grouped_sims = torch.bmm(grouped_if, grouped_if.transpose(1, 2))\n",
    "print(grouped_sims.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c8be067",
   "metadata": {},
   "outputs": [],
   "source": [
    "grouped_sims[0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbe1554e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# image = preprocess(Image.open(\"pics/golden-retriever.png\")).unsqueeze(0).to(device)\n",
    "# image = preprocess(Image.open(\"pics/CLIP.png\")).unsqueeze(0).to(device)\n",
    "base_text = [\"a diagram\", \"a dog\", \"a cat\"]\n",
    "cc_model.eval()\n",
    "dog_image = preprocess(Image.open(\"pics/golden-retriever.png\"))\n",
    "diagram_image = preprocess(Image.open(\"pics/CLIP.png\").convert(\"RGB\"))\n",
    "cat_image = preprocess(Image.open(\"pics/cat.jpg\"))\n",
    "images = torch.stack([diagram_image, dog_image, cat_image]).cuda()\n",
    "\n",
    "text = tokenizer([f\"a picture of {s}\" for s in base_text]).cuda()\n",
    "\n",
    "with torch.no_grad():\n",
    "    image_features = cc_model.encode_image(images)\n",
    "    text_features = cc_model.encode_text(text)\n",
    "    \n",
    "    image_features = image_features / image_features.norm(dim=-1, keepdim=True)\n",
    "    text_features = text_features / text_features.norm(dim=-1, keepdim=True)\n",
    "    \n",
    "    \n",
    "    logits_per_image = cc_model.logit_scale.exp() * image_features @ text_features.t()\n",
    "    # logits_per_image, logits_per_text = model(images, text)\n",
    "    probs = logits_per_image.softmax(dim=-1).cpu().numpy()\n",
    "\n",
    "print(\"Label probs:\", probs)  # prints: [[0.9927937  0.00421068 0.00299572]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ec9dd18",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_features @ text_features.t()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baed2a0c",
   "metadata": {},
   "source": [
    "# Basic ImageNet and CIFAR checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a819b31",
   "metadata": {},
   "outputs": [],
   "source": [
    "per_group_norm=False\n",
    "\n",
    "# Download the dataset\n",
    "cifar10 = CIFAR10(root=\"/tmp/\", transform=preprocess, download=True, train=False)\n",
    "loader = torch.utils.data.DataLoader(cifar10, batch_size=32, shuffle=True)\n",
    "# text_inputs = torch.stack([tokenizer(f\"a photo of a {c}\") for c in cifar10.classes]).cuda()\n",
    "\n",
    "\n",
    "text_inputs = torch.stack([tokenizer(f\"a photo of a {c}\") for c in cifar10.classes]).cuda()\n",
    "\n",
    "\n",
    "num_correct = 0\n",
    "num_seen = 0\n",
    "with torch.no_grad():\n",
    "    #  text_features = model.encode_text(text_inputs)\n",
    "    for imgs, targets in loader:\n",
    "        imgs = imgs.cuda()\n",
    "        targets = targets.cuda()\n",
    "        \n",
    "        # image_features = model.encode_image(imgs)\n",
    "        # results = model(imgs, text_inputs, lang_prompt_viz=True, sharded_computation=False)\n",
    "        # results = cc_model(imgs, text_inputs)\n",
    "        image_features = cc_model.encode_image(imgs)\n",
    "        text_features = cc_model.encode_text(text_inputs)\n",
    "\n",
    "        if per_group_norm:\n",
    "            image_features = image_features.view(image_features.shape[0], 64, -1)\n",
    "            text_features = text_features.view(text_features.shape[0], 64, -1)\n",
    "        image_features = image_features / image_features.norm(dim=-1, keepdim=True)\n",
    "        text_features = text_features / text_features.norm(dim=-1, keepdim=True)\n",
    "\n",
    "        image_features = image_features.flatten(start_dim=1)\n",
    "        text_features = text_features.flatten(start_dim=1)\n",
    "        # image_features = (image_features - image_features.mean(0)) / torch.sqrt(image_features.var(0))\n",
    "        # print(\"Using barlow\"); text_features = (text_features - text_features.mean(0)) / torch.sqrt(text_features.var(0)) \n",
    "        # image_features = (image_features - coco_image_mean) / torch.sqrt(coco_image_var)\n",
    "        # print(\"Using barlow\"); text_features = (text_features - coco_text_mean) / torch.sqrt(coco_text_var)  \n",
    "        \n",
    "    \n",
    "        similarity = (image_features @ text_features.T).softmax(dim=-1)\n",
    "        max_sims = similarity.max(dim=-1)[0]\n",
    "        # print(max_sims.min(), max_sims.max(), max_sims.mean())\n",
    "        num_correct += (similarity.argmax(dim=-1)==targets).sum().item()  \n",
    "        num_seen += imgs.shape[0]\n",
    "        curr_acc = num_correct / num_seen\n",
    "        print(f\"{num_seen} / {len(cifar10)} // Current Acc: {curr_acc}\")\n",
    "        if num_seen > 100000: break\n",
    "acc = num_correct / num_seen\n",
    "print(f\"(PARTIAL) Final Acc {acc}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b6734b2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Download the dataset\n",
    "cifar10 = CIFAR10(root=\"/tmp/\", transform=preprocess, download=True, train=False)\n",
    "loader = torch.utils.data.DataLoader(cifar10, batch_size=32, shuffle=True)\n",
    "# text_inputs = torch.stack([tokenizer(f\"a photo of a {c}\") for c in cifar10.classes]).cuda()\n",
    "\n",
    "\n",
    "text_inputs = torch.stack([tokenizer(f\"a photo of a {c}\") for c in cifar10.classes]).cuda()\n",
    "\n",
    "\n",
    "num_correct = 0\n",
    "num_seen = 0\n",
    "with torch.no_grad():\n",
    "    #  text_features = model.encode_text(text_inputs)\n",
    "    for imgs, targets in loader:\n",
    "        imgs = imgs.cuda()\n",
    "        targets = targets.cuda()\n",
    "        \n",
    "        # image_features = model.encode_image(imgs)\n",
    "        # results = model(imgs, text_inputs, lang_prompt_viz=True, sharded_computation=False)\n",
    "        # results = cc_model(imgs, text_inputs)\n",
    "        image_features = model.encode_image(imgs)\n",
    "        text_features = model.encode_text(text_inputs)\n",
    "\n",
    "        if per_group_norm:\n",
    "            image_features = image_features.view(image_features.shape[0], 64, -1)\n",
    "            text_features = text_features.view(text_features.shape[0], 64, -1)\n",
    "        image_features = image_features / image_features.norm(dim=-1, keepdim=True)\n",
    "        text_features = text_features / text_features.norm(dim=-1, keepdim=True)\n",
    "\n",
    "        image_features = image_features.flatten(start_dim=1)\n",
    "        text_features = text_features.flatten(start_dim=1)\n",
    "        # image_features = (image_features - image_features.mean(0)) / torch.sqrt(image_features.var(0))\n",
    "        # print(\"Using barlow\"); text_features = (text_features - text_features.mean(0)) / torch.sqrt(text_features.var(0)) \n",
    "        # image_features = (image_features - coco_image_mean) / torch.sqrt(coco_image_var)\n",
    "        # print(\"Using barlow\"); text_features = (text_features - coco_text_mean) / torch.sqrt(coco_text_var)  \n",
    "        \n",
    "    \n",
    "        similarity = (image_features @ text_features.T).softmax(dim=-1)\n",
    "        max_sims = similarity.max(dim=-1)[0]\n",
    "        # print(max_sims.min(), max_sims.max(), max_sims.mean())\n",
    "        num_correct += (similarity.argmax(dim=-1)==targets).sum().item()  \n",
    "        num_seen += imgs.shape[0]\n",
    "        curr_acc = num_correct / num_seen\n",
    "        print(f\"{num_seen} / {len(cifar10)} // Current Acc: {curr_acc}\")\n",
    "        if num_seen > 100000: break\n",
    "acc = num_correct / num_seen\n",
    "print(f\"(PARTIAL) Final Acc {acc}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10b41c8d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "imagenet = ImageNet(root=\"/export/share/datasets/vision/imagenet\", transform=preprocess, split='val')\n",
    "loader = torch.utils.data.DataLoader(imagenet, batch_size=32, num_workers=4, shuffle=True)\n",
    "text_inputs = torch.stack([tokenizer(f\"a photo of a {c}\") for c in imagenet.classes]).cuda()\n",
    "\n",
    "num_correct = 0\n",
    "num_seen = 0\n",
    "with torch.no_grad():\n",
    "    text_features = cc_model.encode_text(text_inputs)\n",
    "    if per_group_norm:\n",
    "        text_features = text_features.view(text_features.shape[0], 64, -1)\n",
    "    text_features = text_features / text_features.norm(dim=-1, keepdim=True)\n",
    "    text_features = text_features.flatten(start_dim=1)\n",
    "    # print(\"Using barlow\"); text_features = (text_features - text_features.mean(0)) / torch.sqrt(text_features.var(0))\n",
    "    # print(\"Using barlow\"); text_features = (text_features - coco_text_mean) / torch.sqrt(coco_text_var)  \n",
    "    \n",
    "    for imgs, targets in loader:\n",
    "        imgs = imgs.cuda()\n",
    "        targets = targets.cuda()\n",
    "        \n",
    "        image_features = cc_model.encode_image(imgs)\n",
    "        if per_group_norm:\n",
    "            image_features = image_features.view(image_features.shape[0], 64, -1)\n",
    "        image_features = image_features / image_features.norm(dim=-1, keepdim=True)\n",
    "        image_features = image_features.flatten(start_dim=1)\n",
    "        # print(\"using barlow\"); image_features = (image_features - image_features.mean(0)) / torch.sqrt(image_features.var(0))\n",
    "        # results = model(imgs, text_inputs, lang_prompt_viz=True, sharded_computation=False)\n",
    "        # image_features = results['image_embed']\n",
    "        # text_features = results['text_embed']\n",
    "        # image_features = (image_features - coco_image_mean) / torch.sqrt(coco_image_var) \n",
    "        \n",
    "        \n",
    "        similarity = (image_features @ text_features.T).softmax(dim=-1)\n",
    "        num_correct += (similarity.argmax(dim=-1)==targets).sum().item()  \n",
    "        max_sims = similarity.max(dim=-1)[0]\n",
    "        # print(max_sims.min(), max_sims.max(), max_sims.mean())\n",
    "        num_seen += imgs.shape[0]\n",
    "        curr_acc = num_correct / num_seen\n",
    "        print(f\"{num_seen} / {len(imagenet)} // Current Acc: {curr_acc}\")\n",
    "        if num_seen > 100000: break\n",
    "        \n",
    "print(f\"(PARTIAL) Final Acc {curr_acc}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c80ec693",
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity.shape"
   ]
  },
  {
   "cell_type": "raw",
   "id": "e35fbfd8",
   "metadata": {},
   "source": [
    "Baseline CC12M model w/ ViT is ALREADY at 44% CIFAR10 accuracy when at loss 2 at epoch 5\n",
    "And at ImageNEt loss of 18.8!!\n",
    "\n",
    "ResNet18 is worse @ epoch 5 w/ ~18% cifar and ~0.14 imagent\n",
    "    Difference is funny, this is @ epoch 5 w/ 2.6 loss\n",
    "    Epoch 10 CIFAR is 35% and loss is ~2.15, with ImageNEt at ~15%\n",
    "        DTD peformance is down to 0.05, so really seems like model is struggling\n",
    "     Epoch 35 only had 33% CIFAR accuracy, so that's on par with us! Really is parameter thing!\n",
    "         Did have 18% ImageNEt so slight edge there but nothing to get excited about\n",
    "         \n",
    "      This is a much better baseline to show to Nikhil\n",
    "     \n",
    "     So really need specialization, maybe even more flexible models than just the CNNs (not sure if those are better for low raining than low parameter)\n",
    "        \n",
    "     All signs point to lots of parameters with true specialization being key\n",
    "    \n",
    "    not sure how to interpret. We are seeing large dropoff, these #s are pretty similar to our expert-based models, but they're probably still improving a bit? Overall this seems to point and really needing the capacity of bigger models\n",
    "    \n",
    "    Wondering whether I should look @ small transformer instead of small visual? Does inductive bias help at all when getting overwhelmed?\n",
    "    Might just want to go with transformer to increase bang per buck for parameter compared to CNN\n",
    "\n",
    "So yes, MAD overfitting is most likely answer if loss is going down that much but not seeing imrpvoemetn\n",
    "\n",
    "Sooo..... Are we mad overfitting if our loss is that low with bad accuracy?\n",
    "\n",
    "Loss convergence rate is similar which is weird\n",
    "\n",
    "AH, note that our loss is for easier version of problem (bs=1024) so actually converg\n",
    "\n",
    "----\n",
    "epoch 35 of my orig CLIP model (both vanilla prompting)\n",
    "CIFAR 69.2%\n",
    "ImageNEt 28.2\n",
    "\n",
    "epoch 2 of parallel vision standard text\n",
    "got 24% CIFAR  definitely nontrivial signal\n",
    "Also got 7.3% on ImageNet\n",
    "\n",
    "Epoch 15\n",
    "30 % / 3%\n",
    "\n",
    "So definitely have some overfitting (probably from removing weight decay)\n",
    "\n",
    "Epoch 21\n",
    "0.292 CIFAR\n",
    "ImageNet: only 2%, so got worse\n",
    "\n",
    "\n",
    "Potential of overfitting here? Don't have a val set, also going to load earlier ckpts\n",
    "\n",
    "\n",
    "-----\n",
    "Epoch 15 of pretrained + WD (loss ~1.14)\n",
    "10.5% ImageNet\n",
    "36.7% CIFAR10\n",
    "\n",
    "Epoch 21 (loss 0.7688705921173096) \n",
    "CIFAR10 dropped to 30.9\n",
    "ImageNet down to \n",
    "\n",
    "Epoch 35 (loss 0.14)\n",
    "\n",
    "CIFAR 32.3%\n",
    "ImageNet still 9%\n",
    "\n",
    "Epoch 35 w/ lower LR and higher WD\n",
    "~29% CIFAR\n",
    "ImageNet ~11%\n",
    "\n",
    "\n",
    "-----\n",
    "\n",
    "Epoch 5 of viz paralle lstanadard text 64 model head size 64\n",
    "CIFAR ~34% ImageNet 10%\n",
    "    So this is already better than end of previous version w/ global norm on head size 8\n",
    "    Up with epoch 15 of pretrained + full WD for performance\n",
    "        If acc drops as well as time goes on then definitely need to add more regularization or oemthing\n",
    "        \n",
    "        \n",
    "Epoch 30\n",
    "CIFAR ~32%\n",
    "ImageNEt ~13%\n",
    "\n",
    "OK, DEFINTELY getting some overfitting here. Weird that CIFAR is worse than DTD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fee44dc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "imagenet = ImageNet(root=\"/export/share/datasets/vision/imagenet\", transform=preprocess, split='val')\n",
    "loader = torch.utils.data.DataLoader(imagenet, batch_size=32, num_workers=4, shuffle=True)\n",
    "text_inputs = torch.stack([tokenizer(f\"a photo of a {c}\") for c in imagenet.classes]).cuda()\n",
    "\n",
    "num_correct = 0\n",
    "num_seen = 0\n",
    "with torch.no_grad():\n",
    "    text_features = model.encode_text(text_inputs)\n",
    "    if per_group_norm:\n",
    "        text_features = text_features.view(text_features.shape[0], 64, -1)\n",
    "    text_features = text_features / text_features.norm(dim=-1, keepdim=True)\n",
    "    text_features = text_features.flatten(start_dim=1)\n",
    "    # print(\"Using barlow\"); text_features = (text_features - text_features.mean(0)) / torch.sqrt(text_features.var(0))\n",
    "    # print(\"Using barlow\"); text_features = (text_features - coco_text_mean) / torch.sqrt(coco_text_var)  \n",
    "    \n",
    "    for imgs, targets in loader:\n",
    "        imgs = imgs.cuda()\n",
    "        targets = targets.cuda()\n",
    "        \n",
    "        image_features = model.encode_image(imgs)\n",
    "        if per_group_norm:\n",
    "            image_features = image_features.view(image_features.shape[0], 64, -1)\n",
    "        image_features = image_features / image_features.norm(dim=-1, keepdim=True)\n",
    "        image_features = image_features.flatten(start_dim=1)\n",
    "        # print(\"using barlow\"); image_features = (image_features - image_features.mean(0)) / torch.sqrt(image_features.var(0))\n",
    "        # results = model(imgs, text_inputs, lang_prompt_viz=True, sharded_computation=False)\n",
    "        # image_features = results['image_embed']\n",
    "        # text_features = results['text_embed']\n",
    "        # image_features = (image_features - coco_image_mean) / torch.sqrt(coco_image_var) \n",
    "        \n",
    "        \n",
    "        similarity = (image_features @ text_features.T).softmax(dim=-1)\n",
    "        num_correct += (similarity.argmax(dim=-1)==targets).sum().item()  \n",
    "        max_sims = similarity.max(dim=-1)[0]\n",
    "        # print(max_sims.min(), max_sims.max(), max_sims.mean())\n",
    "        num_seen += imgs.shape[0]\n",
    "        curr_acc = num_correct / num_seen\n",
    "        print(f\"{num_seen} / {len(imagenet)} // Current Acc: {curr_acc}\")\n",
    "        if num_seen > 1000: break\n",
    "        \n",
    "print(f\"(PARTIAL) Final Acc {curr_acc}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b9564c4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cd182a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def dataset_constructor(name, split, transform):\n",
    "    \"\"\"\n",
    "    inputs\n",
    "        name: string that we'll run if/else on\n",
    "        split: \"train\" or \"val\" for now\n",
    "    \"\"\"\n",
    "    assert split in ['train', 'val']\n",
    "\n",
    "    # Need to figure out way to override the getitem call for strong weak augmentatin\n",
    "    if name == 'cifar10':\n",
    "        # Has class names\n",
    "        return datasets.CIFAR10(root=os.path.expanduser(\"~/.cache\"),\n",
    "                                download=True,\n",
    "                                train=(split=='train'),\n",
    "                               transform=transform)\n",
    "    elif name == 'cifar100':\n",
    "        # Has class names\n",
    "        return datasets.CIFAR100(root=os.path.expanduser(\"~/.cache\"),\n",
    "                                download=True,\n",
    "                                train=(split=='train'),\n",
    "                                transform=transform)\n",
    "\n",
    "    elif name == 'svhn':\n",
    "        # Doesn't have class names\n",
    "        return datasets.SVHN(root=os.path.expanduser(\"~/.cache\"),\n",
    "                             download=True,\n",
    "                             transform=transform,\n",
    "                             split=split)\n",
    "    elif name == 'gtsrb': #####\n",
    "        assert split=='val'\n",
    "        return datasets.ImageFolder(root=\"/export/share/bwallace/datasets/gtsrb/pytorch_format\",\n",
    "                                   transform=transform)\n",
    "        return datasets.GTSRB(root=os.path.expanduser(\"~/.cache\"), split='test', download=True)\n",
    "    elif name == 'food101':\n",
    "        # name is in .classes\n",
    "        assert split=='val'\n",
    "        return datasets.ImageFolder(root=\"/export/share/bwallace/datasets/food101/images/\",\n",
    "                   transform=transform)\n",
    "    elif name == 'merced':\n",
    "        # name is in .classes\n",
    "        assert split=='val'\n",
    "        return datasets.ImageFolder(root=\"/export/share/bwallace/datasets/UCMerced_LandUse/Images/\",\n",
    "                   transform=transform)\n",
    "    elif name == 'mnist':\n",
    "        assert split=='val'\n",
    "        return datasets.MNIST(root=os.path.expanduser(\"~/.cache\"), train=False, download=True, transform=transform)\n",
    "    elif name == 'flowers': ####\n",
    "        assert split=='val'\n",
    "        print(\"NOTE PROPER METRIC IS MEAN PER CLASS\")\n",
    "        return datasets.ImageFolder(root=\"/export/share/bwallace/datasets/flowers/test/\",\n",
    "                                    transform=transform)\n",
    "    elif name == 'aircraft': ####\n",
    "        assert split=='val'\n",
    "        print(\"NOTE PROPER METRIC IS MEAN PER CLASS\")\n",
    "        return datasets.ImageFolder(root=\"/export/share/bwallace/datasets/aircraft/data/test_pytorch_format/\",\n",
    "                                   transform=transform)\n",
    "    elif name == 'cars':\n",
    "        assert split=='val'\n",
    "        return datasets.ImageFolder(root=\"/export/share/datasets/vision/stanford_cars/car_data/test/\",\n",
    "                                   transform=transform)\n",
    "    elif name == 'eurosat':\n",
    "        assert split=='val'\n",
    "        return datasets.ImageFolder(root=\"/export/share/datasets/vision/euro_sat/2750/\",\n",
    "                                   transform=transform)\n",
    "    elif name == 'dtd':\n",
    "        # name is in .classes\n",
    "        return datasets.ImageFolder(root=f\"/export/share/bwallace/datasets/dtd/{split}\",\n",
    "                   transform=transform)\n",
    "    elif name == 'cub':\n",
    "        # name is in .classes\n",
    "        return datasets.ImageFolder(root=f\"/export/share/bwallace/datasets/CUB_2011_formatted/{split}\",\n",
    "                   transform=transform)\n",
    "    elif name == 'places365':\n",
    "        # class names works in \n",
    "        return datasets.Places365(\n",
    "                    root=\"/export/share/datasets/vision/Places365/\",\n",
    "                    split='train-standard' if split=='train' else split,\n",
    "                   transform=transform)\n",
    "\n",
    "    elif name == 'imagenet':\n",
    "        # Doesn't have class names, have separate call\n",
    "        return datasets.ImageNet(\n",
    "                    root=\"/export/share/datasets/vision/imagenet/\",\n",
    "                    split=split,\n",
    "                   transform=transform)\n",
    "    elif name == 'imagenet_val':\n",
    "        assert split=='train'\n",
    "        # Doesn't have class names, only use is having Imagenet val be trainset \n",
    "        return datasets.ImageNet(\n",
    "                    root=\"/export/share/datasets/vision/imagenet/\",\n",
    "                    split='val',\n",
    "                   transform=transform)\n",
    "    else:\n",
    "        raise NotImplementedError\n",
    "\n",
    "\n",
    "def get_imagenet_class_dict():\n",
    "        idx_to_word_id_and_name_tuple = json.load(open('imagenet_class_index.json'))\n",
    "        word_id_to_name_and_idx = {v[0]:(v[1],int(k))\n",
    "                           for k,v in idx_to_word_id_and_name_tuple.items()}\n",
    "        # e.g. n0023923939 to ('unicorn', 123)\n",
    "        return word_id_to_name_and_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd4b4410",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_on_dataset(model, dataset_name, transform,\n",
    "                   prompt_template=None,\n",
    "                   test_components_mode='standard',\n",
    "                   normalization='mask',\n",
    "                   component_grouping=1):\n",
    "    \"\"\"\n",
    "    Options for test components mode\n",
    "        standard : standard\n",
    "        test_all : scale through\n",
    "        random_all : As all but random sorting\n",
    "        \n",
    "    Normalization is whether to normalize from all components or just ones in computation\n",
    "    \n",
    "    \"\"\"\n",
    "    dataset = dataset_constructor(dataset_name, 'val', transform)\n",
    "    loader = torch.utils.data.DataLoader(dataset, batch_size=32, num_workers=4)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        \n",
    "        \n",
    "        if prompt_template is not None:\n",
    "            class_names = [prompt_template.format(c.replace(\"'\",\"\")) for c in dataset.classes]\n",
    "            # print(class_names)\n",
    "            text_inputs = torch.stack([tokenizer(txt) for txt in class_names]).cuda()\n",
    "            text_features = model.encode_text(text_inputs)\n",
    "            if normalization in ['mask', 'global']:\n",
    "                text_features = text_features / text_features.norm(dim=-1, keepdim=True)\n",
    "            elif normalization == 'grouped':\n",
    "                text_features = text_features.view(text_features.shape[0],\n",
    "                                                                -1,\n",
    "                                                                component_grouping)\n",
    "                text_features =  text_features / text_features.norm(dim=-1, keepdim=True)\n",
    "                text_features = text_features.flatten(start_dim=1)\n",
    "            else:\n",
    "                raise NotImplementedError\n",
    "        else:\n",
    "            with open('templates.json') as f:\n",
    "                all_templates = json.load(f)\n",
    "            templates = all_templates[dataset_name]\n",
    "            text_features = []\n",
    "            class_names = [c.replace(\"'\", \"\") for c in dataset.classes]\n",
    "            for class_name in class_names:\n",
    "                text_inputs = torch.stack([tokenizer(t.format(class_name)) for t in templates]).cuda()\n",
    "                class_features = model.encode_text(text_inputs)\n",
    "                # class_features = class_features / class_features.norm(dim=-1, keepdim=True)\n",
    "                if normalization in ['mask', 'global']:\n",
    "                    class_features = class_features / class_features.norm(dim=-1, keepdim=True)\n",
    "                elif normalization == 'grouped':\n",
    "                    class_features = class_features.view(class_features.shape[0],\n",
    "                                                                    -1,\n",
    "                                                                    component_grouping)\n",
    "                    class_features =  class_features / class_features.norm(dim=-1, keepdim=True)\n",
    "                    class_features = class_features.flatten(start_dim=1)\n",
    "                else:\n",
    "                    raise NotImplementedError\n",
    "                # print(class_features.shape)\n",
    "                class_features = class_features.mean(dim=0)\n",
    "                # print(class_features.shape)\n",
    "                if normalization in ['mask', 'global']:\n",
    "                    class_features = class_features / class_features.norm(dim=-1, keepdim=False)\n",
    "                    # text_features = text_features / text_features.norm(dim=-1, keepdim=True)\n",
    "                elif normalization == 'grouped':\n",
    "                    class_features = class_features.view(-1, component_grouping)\n",
    "                    class_features =  class_features / class_features.norm(dim=-1, keepdim=True)\n",
    "                    class_features = class_features.flatten()\n",
    "                else:\n",
    "                    raise NotImplementedError\n",
    "                # print(class_features.shape)\n",
    "                text_features.append(class_features)\n",
    "            text_features = torch.stack(text_features)\n",
    "            # print(text_features.shape)\n",
    "                \n",
    "        if test_components_mode in ['standard']:\n",
    "            sorted_feature_variances, sorted_variance_idx = None, None\n",
    "        elif test_components_mode in ['test_all', 'ensemble']:\n",
    "            # grouping will return indices in blocks (idx //component_grouping will be repeated numbers in blocks of component_grouping)\n",
    "            sorted_feature_variances, sorted_variance_idx = sorted_feature_component_variances(text_features,\n",
    "                                                                                              group_size=component_grouping)\n",
    "        elif test_components_mode in ['all_indiv']:\n",
    "            sorted_variance_idx = torch.arange(text_features.shape[1])\n",
    "        elif test_components_mode in ['random_all']:\n",
    "            # Want random to be grouped as well\n",
    "            sorted_feature_variances, sorted_variance_idx = sorted_feature_component_variances(text_features,\n",
    "                                                                                              group_size=component_grouping)\n",
    "            grouped_idx = torch.randperm(sorted_variance_idx.shape[0] // component_grouping)\n",
    "            sorted_variance_idx = torch.cat([torch.arange(component_grouping).cuda() + (component_grouping*i) for i in grouped_idx])\n",
    "            # print(sorted_variance_idx)\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "        \n",
    "        num_correct = 0 if test_components_mode =='standard' else torch.zeros(text_features.shape[1] // component_grouping)\n",
    "        num_seen = 0\n",
    "        for imgs, targets in tqdm(loader):\n",
    "            imgs = imgs.cuda()\n",
    "            targets = targets.cuda()\n",
    "\n",
    "            image_features = model.encode_image(imgs)\n",
    "            image_features = image_features / image_features.norm(dim=1, keepdim=True)\n",
    "            # results = model(imgs, text_inputs, lang_prompt_viz=True, sharded_computation=False)\n",
    "            # image_features = results['image_embed']\n",
    "            # text_features = results['text_embed']\n",
    "\n",
    "\n",
    "            if test_components_mode==\"standard\":\n",
    "                if normalization == 'grouped':\n",
    "                    image_features = image_features.view(image_features.shape[0],\n",
    "                                                                    -1,\n",
    "                                                                    component_grouping)\n",
    "                    image_features =  image_features / image_features.norm(dim=-1, keepdim=True)\n",
    "                    image_features = image_features.flatten(start_dim=1)\n",
    "                similarity = (image_features @ text_features.T).softmax(dim=-1)\n",
    "                num_correct += (similarity.argmax(dim=-1)==targets).sum().item() \n",
    "            elif test_components_mode in [\"test_all\", \"random_all\", \"ensemble\", \"all_indiv\"]:\n",
    "                for num_components in range(component_grouping, text_features.shape[1]+1,\n",
    "                                           component_grouping):\n",
    "                    \n",
    "                    start_idx = num_components - component_grouping if test_components_mode == 'all_indiv' else 0\n",
    "                    idx_to_use = sorted_variance_idx[start_idx:num_components]\n",
    "                    masked_text_features = text_features[:, idx_to_use]\n",
    "                    masked_image_features = image_features[:, idx_to_use]\n",
    "                    if normalization == 'mask':\n",
    "                        masked_text_features =  masked_text_features / masked_text_features.norm(dim=-1, keepdim=True)\n",
    "                        masked_image_features =  masked_image_features / masked_image_features.norm(dim=-1, keepdim=True)\n",
    "                    elif normalization == 'global':\n",
    "                        masked_image_features = masked_image_features / image_features.norm(dim=-1, keepdim=True)\n",
    "                    elif normalization == 'grouped':\n",
    "                        masked_text_features = masked_text_features.view(masked_text_features.shape[0],\n",
    "                                                                        -1,\n",
    "                                                                        component_grouping)\n",
    "                        masked_text_features =  masked_text_features / masked_text_features.norm(dim=-1, keepdim=True)\n",
    "                        masked_text_features = masked_text_features.flatten(start_dim=1)\n",
    "                        masked_image_features = masked_image_features.view(masked_image_features.shape[0],\n",
    "                                                                        -1,\n",
    "                                                                        component_grouping)\n",
    "                        masked_image_features =  masked_image_features / masked_image_features.norm(dim=-1, keepdim=True)\n",
    "                        masked_image_features = masked_image_features.flatten(start_dim=1)\n",
    "                    else:\n",
    "                        raise NotImplementedError\n",
    "                    if test_components_mode in [\"test_all\", \"random_all\", \"all_indiv\"]:\n",
    "                        masked_similarity = (masked_image_features @ masked_text_features.T).softmax(dim=-1)\n",
    "                        num_correct[(num_components//component_grouping)-1] += (masked_similarity.argmax(dim=-1)==targets).sum().item()\n",
    "                    elif test_components_mode == 'ensemble':\n",
    "                        # right now features are bs x (head_dim * num_components)\n",
    "                        # put as bs x num_components x head_dim\n",
    "                        # permute to num_components x bs x head_dim\n",
    "                        # BMM against num_components x head_dim x bs\n",
    "                        # get num_components bs x bs selection matrices\n",
    "                        # print(masked_image_features.shape)\n",
    "                        # print(masked_text_features.shape)\n",
    "                        num_texts = masked_text_features.shape[0]\n",
    "                        num_images = masked_image_features.shape[0]\n",
    "                        masked_image_features = masked_image_features.view(num_images, -1, component_grouping).permute(1, 0, 2)\n",
    "                        masked_text_features = masked_text_features.view(num_texts, -1, component_grouping).permute(1, 2, 0)\n",
    "                        per_component_sims = torch.bmm(masked_image_features, masked_text_features) # num_components x num_images x num_texts\n",
    "                        per_image_per_component_preds = per_component_sims.argmax(dim=-1) # num_components x num_images\n",
    "                        # print(per_image_per_component_preds)\n",
    "                        # print(per_image_per_component_preds.shape)\n",
    "                        ensemble_pred = per_image_per_component_preds.cpu().mode(0)[0]\n",
    "                        # print(ensemble_pred.shape, targets.shape)\n",
    "                        num_correct[(num_components//component_grouping)-1] += (ensemble_pred == targets.cpu()).sum().item()\n",
    "                        # BMM against bs x head_dim x num_components and get \n",
    "                        # bs x num_components x num_components\n",
    "            else:\n",
    "                 raise NotImplementedError\n",
    "            # max_sims = similarity.max(dim=-1)[0]\n",
    "            # print(max_sims.min(), max_sims.max(), max_sims.mean())\n",
    "            num_seen += imgs.shape[0]\n",
    "            curr_acc = num_correct / num_seen\n",
    "    return curr_acc\n",
    "\n",
    "\n",
    "def get_dataset_text_features(model, dataset_name, norm=True, print_class_names=False):\n",
    "    dataset = dataset_constructor(dataset_name, 'val', None)\n",
    "    class_names = [f\"a photo of a {c}\" for c in dataset.classes]\n",
    "    if print_class_names:\n",
    "            print(dataset.classes)\n",
    "    text_inputs = torch.stack([tokenizer(txt) for txt in class_names]).cuda()\n",
    "    with torch.no_grad():\n",
    "        text_features = model.encode_text(text_inputs)\n",
    "        if norm: text_features = text_features / text_features.norm(dim=1, keepdim=True)\n",
    "    return text_features\n",
    "\n",
    "def sorted_feature_component_variances(tensor, group_size=1):\n",
    "    # NOTE: sorting max first so [0] is intuitively important\n",
    "    feature_variances = tensor.var(dim=0)\n",
    "    grouped_feature_variances = feature_variances.view(-1, group_size).mean(1)\n",
    "    grouped_idx = grouped_feature_variances.argsort().flip(0)\n",
    "    # will be # components / group_size\n",
    "    idx = torch.cat([torch.arange(group_size).cuda() + (group_size*i) for i in grouped_idx])\n",
    "    return feature_variances[idx], idx\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d77f1647",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_ordered_vs_random_acc_samples(ordered_accs, random_acc_sampling,\n",
    "                                       baseline_ordered_accs=None, baseline_random_acc_sampling=None,\n",
    "                                       title=None):\n",
    "    random_accs = torch.mean(torch.stack(random_acc_sampling, dim=0), dim=0)\n",
    "    plt.figure()\n",
    "    plt.title(title)\n",
    "    plt.ylabel('accuracy')\n",
    "    plt.xlabel('# components used')\n",
    "    plt.plot(ordered_accs.detach().cpu().numpy(), label=\"Sorted Variances\")\n",
    "    plt.plot(random_accs.detach().cpu().numpy(), label=f\"Random (average {len(random_acc_sampling)} runs)\")\n",
    "    if baseline_ordered_accs is not None:\n",
    "        plt.plot(baseline_ordered_accs.detach().cpu().numpy(), label=\"Baseline sorted\")\n",
    "    if baseline_random_acc_sampling is not None:\n",
    "        plt.plot(torch.stack(baseline_random_acc_sampling).mean(0).detach().cpu().numpy(), label=f\"Baseline Random (average {len(baseline_random_acc_sampling)} runs)\")\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    \n",
    "    \n",
    "def generate_order_variance_plot(model, ds_name, preprocess, num_samples=2, group_size=1, normalization='mask',\n",
    "                                    baseline_model=None):\n",
    "    random_acc_sampling = [test_on_dataset(model, ds_name, preprocess,\n",
    "                                           test_components_mode='random_all',\n",
    "                                          component_grouping=group_size,\n",
    "                                          normalization=normalization) for _ in range(num_samples)]\n",
    "    ordered_accs = test_on_dataset(model, ds_name, preprocess,\n",
    "                                   test_components_mode='test_all', component_grouping=group_size,\n",
    "                                          normalization=normalization)\n",
    "    \n",
    "    baseline_random_acc_sampling = None\n",
    "    baseline_ordered_accs = None\n",
    "    if baseline_model is not None:\n",
    "        baseline_random_acc_sampling = [test_on_dataset(baseline_model, ds_name, preprocess,\n",
    "                                               test_components_mode='random_all',\n",
    "                                              component_grouping=group_size,\n",
    "                                              normalization=normalization) for _ in range(num_samples)]\n",
    "        baseline_ordered_accs = test_on_dataset(baseline_model, ds_name, preprocess,\n",
    "                                       test_components_mode='test_all', component_grouping=group_size,\n",
    "                                              normalization=normalization)\n",
    "    plot_ordered_vs_random_acc_samples(ordered_accs, random_acc_sampling, title=ds_name,\n",
    "                                      baseline_ordered_accs=baseline_ordered_accs,\n",
    "                                      baseline_random_acc_sampling=baseline_random_acc_sampling)\n",
    "    \n",
    "    \n",
    "def simple_ordered_plot(model, ds_name, preprocess, group_size=1, normalization='mask',\n",
    "                                    baseline_model=None,\n",
    "                       parameter_x_axis=False):\n",
    "    ordered_accs = test_on_dataset(model, ds_name, preprocess,\n",
    "                                   test_components_mode='test_all', component_grouping=group_size,\n",
    "                                          normalization=normalization)\n",
    "    \n",
    "    baseline_random_acc_sampling = None\n",
    "    baseline_ordered_accs = None\n",
    "    if baseline_model is not None:\n",
    "        baseline_performance = test_on_dataset(baseline_model, ds_name, preprocess,\n",
    "                                       test_components_mode='standard', component_grouping=group_size,\n",
    "                                              normalization=normalization)\n",
    "    plt.figure()\n",
    "    plt.title(ds_name)\n",
    "    plt.ylabel('accuracy')\n",
    "    plt.plot(ordered_accs.detach().cpu().numpy(), label=\"Experts CLIP (Ours)\")\n",
    "    if baseline_performance is not None:\n",
    "        plt.axhline(baseline_performance, label=\"Baseline CLIP (CC12M) Performance\",\n",
    "                   linestyle='dashed',\n",
    "                   color='gray')\n",
    "    plt.legend()\n",
    "    if parameter_x_axis:\n",
    "        plt.xlabel(\"Percentage of parameters used\")\n",
    "        num_baseline_params = sum([p.numel() for p in baseline_model.visual.parameters()])\n",
    "        total_new_params = sum([p.numel() for p in model.visual.parameters()])\n",
    "        locs, labels = plt.xticks()\n",
    "        # locs are # of groups being used out of num_models\n",
    "        num_models = model.visual[-2].groups\n",
    "        labels = [f\"{int(100*(loc/num_models)*(total_new_params/num_baseline_params))}%\"\n",
    "                 for loc in locs]\n",
    "        plt.xticks(locs, labels)\n",
    "    else:\n",
    "        plt.xlabel('# groups used')\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03675007",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2134c831",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(cc_model, 'dtd', preprocess)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5692cc42",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'dtd', preprocess, test_components_mode='test_all',\n",
    "               component_grouping=64, normalization='grouped')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2171de17",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'dtd', preprocess, test_components_mode='all_indiv',\n",
    "               component_grouping=64, normalization='mask')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "048965ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'dtd', preprocess, test_components_mode='ensemble',\n",
    "               component_grouping=64, normalization='grouped')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04597cbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'aircraft', preprocess, test_components_mode='test_all',\n",
    "               component_grouping=64, normalization='grouped')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "247de749",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'cifar10', preprocess, test_components_mode='all_indiv',\n",
    "               component_grouping=64, normalization='mask')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7734c818",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'cifar10', preprocess, test_components_mode='test_all',\n",
    "               component_grouping=64, normalization='mask')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6c08a45",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'aircraft', preprocess, test_components_mode='all_indiv',\n",
    "               component_grouping=64, normalization='mask')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e043140",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'eurosat', preprocess, test_components_mode='all_indiv',\n",
    "               component_grouping=64, normalization='mask')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c09709d2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dfc1b52",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'cifar100', preprocess, test_components_mode='test_all',\n",
    "               component_grouping=64, normalization='grouped')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6afb574",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'cub', preprocess, test_components_mode='standard',\n",
    "               component_grouping=64, normalization='grouped')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "130c45cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'cub', preprocess, test_components_mode='test_all',\n",
    "               component_grouping=64, normalization='grouped')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd67cc35",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'dtd', preprocess, test_components_mode='random_all',\n",
    "               component_grouping=64, normalization='grouped')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2601f29b",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_on_dataset(model, 'cifar10', preprocess, test_components_mode='test_all',\n",
    "               component_grouping=64, normalization='grouped')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2713f20",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(model, 'dtd', preprocess, group_size=64, normalization='grouped', baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a2f4597",
   "metadata": {},
   "outputs": [],
   "source": [
    "simple_ordered_plot(model, 'dtd', preprocess, group_size=64, normalization='grouped', baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "500f123e",
   "metadata": {},
   "outputs": [],
   "source": [
    "simple_ordered_plot(model, 'dtd', preprocess, group_size=8, normalization='mask', baseline_model=cc_model,\n",
    "                   parameter_x_axis=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "196434a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "for ds in ds_names:\n",
    "    if ds=='mnist': continue\n",
    "    simple_ordered_plot(model, ds, preprocess, group_size=8, normalization='mask', baseline_model=cc_model,\n",
    "                   parameter_x_axis=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23f56167",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(model, 'mnist', grayscale_preprocess, group_size=8, global_norm=False, baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a6aa8f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(model, 'eurosat', preprocess, group_size=8, global_norm=False, baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61c32d05",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(model, 'cars', preprocess, group_size=8, global_norm=False, baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14be850b",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(model, 'food101', preprocess, group_size=8, global_norm=False, baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12260507",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(model, 'cifar10', preprocess, group_size=8, global_norm=False, baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "311f4faa",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(model, 'cifar100', preprocess, group_size=8, global_norm=False, baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9162893d",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(model, 'cub', preprocess, group_size=8, global_norm=False, baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64bb88cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(model, 'gtsrb', preprocess, group_size=8, global_norm=False, baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b10602c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(model, 'flowers', preprocess, group_size=64, global_norm=False, baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d054e633",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(model, 'aircraft', preprocess, group_size=8, global_norm=False, baseline_model=cc_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6527e0b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_order_variance_plot(cc_model, 'aircraft', preprocess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fcdf2eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_names = ['dtd', 'mnist', 'eurosat', 'cars', 'food101', 'cifar10',\n",
    "           'cifar100', 'cub', 'gtsrb', 'flowers', 'aircraft']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16308f71",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for ds_name in ds_names:\n",
    "    for net in [model, cc_model]:\n",
    "        print(ds_name)\n",
    "        print(test_on_dataset(net, ds_name,\n",
    "                              grayscale_preprocess if ds_name=='mnist' else preprocess))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52481c06",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9903d2de",
   "metadata": {},
   "outputs": [],
   "source": [
    "for ds_name in ds_names:\n",
    "    print(ds_name, len(dataset_constructor(ds_name, 'val', preprocess)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b1e7313",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bc64f30",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d34e4bf",
   "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}