finite_differences_medium_greyscale.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f59161e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import time\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.backends.backend_pdf import PdfPages\n",
    "from collections import defaultdict\n",
    "\n",
    "import color_regions, network, visualizations, utils\n",
    "from color_regions import *\n",
    "from network import *\n",
    "from visualizations import *\n",
    "from utils import *\n",
    "from hooks import *\n",
    "\n",
    "torch.backends.cudnn.benchmark = True\n",
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6613c11b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up autoreloading of shared code\n",
    "%load_ext autoreload\n",
    "%autoreload 1\n",
    "%aimport color_regions,network,visualizations,utils,hooks\n",
    "%aimport"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f4eadc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose(\n",
    "     [transforms.ToTensor()])#,\n",
    "    #transforms.Normalize((0.5), (0.5))])\n",
    "\n",
    "batch_size = 128  # seems to be the fastest batch size\n",
    "train_indices = (0, 200_000) # size of training set\n",
    "valid_indices = (1_250_000, 1_255_000)\n",
    "test_indices = (260_000, 310_000)\n",
    "\n",
    "def color_classifier(color):\n",
    "    if color <= 100:  # medium difficulty (width = 75)\n",
    "        return 0\n",
    "    if 100 < color <= 150:  # hard difficulty (width = 50)\n",
    "        return 1\n",
    "    if 150 < color <= 200:  # hard difficulty (width = 50)\n",
    "        return 2\n",
    "    if 200 < color:  # hard difficulty (width = 50)\n",
    "        return 1\n",
    "critical_color_values = [100, 150, 200]\n",
    "\n",
    "def set_loader_helper(indices):\n",
    "    data_set = ColorDatasetGenerator(color_classifier=color_classifier,\n",
    "                                    image_indices=indices,\n",
    "                                    transform=transform,\n",
    "                                    color_range=(25, 250),\n",
    "                                    noise_size=(1,7),\n",
    "                                    num_classes=3,\n",
    "                                    size=128,\n",
    "                                    radius=(128//6, 128//3))\n",
    "    loader = torch.utils.data.DataLoader(data_set, batch_size=batch_size, \n",
    "                                          shuffle=True, num_workers=4, pin_memory=True)\n",
    "    return data_set, loader\n",
    "train_set, train_loader = set_loader_helper(train_indices)\n",
    "valid_set, valid_loader = set_loader_helper(valid_indices)\n",
    "test_set, test_loader = set_loader_helper(test_indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9d39721",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the \"medium\" task\n",
    "color_probe = np.linspace(0, 255, 255)\n",
    "color_class = [color_classifier(x) for x in color_probe]\n",
    "plt.plot(color_probe, color_class)\n",
    "plt.xlabel(\"Color\")\n",
    "plt.yticks([0, 1, 2])\n",
    "plt.ylabel(\"Class\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4ff8c51",
   "metadata": {},
   "outputs": [],
   "source": [
    "small_net = ResNet([[16, 3, 2],  # num_channels (input and output), kernel_size, stride\n",
    "                    [32, 3, 2],\n",
    "                    [64, 3, 2]], 3, [128, 128, 1], \n",
    "                   \"small_net_noise_medium_grey.dict\", fc_layers=[32]).to(device)\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "small_optim = torch.optim.Adam(small_net.parameters())\n",
    "print(small_net.num_params())\n",
    "small_net.load_model_state_dict(optim=small_optim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61abb7d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "small_net2 = ResNet([[16, 3, 2],  # num_channels (input and output), kernel_size, stride\n",
    "                    [32, 3, 2],\n",
    "                    [64, 3, 2]], 3, [128, 128, 1], \n",
    "                   \"small_net_noise_medium_grey_2.dict\", fc_layers=[32]).to(device)\n",
    "# if 0:\n",
    "#     small_net2.save_model_state_dict(\"small_net_init_noise_medium_grey.dict\")\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "small_optim2 = torch.optim.Adam(small_net2.parameters())\n",
    "print(small_net2.num_params())\n",
    "small_net2.load_model_state_dict(optim=small_optim2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f9f2c98",
   "metadata": {},
   "outputs": [],
   "source": [
    "tiny_net_med = ResNet([[2, 3, 4],  # num_channels (input and output), kernel_size, stride\n",
    "                   [6, 3, 4]], 3, [128, 128, 1], \n",
    "                   \"tiny_net_noise_medium_grey.dict\", fc_layers=[32]).to(device)\n",
    "# if 0:\n",
    "#     tiny_net_med.save_model_state_dict(\"tiny_net_noise_init_medium_grey.dict\")\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "tiny_optim_med = torch.optim.Adam(tiny_net_med.parameters())\n",
    "print(tiny_net_med.num_params())\n",
    "tiny_net_med.load_model_state_dict(optim=tiny_optim_med)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95cce16b",
   "metadata": {},
   "outputs": [],
   "source": [
    "tiny_net_med2 = ResNet([[2, 3, 4],  # num_channels (input and output), kernel_size, stride\n",
    "                   [6, 3, 4]], 3, [128, 128, 1], \n",
    "                   \"tiny_net_noise_medium_grey_2.dict\", fc_layers=[32]).to(device)\n",
    "tiny_net_med2.load_model_state_dict(\"tiny_net_noise_init_medium_grey.dict\")\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "tiny_optim_med2 = torch.optim.Adam(tiny_net_med2.parameters())\n",
    "print(tiny_net_med2.num_params())\n",
    "tiny_net_med2.load_model_state_dict(optim=tiny_optim_med2)\n",
    "# didnt actually train this one for some reason"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c90a760",
   "metadata": {},
   "outputs": [],
   "source": [
    "tiny_net_med3 = ResNet([[2, 3, 4],  # num_channels (input and output), kernel_size, stride\n",
    "                   [6, 3, 4]], 3, [128, 128, 1], \n",
    "                   \"tiny_net_noise_medium_grey_3.dict\", fc_layers=[32]).to(device)\n",
    "tiny_net_med3.load_model_state_dict(\"tiny_net_noise_init_medium_grey.dict\")\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "tiny_optim_med3 = torch.optim.Adam(tiny_net_med3.parameters())\n",
    "print(tiny_net_med3.num_params())\n",
    "tiny_net_med3.load_model_state_dict(optim=tiny_optim_med3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eea99043",
   "metadata": {},
   "outputs": [],
   "source": [
    "unstrided_net = ResNet([[16, 3, 1],  # num_channels (input and output), kernel_size, stride\n",
    "                        [32, 3, 1],\n",
    "                        [64, 3, 2]], 3, [128, 128, 1], \n",
    "                   \"unstrided_net_noise_medium_grey.dict\", fc_layers=[32]).to(device)\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "unstrided_optim = torch.optim.Adam(unstrided_net.parameters())\n",
    "print(unstrided_net.num_params())\n",
    "unstrided_net.load_model_state_dict(optim=unstrided_optim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b35a5b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_gram = visualizations.fc_conv_feature_angles(unstrided_net, \n",
    "                            \"fully_connected.0.act_func\", num_embed=50, normalize=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c209019",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualizations.display_fc_conv_grams(feature_gram, selection=[1,4,17])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee214b09",
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_gram.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07a295d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "imshow_centered_colorbar(np.arcsinh(feature_gram[5]), cmap=\"bwr\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed0b4cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_net = ResNet([[16, 3, 1],  # num_channels (input and output), kernel_size, stride\n",
    "                        [32, 3, 1],\n",
    "                        [64, 3, 2]], 3, [128, 128, 1], \n",
    "                   \"random_init_net_noise_medium_grey.dict\", fc_layers=[32]).to(device)\n",
    "# if 0:  # changed right away, so its good\n",
    "#     random_net.save_model_state_dict()\n",
    "random_net.path = \"random_net_noise_medium_grey.dict\"\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "random_optim = torch.optim.Adam(random_net.parameters())\n",
    "print(random_net.num_params())\n",
    "random_net.load_model_state_dict(optim=random_optim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24b7fbb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_net2 = ResNet([[16, 3, 1],  # num_channels (input and output), kernel_size, stride\n",
    "                        [32, 3, 1],\n",
    "                        [64, 3, 2]], 3, [128, 128, 1], \n",
    "                   \"random_init_net_noise_medium_grey.dict\", fc_layers=[32]).to(device)\n",
    "random_net2.load_model_state_dict()  # start from same initialization\n",
    "random_net2.path = \"random_net_noise_medium_grey_2.dict\"\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "random_optim2 = torch.optim.Adam(random_net2.parameters())\n",
    "print(random_net2.num_params())\n",
    "random_net2.load_model_state_dict(optim=random_optim2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "531fc3e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_net3 = ResNet([[16, 3, 1],  # num_channels (input and output), kernel_size, stride\n",
    "                      [32, 3, 1],\n",
    "                      [64, 3, 2]], 3, [128, 128, 1], \n",
    "                   \"random_init_net_noise_medium_grey.dict\", fc_layers=[32]).to(device)\n",
    "random_net3.load_model_state_dict()  # start from same initialization\n",
    "random_net3.path = \"random_net_noise_medium_grey_3.dict\"\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "random_optim3 = torch.optim.Adam(random_net3.parameters())\n",
    "print(random_net3.num_params())\n",
    "random_net3.load_model_state_dict(optim=random_optim3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44484cba",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_net4 = ResNet([[16, 3, 1],  # num_channels (input and output), kernel_size, stride\n",
    "                      [32, 3, 1],\n",
    "                      [64, 3, 2]], 3, [128, 128, 1], \n",
    "                   \"random_init_net_noise_medium_grey.dict\", fc_layers=[32]).to(device)\n",
    "random_net4.load_model_state_dict()  # start from same initialization\n",
    "random_net4.path = \"random_net_noise_medium_grey_4.dict\"\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "random_optim4 = torch.optim.Adam(random_net4.parameters())\n",
    "print(random_net4.num_params())\n",
    "random_net4.load_model_state_dict(optim=random_optim4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "526bb251",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_net5 = ResNet([[16, 3, 1],  # num_channels (input and output), kernel_size, stride\n",
    "                      [32, 3, 1],\n",
    "                      [64, 3, 2]], 3, [128, 128, 1], \n",
    "                   \"random_net_noise_medium_grey_5.dict\", fc_layers=[32]).to(device)\n",
    "random_net5.load_model_state_dict(\"random_init_net_noise_medium_grey.dict\")  # start from same initialization\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "random_optim5 = torch.optim.Adam(random_net5.parameters())\n",
    "print(random_net5.num_params())\n",
    "random_net5.load_model_state_dict(optim=random_optim5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98ea1ea7",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_net6 = ResNet([[16, 3, 1],  # num_channels (input and output), kernel_size, stride\n",
    "                      [32, 3, 1],\n",
    "                      [64, 3, 2]], 3, [128, 128, 1], \n",
    "                   \"random_net_noise_medium_grey_6.dict\", fc_layers=[32]).to(device)\n",
    "random_net6.load_model_state_dict(\"random_init_net_noise_medium_grey.dict\")  # start from same initialization\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "random_optim6 = torch.optim.Adam(random_net6.parameters())\n",
    "print(random_net6.num_params())\n",
    "random_net6.load_model_state_dict(optim=random_optim6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc11393c",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_net7 = ResNet([[16, 3, 1],  # num_channels (input and output), kernel_size, stride\n",
    "                      [32, 3, 1],\n",
    "                      [64, 3, 2]], 3, [128, 128, 1], \n",
    "                   \"random_net_noise_medium_grey_7.dict\", fc_layers=[32]).to(device)\n",
    "loss_func = nn.CrossEntropyLoss()  # dont start from same initialization\n",
    "random_optim7 = torch.optim.Adam(random_net7.parameters())\n",
    "print(random_net7.num_params())\n",
    "random_net7.load_model_state_dict(optim=random_optim7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "924d950f",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train(random_net7, random_optim7, loss_func, 100, train_loader, valid_loader, device=device,\n",
    "                track_stat=final_activation_tracker, summarize=count_logit_usage, \n",
    "                log_file=\"logs/random_net7.pkl\", test_loader=test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be1a6c13",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_results(results, \"RandomNet7\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65c2f2e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train(random_net6, random_optim6, loss_func, 100, train_loader, valid_loader, device=device,\n",
    "                track_stat=final_activation_tracker, summarize=count_logit_usage, \n",
    "                log_file=\"logs/random_net6.pkl\", test_loader=test_loader)\n",
    "\n",
    "# double check the correlation computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d53a0c44",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_results(results, \"RandomNet6\", size=0.2, alpha=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13f424a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train(random_net5, random_optim5, loss_func, 100, train_loader, valid_loader, device=device,\n",
    "                track_stat=final_activation_tracker, summarize=count_logit_usage, \n",
    "                log_file=\"logs/random_net5.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f54e64ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_results(results, \"RandomNet5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f87d81a",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_corr_grad_logit_info(results[-2], \"RandomNet5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e993ca19",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train(tiny_net_med3, tiny_optim_med3, loss_func, 100, train_loader, valid_loader, device=device,\n",
    "                track_stat=final_activation_tracker, summarize=count_logit_usage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66528f6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_results(results, \"TinyNetMedium3\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74110658",
   "metadata": {},
   "outputs": [],
   "source": [
    "network.plot_corr_grad_logit_info(results[-2], \"TinyNetMedium3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99547bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "tiny_results = train(tiny_net_med, tiny_optim_med, loss_func, 100, train_loader, valid_loader, device=device,\n",
    "                track_stat=final_activation_histogram, summarize=count_logit_usage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ec6a243",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_results(tiny_results, \"TinyNetMedium\") # initial is 76998"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9649f43",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train(small_net2, small_optim2, loss_func, 100, train_loader, valid_loader, device=device,\n",
    "                track_stat=final_activation_histogram, summarize=count_logit_usage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "768dc743",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_results(results, \"SmallNet2\") # initial is 36031"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "543c1973",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train(random_net4, random_optim4, loss_func, 100, train_loader, valid_loader, device=device,\n",
    "                track_stat=final_activation_histogram, summarize=count_logit_usage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a627193",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_results(results, \"RandomNet4\") # starts to increase num_logits used when tr_loss drops below 3\n",
    "                   # initial num_logits 61498\n",
    "                   # doesnt seem to directly correspond to generalization capacity\n",
    "                   # (overfitting, EMC curve)\n",
    "# human loss curve = concave\n",
    "# AI should be that too\n",
    "\n",
    "# bad hidden units could be same across inputs or amplifying noise in input\n",
    "# already converged and happy (low gradient)\n",
    "# not converged but still ok (high grad)\n",
    "# completely useless, gradient is 0 again\n",
    "\n",
    "# correlation with outputs to check these cases and magnitude of gradient\n",
    "# do correlation before ReLU so there is at least some signal\n",
    "# guess: start with useful units, \n",
    "# try more times\n",
    "\n",
    "# somewhat related to grokking?\n",
    "# fourier components, need small data to get first components, need more to get later ones\n",
    "# => physicist!\n",
    "# also inverse graphics (learning model from image joshua b tenenbaum, then do physics on it)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b011920d",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train(random_net3, random_optim3, loss_func, 200, train_loader, valid_loader, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6310f4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train(random_net2, random_optim2, loss_func, 200, train_loader, valid_loader, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01645ad2",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train(random_net, random_optim, loss_func, 200, train_loader, valid_loader, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f46fa421",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train(unstrided_net, unstrided_optim, loss_func, 200)\n",
    "# never actually reached the tr_loss < 3 regime, so wouldnt have started to increase \n",
    "# its num_logits_used!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40afa409",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = train(small_net, small_optim, loss_func, 200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c63984f",
   "metadata": {},
   "outputs": [],
   "source": [
    "rate_distribution(unstrided_net, valid_loader, valid_set,\n",
    "                  critical_values=critical_color_values, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b17f7269",
   "metadata": {},
   "outputs": [],
   "source": [
    "response_graph(unstrided_net, valid_set, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b5987e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "unstrided_net.eval()\n",
    "for sample in valid_loader:\n",
    "    #explain_img, explain_target_logit, *__ = valid_set.generate_one()\n",
    "    imgs = sample[\"image\"].to(device).float()\n",
    "    labels = sample[\"label\"].to(device).float()\n",
    "    model_outpt = unstrided_net(imgs)\n",
    "    print(model_outpt.argmax(dim=1))\n",
    "    print(labels.argmax(dim=1))\n",
    "    print(correct(model_outpt, labels))\n",
    "    break\n",
    "    #plt.imshow(explain_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de51df40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# to further test the \"using 1 image => bad batchnorm estimates\" lets do the same test\n",
    "# but instead we will average over a sample of responses\n",
    "small_net.eval() # very important!\n",
    "stack_size = 32\n",
    "sampled_indices = 1_250_000 + np.random.choice(1000, stack_size, replace=False)\n",
    "total_images = stack_size * 255\n",
    "correct_num = 0\n",
    "with torch.no_grad():\n",
    "    counterfactual_color_values = np.linspace(0, 255, 255)\n",
    "    responses = []\n",
    "    for color in tqdm(counterfactual_color_values):\n",
    "        stacked_generated_img = []\n",
    "        for sampled_index in sampled_indices:\n",
    "            np.random.seed(sampled_index)\n",
    "            generated_img, lbl, *__ = valid_set.generate_one(set_color=color)\n",
    "            stacked_generated_img.append(generated_img)\n",
    "        stacked_generated_img = np.array(stacked_generated_img).transpose(0, 3, 1, 2)\n",
    "        generated_img = torch.tensor(stacked_generated_img).to(device).float()\n",
    "        response = small_net(generated_img, logits=True)\n",
    "        stacked_lbl = torch.tensor(np.repeat(np.expand_dims(lbl, 0), stack_size, axis=0)).to(device)\n",
    "        correct_num += correct(response, stacked_lbl).sum()\n",
    "        responses.append(np.squeeze(response.cpu().numpy()).mean(axis=0))\n",
    "print(correct_num/total_images, \"total accuracy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a8a4304",
   "metadata": {},
   "outputs": [],
   "source": [
    "responses = np.arcsinh(np.array(responses))  # this graph is quite robust to changes in batch size\n",
    "for output_logit in range(responses.shape[1]): # if we do .eval(), but varies if we do .train()\n",
    "    plt.plot(counterfactual_color_values, responses[:, output_logit], label=f\"class {output_logit}\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"Color value\")\n",
    "plt.ylabel(\"Network output logit\")\n",
    "plt.vlines([100, 150], np.min(responses), np.max(responses), linewidth=0.8,\n",
    "           colors=\"r\", label=\"decision boundary\", # probably because .train() in this case actually gives biased estimates because all the colors are the same\n",
    "           linestyles=\"dashed\")  # logit graphs look bad if doing .train(), and accuracy is lower?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a01fdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_net.eval()   # ---> without this line, it fails, especially with small colors\n",
    "with torch.no_grad():  # => batchnorm updates are very inaccurate if just one image\n",
    "    idx = 1_250_026   # => network expects batchnorm updates to basically be exactly in the \"middle\" 127\n",
    "    print(valid_set[idx])  # => fix the logit response graph by requiring it to be in eval mode\n",
    "    print(torch.softmax(res_net(torch.unsqueeze(valid_set[idx][\"image\"], 0).to(device).float()), 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc75f90f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# to test the above hypothesis, if we just stack the same image a bunch of times, and do .train()\n",
    "# the estimates should still be bad because batchnorm estimates would be just as bad as with \n",
    "# a single image in the batch\n",
    "res_net.train()\n",
    "idx = 1_250_024  # => hypothesis seems to be confirmed\n",
    "test_image = valid_set[idx][\"image\"].numpy()\n",
    "stacked_test = np.repeat(np.expand_dims(test_image, 0), 32, axis=0)\n",
    "print(valid_set[idx])\n",
    "print(torch.softmax(res_net(torch.tensor(stacked_test).to(device).float()), 1)[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "493f1f17",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(np.squeeze(generated_img.cpu().numpy()), cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcb9a2a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_net.train()\n",
    "with torch.no_grad():\n",
    "    for i, sample in enumerate(valid_loader):\n",
    "        imgs = sample[\"image\"].to(device).float()\n",
    "        labels = sample[\"label\"].to(device).float()\n",
    "        print(sample[\"color\"])\n",
    "        outputs = res_net(imgs)\n",
    "        print(loss_func(outputs, labels).item())\n",
    "        print(torch.argmax(labels,dim=1), torch.argmax(outputs, dim=1))\n",
    "        print(correct(outputs, labels).sum().item())\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "362c7276",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(500_001)\n",
    "explain_img, lbl, *_ = valid_set.generate_one()\n",
    "heat_map = finite_differences_map(small_net, valid_set, lbl.argmax(), explain_img, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "725c5459",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(explain_img, cmap=\"gray\")\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(heat_map, cmap=\"bwr\", interpolation=\"bilinear\")\n",
    "plt.colorbar()\n",
    "# generated with strides = 2 everywhere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cc1e0d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([explain_img, heat_map], transpose=False, colorbar=True)\n",
    "# generated with strides = 1, strides = 8 for lats layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c9c62d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_ids = [20_000, 25_000, 30_000, 600_000, 600_001]\n",
    "heat_maps = []\n",
    "explain_imgs = []\n",
    "for i, image_id in enumerate(image_ids):\n",
    "    np.random.seed(image_id)\n",
    "    explain_img_i, target_i, *__ = valid_set.generate_one()\n",
    "    heat_map_i = finite_differences_map(small_net, valid_set, target_i.argmax(), explain_img_i)\n",
    "    heat_maps.append(heat_map_i)\n",
    "    explain_imgs.append(explain_img_i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78ff450a",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([explain_imgs, heat_maps], transpose=True, colorbar=True)\n",
    "# generated with strides=2 everywhere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65bdc89c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([explain_imgs, heat_maps], transpose=True, colorbar=True)\n",
    "# generated with strides = 1, strides = 8 for lats layer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38f8060b",
   "metadata": {},
   "source": [
    "What if we run the same experiment, but cheat with a prior on pixel values that we know *should* be informative to the output logit, namely values closest to the decision boundary?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1f1e957",
   "metadata": {},
   "outputs": [],
   "source": [
    "unfair_prior = np.array([90, 110, 140, 160])  #  close to the critical values of 100, 150\n",
    "unfair_heat_maps = []\n",
    "plt.figure(figsize=(12, 5*len(image_ids)))\n",
    "for i, image_id in enumerate(image_ids):\n",
    "    np.random.seed(image_id)\n",
    "    explain_img_i, target_i, __ = valid_set.generate_one()\n",
    "    unfair_map_i = finite_differences_map(res_net, valid_set, target_i.argmax(), explain_img_i, unfairness=\"unfair\", values_prior=unfair_prior)\n",
    "    unfair_heat_maps.append(unfair_map_i)\n",
    "    plt.subplot(len(image_ids), 3, 3*i+1)\n",
    "    plt.imshow(explain_img_i, cmap=\"gray\")\n",
    "    plt.subplot(len(image_ids), 3, 3*i+2)\n",
    "    heat_max = np.max(abs(unfair_map_i))\n",
    "    plt.imshow(unfair_map_i, cmap=\"bwr\", vmax=heat_max, vmin=-heat_max)\n",
    "    plt.colorbar(shrink=0.5)\n",
    "    plt.subplot(len(image_ids), 3, 3*i+3)\n",
    "    heat_max = np.max(abs(heat_maps[i]))\n",
    "    plt.imshow(heat_maps[i], cmap=\"bwr\", vmax=heat_max, vmin=-heat_max)\n",
    "    plt.colorbar(shrink=0.5)\n",
    "plt.show()  # => very similar results, but with a 4x speedup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "094020d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# regenerate unfair FD maps\n",
    "unfair_prior = np.array([90, 110, 140, 160])  #  close to the critical values of 100, 150\n",
    "image_ids = [20_000, 25_000, 30_000, 600_000, 600_001, 227_662, 998_102, 106_758]\n",
    "unfair_heat_maps = []\n",
    "#plt.figure(figsize=(12, 5*len(image_ids)))\n",
    "for i, image_id in enumerate(image_ids):\n",
    "    np.random.seed(image_id)\n",
    "    explain_img_i, target_i, __ = valid_set.generate_one()\n",
    "    unfair_map_i = finite_differences_map(res_net, valid_set, target_i.argmax(), explain_img_i, unfairness=\"unfair\", values_prior=unfair_prior)\n",
    "    unfair_heat_maps.append(unfair_map_i)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f292125",
   "metadata": {},
   "source": [
    "We can do even better by taking the \"closest value in a different class\" for our prior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb5440d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "very_unfair_heat_maps = []\n",
    "plt.figure(figsize=(20, 5*len(image_ids)))\n",
    "for i, image_id in enumerate(image_ids):\n",
    "    np.random.seed(image_id)\n",
    "    explain_img_i, target_i, color_i = valid_set.generate_one()\n",
    "    very_unfair_map_i = finite_differences_map(res_net, valid_set, target_i.argmax(), explain_img_i, unfairness=\"very unfair\", values_prior=[100, 150])\n",
    "    very_unfair_heat_maps.append(very_unfair_map_i)\n",
    "    \n",
    "    plt.subplot(len(image_ids), 5, 5*i+1)\n",
    "    if i == 0:\n",
    "        plt.title(\"Image\")\n",
    "    plt.imshow(explain_img_i, cmap=\"gray\")\n",
    "    \n",
    "    plt.subplot(len(image_ids), 5, 5*i+2)\n",
    "    if i == 0:\n",
    "        plt.title(\"Very unfair FD map\")\n",
    "    heat_max = np.max(abs(very_unfair_map_i))\n",
    "    plt.imshow(very_unfair_map_i, cmap=\"bwr\", vmax=heat_max, vmin=-heat_max)\n",
    "    plt.colorbar(shrink=0.5)\n",
    "    \n",
    "    plt.subplot(len(image_ids), 5, 5*i+3)\n",
    "    if i == 0:\n",
    "        plt.title(\"Unfair FD map\")\n",
    "    heat_max = np.max(abs(unfair_heat_maps[i]))\n",
    "    plt.imshow(unfair_heat_maps[i], cmap=\"bwr\", vmax=heat_max, vmin=-heat_max)\n",
    "    plt.colorbar(shrink=0.5)\n",
    "    \n",
    "    plt.subplot(len(image_ids), 5, 5*i+4)\n",
    "    if i == 0:\n",
    "        plt.title(\"FD map\")\n",
    "    heat_max = np.max(abs(heat_maps[i]))\n",
    "    plt.imshow(heat_maps[i], cmap=\"bwr\", vmax=heat_max, vmin=-heat_max)\n",
    "    plt.colorbar(shrink=0.5)\n",
    "    \n",
    "    plt.subplot(len(image_ids), 5, 5*i+5)\n",
    "    if i == 0:\n",
    "        plt.title(\"Location in color space\")\n",
    "    plt.plot(color_probe, color_class) \n",
    "    plt.vlines([color_i], 0, valid_set.num_classes-1, linewidth=0.8,\n",
    "           colors=\"r\", label=\"color value\",\n",
    "           linestyles=\"dashed\")\n",
    "plt.show()  # => somewhat similar results (see image 2), but with an overall ~11x speedup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d075e168",
   "metadata": {},
   "outputs": [],
   "source": [
    "grad_heat_maps = []\n",
    "plt.figure(figsize=(20, 5*len(image_ids)))\n",
    "for i, image_id in tqdm(enumerate(image_ids)):\n",
    "    np.random.seed(image_id)\n",
    "    explain_img_i, target_i, color_i = valid_set.generate_one()\n",
    "    batched_explain_img_i = torch.tensor(np.expand_dims(explain_img_i, 0).transpose(0, 3, 1, 2), requires_grad=True).to(device).float()\n",
    "    output_logit_i = res_net(batched_explain_img_i)[0, target_i.argmax()]\n",
    "    \n",
    "    img_grad_i = torch.autograd.grad(output_logit_i, batched_explain_img_i)[0].squeeze().cpu().numpy()\n",
    "    grad_times_input_i = img_grad_i * np.squeeze(explain_img_i)\n",
    "    \n",
    "    plt.subplot(len(image_ids), 5, 5*i+1)\n",
    "    if i == 0:\n",
    "        plt.title(\"Image\")\n",
    "    plt.imshow(explain_img_i, cmap=\"gray\")\n",
    "    \n",
    "    plt.subplot(len(image_ids), 5, 5*i+2)\n",
    "    if i == 0:\n",
    "        plt.title(\"Input*Gradient explanation\")\n",
    "    heat_max = np.max(abs(grad_times_input_i))\n",
    "    plt.imshow(grad_times_input_i, cmap=\"bwr\", vmax=heat_max, vmin=-heat_max)\n",
    "    plt.colorbar(shrink=0.5)\n",
    "    \n",
    "    plt.subplot(len(image_ids), 5, 5*i+3)\n",
    "    if i == 0:\n",
    "        plt.title(\"Gradient explanation\")\n",
    "    heat_max = np.max(abs(img_grad_i))\n",
    "    plt.imshow(img_grad_i, cmap=\"bwr\", vmax=heat_max, vmin=-heat_max)\n",
    "    plt.colorbar(shrink=0.5)\n",
    "    \n",
    "    plt.subplot(len(image_ids), 5, 5*i+4)\n",
    "    if i == 0:\n",
    "        plt.title(\"FD explanation (unfair)\")\n",
    "    heat_max = np.max(abs(unfair_heat_maps[i]))\n",
    "    plt.imshow(unfair_heat_maps[i], cmap=\"bwr\", vmax=heat_max, vmin=-heat_max)\n",
    "    plt.colorbar(shrink=0.5)\n",
    "    \n",
    "    plt.subplot(len(image_ids), 5, 5*i+5)\n",
    "    if i == 0:\n",
    "        plt.title(\"Location in color space\")\n",
    "    plt.plot(color_probe, color_class) \n",
    "    plt.vlines([color_i], 0, valid_set.num_classes-1, linewidth=0.8,\n",
    "           colors=\"r\", label=\"color value\",\n",
    "           linestyles=\"dashed\")\n",
    "plt.show()\n",
    "# gradient should be zero, so double check computations, fix scale on cmap"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "656cd5d0",
   "metadata": {},
   "source": [
    "# PCA Direction Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "489c03e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DummyNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.linear = nn.Linear(2,2)\n",
    "        self.relu = nn.ReLU()\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.linear(x)\n",
    "        print(\"\\tBefore ReLUs\", x)\n",
    "        return self.relu(x)\n",
    "dummy_net = DummyNet()\n",
    "dummy_net.linear._parameters[\"weight\"].data = torch.nn.Parameter(torch.tensor([[1., 0], [0, 1]]))\n",
    "dummy_net.linear._parameters[\"bias\"].data = torch.nn.Parameter(torch.tensor([200., 200]))\n",
    "print(\"Network parameters\", dummy_net.linear._parameters)\n",
    "print(\"WITHOUT GUIDED BACKPROP\")\n",
    "inpt = torch.tensor([1., -1.], requires_grad=True)\n",
    "result = dummy_net(inpt)\n",
    "print(\"\\tNetwork output\", result)\n",
    "print(\"\\t'Loss'\", -result[0]+result[1])\n",
    "print(\"\\tResulting gradients\", torch.autograd.grad(-result[0]+result[1], inpt))\n",
    "print(\"WITH GUIDED BACKPROP\")\n",
    "guided_dummy = GuidedBackprop(dummy_net)\n",
    "result = guided_dummy(inpt, preserve_hooks=False)\n",
    "print(\"\\tNetwork output\", result)\n",
    "print(\"\\t'Loss'\", -result[0]+result[1])\n",
    "print(\"\\tResulting gradients\", torch.autograd.grad(-result[0]+result[1], inpt))\n",
    "print(\"GUIDED BACKPROP AGAIN (should auto-clean now)\")\n",
    "inpt = torch.tensor([1., -1.], requires_grad=True)\n",
    "result = dummy_net(inpt)\n",
    "print(\"\\tNetwork output\", result)\n",
    "print(\"\\t'Loss'\", -result[0]+result[1])\n",
    "print(\"\\tResulting gradients\", torch.autograd.grad(-result[0]+result[1], inpt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d630c95",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_scales = [3,5,7,9,13,15]\n",
    "if 1: \n",
    "    %store -r color_pca_directions_1_stride color_pca_directions_s_stride\n",
    "else:\n",
    "    color_pca_directions_1_stride = find_pca_directions(valid_set, 16384, default_scales, 1)\n",
    "    color_pca_directions_s_stride = find_pca_directions(valid_set, 16384, default_scales, default_scales)\n",
    "    %store color_pca_directions_1_stride color_pca_directions_s_stride"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ca58f25",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_pca_directions(color_pca_directions_1_stride, \"Strides=1\", default_scales, lines=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c64b11bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_pca_directions(color_pca_directions_s_stride, \"Strides=scales\", default_scales, lines=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56fe52ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(200_010)\n",
    "generated_img, label, *__ = valid_set.generate_one()\n",
    "pca_map_strided = pca_direction_grids(small_net, valid_set, label.argmax(), generated_img, \n",
    "                                      default_scales, color_pca_directions_1_stride, strides=default_scales,\n",
    "                                      device=device, batch_size=128, component=0)\n",
    "pca_map_1_stride = pca_direction_grids(small_net, valid_set, label.argmax(), generated_img, \n",
    "                                      default_scales, color_pca_directions_1_stride, component=0, \n",
    "                                      device=device, batch_size=128, strides=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ac78412",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([generated_img, pca_map_strided])\n",
    "# generated on\n",
    "# small_net = ResNet([[16, 3, 2],  # num_channels (input and output), kernel_size, stride\n",
    "#                     [32, 3, 2],\n",
    "#                     [64, 3, 2]], 3, [128, 128, 1], \n",
    "#                    \"small_net_noise_medium_grey.dict\",\n",
    "# => with strides == scales, unguided"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f115175f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([generated_img, pca_map_1_stride])\n",
    "# generated on\n",
    "# small_net = ResNet([[16, 3, 2],  # num_channels (input and output), kernel_size, stride\n",
    "#                     [32, 3, 2],\n",
    "#                     [64, 3, 2]], 3, [128, 128, 1], \n",
    "#                    \"small_net_noise_medium_grey.dict\",\n",
    "# => with strides == 1, unguided"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff5b467a",
   "metadata": {},
   "outputs": [],
   "source": [
    "seeds = [1_2123, 1_40_124, 1_508_559, 1_5_019_258, 1_2_429_852, 9032, 5832, 12, 5014, 92, 42, 52, \n",
    "         52_934, 935_152, 1_000_000, 1_000_001, 27, 24, 512, 999_105]  # 20 \n",
    "# def generate_many_pca(net, seeds, pca_directions_1_stride, scales, dataset, \n",
    "#         component=0, batch_size=128, strides=None, skip_1_stride=False, device=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90a8eb82",
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"unittest\" for pca_direction_maps\n",
    "fake_pca_directions_s_strides = []\n",
    "for i, scale in enumerate(default_scales):\n",
    "    fake_pca_directions_s_strides.append(color_pca_directions_1_stride[i][::scale, ::scale])\n",
    "np.random.seed(12123)\n",
    "test_img, label, *__ = valid_set.generate_one()\n",
    "new_pca_map = pca_direction_grids(unstrided_net, valid_set, label.argmax(), test_img, \n",
    "                    default_scales, color_pca_directions_1_stride, strides=default_scales,\n",
    "                    device=device, batch_size=128, component=0)\n",
    "old_pca_map = old_old_pca_direction_grids(unstrided_net, valid_set, label.argmax(), test_img, \n",
    "                    pca_direction_grids=fake_pca_directions_s_strides, \n",
    "                    scales=default_scales, device=device)\n",
    "print(abs(new_pca_map - old_pca_map).max())  # should be very low\n",
    "print(abs(new_pca_map - old_pca_map).mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1124c0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "pca_map_s_strides, pca_map_1_strides, grad_maps, explain_imgs = generate_many_pca(unstrided_net, component=0, strided_scales=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64a3df7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "guided_net = GuidedBackprop(unstrided_net)\n",
    "guided_pca_map_s_strides, guided_pca_map_1_strides, guided_grad_maps, explain_imgs = generate_many_pca(guided_net, component=0, strided_scales=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a35a8f10",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([explain_imgs, guided_pca_map_s_strides, guided_pca_map_1_strides, guided_grad_maps, pca_map_s_strides, pca_map_1_strides, grad_maps], transpose=True, titles=[\"Image\", \"Guided Strides=scale\", \"Guided strides=1\", \"Guided Gradient\", \"Strides=scale\", \"strides=1\", \"Gradient\"])\n",
    "# on the strides=2 final layer network  (unstrided_net)\n",
    "# with the new, correct PCA implementation\n",
    "# and sample size of 16384 on PCA directions\n",
    "# also it isnt actually strides=scales, but strides=3\n",
    "\n",
    "# how do we use this to make a better classifier (the binary code behaviour)\n",
    "# dead neuron caused by no gradient (not unusual)\n",
    "# cutting them out doesn't work that well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcfff9ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# what if we disable guided backprop for last layer?\n",
    "np.random.seed(seeds[3])\n",
    "generated_img, label, *__ = valid_set.generate_one()\n",
    "tensored_img = tensorize(generated_img, device=device, requires_grad=True)\n",
    "guided_net = GuidedBackprop(unstrided_net, exceptions=[\"fully_connected.0.act_func\"])\n",
    "pca_map_strided = pca_direction_grids(guided_net, valid_set, label.argmax(), generated_img, \n",
    "                                      default_scales, color_pca_directions_1_stride, strides=3,\n",
    "                                      device=device, batch_size=128, component=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1337e1f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([generated_img, pca_map_strided])\n",
    "# initially seems promising"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98ff3b26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# now do the same for all\n",
    "guided_net = GuidedBackprop(unstrided_net, exceptions=[\"fully_connected.0.act_func\"])\n",
    "guided_pca_map_s_strides, _, guided_grad_maps, explain_imgs = generate_many_pca(guided_net, component=0, strided_scales=3, skip_1_stride=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8321eecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([explain_imgs, guided_pca_map_s_strides, guided_grad_maps], transpose=True, titles=[\"Image\", \"EGuided Strides=3\", \"EGuided Gradient\"])\n",
    "# related to when you start getting non-monotonic\n",
    "# GB = only care about positive paths = only monotonic functions\n",
    "# its only the increasing ones\n",
    "# makes sense to put negative weights at the end? \n",
    "# sample how many weights in each layer are +ve/-ve?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d747a2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from captum.attr import IntegratedGradients, DeepLift"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "899bd47a",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(seeds[11])\n",
    "generated_img, label, *__ = valid_set.generate_one()\n",
    "tensored_img = tensorize(generated_img, device=device, requires_grad=True)\n",
    "baseline = torch.zeros_like(tensored_img)\n",
    "print(label.argmax())\n",
    "ig_model = IntegratedGradients(unstrided_net)\n",
    "ig_attrs = ig_model.attribute(tensored_img, baseline, target=2)\n",
    "dl_model = DeepLift(unstrided_net)\n",
    "dl_attrs = dl_model.attribute(tensored_img, baseline, target=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d4c6731",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dl_attrs.shape)\n",
    "ind = torch.topk(dl_attrs.flatten(), 60).indices.cpu()\n",
    "dl_attrs[np.unravel_index(ind, dl_attrs.shape)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad0d3794",
   "metadata": {},
   "outputs": [],
   "source": [
    "imshow_centered_colorbar(ig_attrs.detach().cpu().numpy().squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6563f7a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "imshow_centered_colorbar(dl_attrs.detach().cpu().numpy().squeeze())\n",
    "\n",
    "# measurement of interesection over union of bounding box (top k pixels?)\n",
    "# find a dataset that we can test this on?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69c70afc",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([explain_imgs, guided_pca_map_s_strides, guided_pca_map_1_strides, guided_grad_maps, pca_map_s_strides, pca_map_1_strides, grad_maps], transpose=True, titles=[\"Image\", \"Guided Strides=scale\", \"Guided strides=1\", \"Guided Gradient\", \"Strides=scale\", \"strides=1\", \"Gradient\"])\n",
    "# on the strides=2 final layer network (unstrided_net)\n",
    "# with the old, slightly incorrect PCA implementation\n",
    "# and sample size of only 2048 on PCA directions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ea42375",
   "metadata": {},
   "outputs": [],
   "source": [
    "def recall_k(ground_truth, k, explanation):\n",
    "    # threshold is implicitly chosen such that exactly k values are above the threshold\n",
    "    # choosing k, we could set k such that it is equal to the size of the ground truth target\n",
    "    # however, this approach is arguably biased since there could be other areas of the image \n",
    "    # that are equally salient, and thus the intersection would be split between those\n",
    "    # salient regions. setting k to the max target size biases against smaller targets\n",
    "    # (they will be limited to a max recall). Setting k to the lowest target size means\n",
    "    # larger targets will be biased against since they will again be limited by a max.\n",
    "    # The only other way is with a threshold\n",
    "    # but scales are arbitrary so this may be tricky. Maybe like a scale that switches \n",
    "    # based on the statistics of the explanation? Don't want to do percentiles since that's\n",
    "    # just top k again. \n",
    "    # Maybe do some sort of \"the explanation is a measure, which measure assigns highest \n",
    "    # probablity to the ground truth indicator\" => would have to normalize the explanations to \n",
    "    # sum to 1?\n",
    "    top_k_locs = torch.topk(explanation.flatten(), k).indices.cpu()\n",
    "    indices = np.unravel_index(top_k_locs, explanation.shape)\n",
    "    intersection = ground_truth[indices].sum()\n",
    "\n",
    "    ground_truth[indices] = 1\n",
    "    return intersection/ground_truth.sum()\n",
    "\n",
    "truth_example = torch.zeros((128,128))\n",
    "truth_example[:32, :32] = 1\n",
    "explain_example = torch.rand(128,128)\n",
    "\n",
    "plt.subplot(1,3,1)\n",
    "imshow_centered_colorbar(truth_example.numpy())\n",
    "\n",
    "print(recall_k(truth_example, int((128**2)/16), explain_example))\n",
    "\n",
    "plt.subplot(1,3,3)\n",
    "imshow_centered_colorbar(explain_example.numpy())\n",
    "\n",
    "recall_k(truth_example, int((128**2)/16), dl_attrs.squeeze())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c51d7194",
   "metadata": {},
   "source": [
    "# Investigation into Guided Backprop Failure Cases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f5fdf63",
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_classes(min_y, max_y, ax=None, alpha=0.25):\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    color_probe = np.arange(255)\n",
    "    classified = np.vectorize(valid_set.color_classifier)(color_probe)\n",
    "    classified = classified/(valid_set.num_classes-1)*(max_y-min_y) + min_y\n",
    "    ax.plot(classified, c=\"k\", alpha=alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75fcc649",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def top_k_activating(net, loader, activation_getter, k=50, device=None):\n",
    "    # assume already has AllActivations hooks\n",
    "    net.eval()\n",
    "    best_seeds = np.zeros(k)\n",
    "    best_colors = np.zeros(k)\n",
    "    seed_activations = np.full(k, -np.inf)\n",
    "    for i, sample in tqdm(enumerate(loader)):\n",
    "        imgs = sample[\"image\"].to(device).float()\n",
    "        seeds = sample[\"seeds\"]\n",
    "        colors = sample[\"color\"]\n",
    "        net(imgs)\n",
    "        activs = activation_getter(net._features)\n",
    "        for activ,seed,color in zip(activs, seeds, colors):\n",
    "            if activ > min(seed_activations):\n",
    "                lowest_entry = np.argmin(seed_activations)\n",
    "                seed_activations[lowest_entry] = activ\n",
    "                best_seeds[lowest_entry] = seed\n",
    "                best_colors[lowest_entry] = color\n",
    "    sorted_indices = seed_activations.argsort()[::-1]\n",
    "    data, *_ = plt.hist(best_colors, bins=255)\n",
    "    draw_classes(0, max(data))\n",
    "    return best_seeds[sorted_indices], seed_activations[sorted_indices]\n",
    "    \n",
    "interp_net = AllActivations(unstrided_net)\n",
    "def logit_1(features):\n",
    "    return features[\"fully_connected.0.act_func\"].sum(axis=1).detach().cpu().numpy()\n",
    "best_seeds,_ = top_k_activating(interp_net, valid_loader, logit_1, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f48c6196",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_seeds,_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc23b495",
   "metadata": {},
   "outputs": [],
   "source": [
    "for seed in result[0]:\n",
    "    np.random.seed(int(seed))\n",
    "    explain_img, label, color, *_ = valid_set.generate_one()\n",
    "    print(color)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b06ef1ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "uniq = set()\n",
    "for x in color_distrib.values():     # random_net\n",
    "    uniq = uniq.union(set(x.keys()))    # => did NOT learn the same 3 logit structure\n",
    "print(len(uniq))\n",
    "uniq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01dfe383",
   "metadata": {},
   "outputs": [],
   "source": [
    "uniq = set()\n",
    "for x in color_distrib.values():     # random_net2  learned more compact structure??\n",
    "    uniq = uniq.union(set(x.keys()))\n",
    "print(len(uniq))\n",
    "uniq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f755a08",
   "metadata": {},
   "outputs": [],
   "source": [
    "uniq = set()\n",
    "for x in color_distrib.values():     # random_net3 \n",
    "    uniq = uniq.union(set(x.keys()))    # similarly fails to replicate unstrided_net\n",
    "print(len(uniq))\n",
    "uniq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc3c66ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "uniq = set()\n",
    "for x in color_distrib.values():\n",
    "    uniq = uniq.union(set(x.keys()))\n",
    "pattern_to_names = {}\n",
    "for pattern in uniq:\n",
    "    pattern_to_names[pattern] = pattern[1+2*1] + pattern[1+2*4] + pattern[1+2*17]\n",
    "# these were the only non-zero logits across ALL cases!\n",
    "# seems to be some sort of binary code\n",
    "pattern_to_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9364bf1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pattern_totals = defaultdict(int)\n",
    "for color, values in color_distrib.items():\n",
    "    for pattern, count in values.items():\n",
    "        pattern_totals[pattern] += count\n",
    "# counts of the different \"binary codes\"\n",
    "# top 4 patterns are the 3 \"pure codes\"\n",
    "# 1. pure\n",
    "# 2. all 3\n",
    "# 3. pure\n",
    "# 4. pure\n",
    "pattern_totals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53ff0345",
   "metadata": {},
   "outputs": [],
   "source": [
    "for pattern in uniq:\n",
    "    amounts = np.zeros((255,))\n",
    "    for color, distrib in color_distrib.items():\n",
    "        amounts[color] = distrib[pattern]\n",
    "    plt.plot(amounts, label=pattern_to_names[pattern])\n",
    "color_probe = np.linspace(0, 255, 255)\n",
    "color_class = [color_classifier(x)*400 for x in color_probe]\n",
    "plt.plot(color_probe, color_class, label=\"classes\", linestyle=\"dotted\")\n",
    "plt.legend()\n",
    "# transition regimes\n",
    "# pure regimes\n",
    "# underlying \"noise\"\n",
    "# next plot to do is just the raw activations of the 3 classes\n",
    "# predicts which images will be good and which will not\n",
    "\n",
    "# [0, 100] => Varying hints of a circle, square-ish/sparse behaviour (pure 100, class 0)\n",
    "# [100, 150] => Square-ish/sparse behaviour (all 3, class 1)\n",
    "# [150, 200] => all zeros  (pure 001, class 2)\n",
    "# [200, 250] => excellent quality  (pure 010, class 1)\n",
    "\n",
    "# do activation thing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cf55d31",
   "metadata": {},
   "source": [
    "# Color Profiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2f721b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile_plots,_ = activation_color_profile(AllActivations(unstrided_net), valid_loader, valid_set, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d19f81d",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_profile_plots(profile_plots, [f\"fully_connected.0.act_func_{x}\" for x in [1,4,17]],\n",
    "                   fixed_height=True)\n",
    "# color profiles of hidden units 1, 4, 17, for unstrided_net\n",
    "# see that in the '111' regime, its \"equally units 17 and 4, small unit 1\"\n",
    "# class 2 is very strongly \"no 1 or 4, only 17\"\n",
    "# class 0 is almost entirely unit 1 (some small amount of unit 17, but basically no unit 4)\n",
    "# class 1 has a different expression in the >200 color regime (as compared to its expression\n",
    "#    in the '111', aka 100 < color < 150 regime), in being \n",
    "#    \"mostly unit 4, some unit 17, little unit 1\" (4 increase, 17 decrease)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0443bb6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_conv_weights(unstrided_net, \"conv_blocks.0.act_func1\", color_profile=profile_plots, size_mul=(6,12), fixed_height=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "896c33ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_conv_weights(unstrided_net, \"conv_blocks.0.act_func2\", color_profile=profile_plots, size_mul=(3,6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d98de57",
   "metadata": {},
   "outputs": [],
   "source": [
    "uniform_inpt = torch.full((1,1,32,32), 200.0).to(device)\n",
    "uniform_out = unstrided_net.conv_blocks[0](uniform_inpt)\n",
    "plt.imshow(uniform_out[0,12].detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f43f0567",
   "metadata": {},
   "outputs": [],
   "source": [
    "unstrided_net.eval()\n",
    "c = 12\n",
    "\n",
    "#uniform_inpt = torch.full((1,16,32,32), 100.0).to(device)\n",
    "#plt.imshow(unstrided_net.conv_blocks[0].conv2.weight[c, in_c].detach().cpu().numpy(), cmap=\"bwr\")\n",
    "conv_maps = unstrided_net.conv_blocks[0].conv2.weight[c, :]\n",
    "imshow_centered_colorbar(conv_maps[7].detach().cpu().numpy(), cmap=\"bwr\")\n",
    "conv_scale = conv_maps.sum(axis=-1).sum(axis=-1)\n",
    "conv_shift = unstrided_net.conv_blocks[0].conv2.bias[c]\n",
    "bn_scale = unstrided_net.conv_blocks[0].batch_norm2.weight[c]\n",
    "bn_shift = unstrided_net.conv_blocks[0].batch_norm2.bias[c]\n",
    "bn_var = unstrided_net.conv_blocks[0].batch_norm2.running_var[c]\n",
    "bn_mean = unstrided_net.conv_blocks[0].batch_norm2.running_mean[c]\n",
    "print(conv_shift, bn_scale, bn_shift, bn_var, bn_mean)\n",
    "#(c*conv_scale + conv_shift - bn_mean) / torch.sqrt(bn_var) * bn_scale + bn_shift\n",
    "slope = (conv_scale/torch.sqrt(bn_var)*bn_scale).detach().cpu().numpy()\n",
    "bias = ((conv_shift - bn_mean)/torch.sqrt(bn_var)*bn_scale + bn_shift).detach().cpu().numpy()\n",
    "\n",
    "lines = np.asarray([profile_plots[f\"conv_blocks.0.act_func1_{x}\"][0] for x in range(16)])\n",
    "\n",
    "uniform_scaling = slope.dot(lines) + bias\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba0e7941",
   "metadata": {},
   "outputs": [],
   "source": [
    "uniform_scaling[uniform_scaling < 0] = 0\n",
    "plt.plot(uniform_scaling)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8e677a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "last_layer_weight = unstrided_net._modules[\"fully_connected\"][-1].fully_connected.weight.detach().cpu().numpy()\n",
    "imshow_centered_colorbar(last_layer_weight, cmap=\"bwr\", title=\"Last Layer FC weights\", colorbar=False)\n",
    "plt.vlines([1, 4, 17], ymin=0, ymax=2.5)\n",
    "\n",
    "last_layer_bias = unstrided_net._modules[\"fully_connected\"][-1].fully_connected.bias.detach().cpu().numpy()\n",
    "print(last_layer_bias)\n",
    "# in column 17, all weights are negative, so guided backprop means we immediately zero everything out\n",
    "# relevant columns have been highlighted\n",
    "\n",
    "# column 1 mostly means class 0, very strongly not class 2\n",
    "# column 4 mostly means class 1, equally strongly not class 0 and 2\n",
    "# column 17 means not class 0, not class 1, barely class 2 => problematic\n",
    "\n",
    "# since class 2 is \"default class\" (largest bias), the negative weights in column 17 are fine\n",
    "# rightmost is easiest "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0566f336",
   "metadata": {},
   "outputs": [],
   "source": [
    "seeds = [1_2123, 1_40_124, 1_508_559, 1_5_019_258, 1_2_429_852, 9032, 5832, 12, 5014, 92, 42, 52, \n",
    "         52_934, 935_152, 1_000_000, 1_000_001, 27, 24, 512, 999_105]  # 20 \n",
    "for i, seed in enumerate(seeds):\n",
    "    np.random.seed(seed)\n",
    "    print(i, valid_set.generate_one()[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1864c1a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6*4, 6*8))\n",
    "first_fc = unstrided_net._modules[\"fully_connected\"][0].fully_connected.weight.detach().cpu().numpy()\n",
    "# only bother visualizing outputs 1, 4, and 17 (add others just to see)\n",
    "relevant_outputs = range(32)\n",
    "for i, output_col in enumerate(relevant_outputs):\n",
    "    fc_weights = np.concatenate(np.concatenate(first_fc[output_col].reshape(8, 8, 63, 63),1),1)\n",
    "    plt.subplot(8,4,i+1)\n",
    "    imshow_centered_colorbar(fc_weights, cmap=\"bwr\", title=f\"FC weights of {output_col}\")\n",
    "# weights of final conv -> first fully connected\n",
    "\n",
    "# 4-5 shared \"empty\" maps across useful logits => useless channels??\n",
    "\n",
    "# antipodal-superposition-like structure in weights of useful logits\n",
    "# eg. logit 4 and logit 17, map (0,1)\n",
    "# eg. logit 1 and logit 4, map (1,2), map (-2, -1)\n",
    "\n",
    "# \"useful logits\" are of much higher overall norm (~2) and more positive than other logits\n",
    "# other maps are usually in range of ~0.02, exception is logit 16, which is ~0.2\n",
    "# could indicate a logit that was used in early stages of training but was eventually dropped?\n",
    "\n",
    "\n",
    "# epicycles \n",
    "\n",
    "# how to decide when to use other PCA components?\n",
    "# large checkrboard seperated by diagonal on top its light, bottom its dark\n",
    "# componont you want is \n",
    "\n",
    "# narrative:\n",
    "# investigating classifying by color\n",
    "# investigating attribution methods in cases when global context is important\n",
    "# generalization of eigenfaces to more context windows\n",
    "# guided backprop = edges (conv nets = edges too)\n",
    "# force dataset to not be edges, what does conv net do?\n",
    "# it can do, and heres how it does it\n",
    "# logits not monotonic on color\n",
    "#    if hidden units also non-monotonic, then what is the purpose of last layer?\n",
    "#     initial conv layers are monotonic (almost linear transformation)\n",
    "\n",
    "# classifying wiht conv nets when global context matters\n",
    "# object detection = hierarchical but still template matching\n",
    "# this task is not template matching at all\n",
    "# patches of color cant be capture with edge/template at all\n",
    "# write up mechanistic_interp.ipynb, combine with this stuff\n",
    "# in addition, if we can find good way to use PCA, add to this paper or for a future one\n",
    "\n",
    "# large network random weights some get lucky to be useful\n",
    "\n",
    "# 2 pagse, unlimited appendix\n",
    "# template for ICLR\n",
    "\n",
    "# claim: learning non-monotonic functions is difficult, superposition\n",
    "# reason most of the neurons are initialized so poorly that they can't become useful thru\n",
    "# gradient descent (or at least the other ones [1, 4, 17]) were close enough that they \n",
    "# came to dominate over them\n",
    "# how to check, keep snapshot of initial random network, see what happens if you zero-out\n",
    "# the weights that ended up being good"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "739ab44b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6*4, 6*9))\n",
    "first_fc = unstrided_net._modules[\"fully_connected\"][0].fully_connected.weight.detach().cpu().numpy()\n",
    "# only bother visualizing outputs 1, 4, and 17 (add others just to see)\n",
    "relevant_outputs = [1, 4, 17]\n",
    "for i, output_col in enumerate(relevant_outputs):\n",
    "    fc_weights = np.concatenate(np.concatenate(first_fc[output_col].reshape(8, 8, 63, 63),1),1)\n",
    "    plt.subplot(1,3,i+1)\n",
    "    imshow_centered_colorbar(fc_weights, cmap=\"bwr\", title=f\"FC weights of {output_col}\")\n",
    "# same figure but only the relevant ones\n",
    "# maybe the zero-ed out channels are dead neurons (old idea)\n",
    "# without batchnorm/dropout style things, dead will stay dead (so some get unlucky)\n",
    "# batchnorm is supposed to resample => sometimes allows positive grads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a3b6732",
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate network activations, display tanh(activations) (for better scaling)\n",
    "def visualize_network_activations(net, seed):\n",
    "    net.eval()\n",
    "    np.random.seed(seed)\n",
    "    working_img, label, color, *_____ = valid_set.generate_one()\n",
    "    print(color, label.argmax())\n",
    "    tensored_img = tensorize(working_img, device=device)\n",
    "    debug_net = AllActivations(net)\n",
    "    debug_net(tensored_img)\n",
    "\n",
    "    plt.figure(figsize=(6*4, 12))\n",
    "    final_relu = debug_net._features[\"conv_blocks.2.act_func2\"].detach().cpu().numpy()[0].reshape(8,8,63,63)\n",
    "    print(\"Logits 1,4,17\", debug_net._features[\"fully_connected.0.act_func\"][0, (1,4,17)])\n",
    "    compressed_results = np.tanh(np.concatenate(np.concatenate(final_relu, 1), 1))\n",
    "\n",
    "    imshow_centered_colorbar(compressed_results, cmap=\"bwr\", title=f\"Post ReLU final conv layer activations (color={color})\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6ac40ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_network_activations(unstrided_net, seeds[9])\n",
    "# 111 regime, brightest map is (2, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29bb9812",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_network_activations(unstrided_net, seeds[19])\n",
    "# good regime, 010\n",
    "# bright areas are (4, 2) and (6, 4)  => matches FC weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c91e237",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_network_activations(unstrided_net, seeds[0])\n",
    "# in 100 regime, bright maps are (2, 1), (3, 2), (1, 3), and (-2, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90f21e6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_network_activations(unstrided_net, seeds[11])\n",
    "# in 001 regime, bright maps are  (-2, 4)\n",
    "# both logit 1 and logit 17 have big negatives there"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9add34b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6*4, 12))\n",
    "final_relu = debug_net._features[\"conv_blocks.2.act_func2\"].detach().cpu().numpy()[0].reshape(8,8,63,63)\n",
    "compressed_results = np.tanh(np.concatenate(np.concatenate(final_relu, 1), 1))\n",
    "\n",
    "imshow_centered_colorbar(compressed_results, cmap=\"bwr\", title=\"Post ReLU final conv layer activations\")\n",
    "# with tanh applied to activations (note that ReLU is still the actual activation function)\n",
    "# only certain maps seem important, this is "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "746fed49",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6*4, 12))\n",
    "final_conv_map = debug_net._features[\"conv_blocks.2.batch_norm2\"].detach().cpu().numpy()[0].reshape(8,8,63,63)\n",
    "conv_max = abs(final_conv_map).max()\n",
    "\n",
    "compressed_results = np.concatenate(np.concatenate(final_conv_map, 1), 1)\n",
    "imshow_cenered_colorbar(compressed_results, cmap=\"bwr\", title=\"Pre ReLU final conv layer activations\")\n",
    "# useless"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcfb7e0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def adversarial_generate(img, lbl, net, alpha, lr, runs):\n",
    "    # alpha is maximum norm that the adversarial can be\n",
    "    adversarial_direction = np.random.uniform(-alpha, alpha, size=(1, valid_set.size, valid_set.size))\n",
    "    adversarial_direction = adversarial_direction/np.linalg.norm(adversarial_direction)*alpha\n",
    "    \n",
    "    # pick arbitrary target\n",
    "    bad_lbl = (lbl.argmax() + 1) % valid_set.num_classes\n",
    "    target = torch.tensor(np.zeros_like(lbl)).to(device).unsqueeze(0).float()\n",
    "    target[0,2] = 1.\n",
    "    print(target, lbl)\n",
    "    \n",
    "    tensor_img = torch.tensor(img.transpose(2,0,1)).unsqueeze(0).to(device).float()\n",
    "    tensor_adv_dir = torch.tensor(adversarial_direction, requires_grad=True).unsqueeze(0).to(device).float()\n",
    "    \n",
    "    loss_func = nn.CrossEntropyLoss()\n",
    "    for i in range(runs):\n",
    "        curr_img = tensor_img + tensor_adv_dir\n",
    "        curr_net_out = net(curr_img)\n",
    "        curr_loss = loss_func(curr_net_out, target)# + 5e-10*torch.linalg.norm(tensor_adv_dir)\n",
    "        grad_dir = torch.autograd.grad(curr_loss, tensor_adv_dir)[0]\n",
    "        tensor_adv_dir -= lr*grad_dir\n",
    "        tensor_adv_dir = torch.clamp(tensor_adv_dir, min=-alpha, max=alpha)\n",
    "\n",
    "        if i % (runs//5) == (runs//5-1):\n",
    "            print(curr_loss.item(), curr_net_out)\n",
    "    return tensor_img + tensor_adv_dir\n",
    "np.random.seed(58)\n",
    "generated_img, gen_label, color, *_ = valid_set.generate_one()\n",
    "print(color)\n",
    "adv_example = adversarial_generate(generated_img, gen_label, unstrided_net, 6.4, 1e2, 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0a138c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# highly variable results in terms of alpha (max pixel diff)\n",
    "# seems to work \"best\" (lowest alpha needed) when going to class 2??? (shouldn't it be class 1\n",
    "# since that one in the most polysemantic (at least in one of its regimes))\n",
    "\n",
    "# 230 color -> alpha very close to 5, works on 1e2, 500 (seed 55)\n",
    "# 136 color -> alpha of 3 (seed 54)\n",
    "# 50 color -> alpha of 9-10 (seed 53)\n",
    "# 181 color (in class 2 already) -> adv_dir is basically random noise, weird cyclic structure to it (seed 52)\n",
    "# 82 color -> alpha of 3.8 (high) (seed 51)\n",
    "# 201 color -> alpha of 0-1 (high) (seed 50)\n",
    "# 110 color -> alpha of 2-3 (low) (seed 56)\n",
    "# 232 color -> alpha close to 5 (seed 57)\n",
    "# 60 color -> alpha 6-7 (mid) (seed 58)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71399de7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplot(1,3,1)\n",
    "plt.imshow(generated_img, cmap=\"gray\")\n",
    "plt.subplot(1,3,2)\n",
    "np_adv = adv_example.detach().cpu().numpy().squeeze()\n",
    "plt.imshow(np_adv, cmap=\"gray\")\n",
    "plt.subplot(1,3,3)\n",
    "imshow_centered_colorbar(generated_img.squeeze()-np_adv, cmap=\"bwr\", title=\"Adv direction\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c09307ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([explain_imgs, pca_map_s_strides, pca_map_1_strides], transpose=True, titles=[\"Image\", \"Strides=scale\", \"strides=1\"])\n",
    "# add comparison to regular gradient\n",
    "# smaller circles = bad?\n",
    "# manually test if the gradient changes make sense\n",
    "# model lerans weird stuff about the noise\n",
    "# interpertation of its algo is interesting\n",
    "\n",
    "# make texture dataset and test the methods on it\n",
    "# texture generation: emerging conv?\n",
    "# heuristic = lots of code\n",
    "# dataset\n",
    "# test on natual images eventually\n",
    "\n",
    "# show it works when edeges important too (guided backprop first)\n",
    "# could do saliency checks\n",
    "\n",
    "# fourier transform could work for texture, if its the whole image (to get window size)\n",
    "# do fft on quadrants of image to guess at scale, look at max fourier coeff =>\n",
    "# should give rough idea of window size (top k coeffs?)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "296b8ec7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([explain_imgs, pca_map_s_strides, pca_map_1_strides, grad_maps], transpose=True, titles=[\"Image\", \"Strides=scale\", \"strides=1\", \"Gradient\"])\n",
    "# on the strides=8 final layer network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74688f12",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt_grid_figure([explain_imgs, pca_map_s_strides, pca_map_1_strides, grad_maps], transpose=True, titles=[\"Image\", \"Strides=scale\", \"strides=1\", \"Gradient\"])\n",
    "# on the strides=2 final layer network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85fd60bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(valid_set.generate_one()[0], cmap=\"gray\")\n",
    "plt.xticks([])\n",
    "plt.yticks([])\n",
    "remove_borders(plt.gca())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95c36226",
   "metadata": {},
   "outputs": [],
   "source": [
    "# surely this is the best way to do this :)\n",
    "x = np.arange(400).reshape(5,5,4,4)\n",
    "from itertools import permutations\n",
    "for transp_1 in permutations([0,1,2,3]):\n",
    "    for axis1 in range(3):\n",
    "        for transp_2 in permutations([0,1,2]):\n",
    "            for axis2 in range(2):\n",
    "                for transp_3 in permutations([0,1]):\n",
    "                    try:\n",
    "                        t1 = x.transpose(*transp_1)\n",
    "                        t2 = np.concatenate(t1, axis1)\n",
    "                        t3 = t2.transpose(*transp_2)\n",
    "                        t4 = np.concatenate(t3, axis2)\n",
    "                        t5 = t4.transpose(*transp_3)\n",
    "                        if t5.shape == (20,20):\n",
    "                            if all(t5[0,:4] == np.arange(4)) and t5[0,4] == 16:\n",
    "                                pass\n",
    "                                #print(transp_1, axis1, transp_2, axis2, transp_3)\n",
    "                    except:\n",
    "                        continue\n",
    "np.concatenate(np.concatenate(x,1),1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57e91693",
   "metadata": {},
   "source": [
    "So the question then becomes, how do we search for useful reference images/pixel values in general? We want the distance to be close to the image (small denominator), but also lead to large differences in output logits. This is dangerously close to finding adversarial directions, so we need to make sure we stay in the data manifold => need to establish some sort of distance metric, and potentially a way of detecting whether we are in manifold or not, so we can project into manifold space if needed. This also allows us to switch to a DeepLIFT style type thing."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f86a46e",
   "metadata": {},
   "source": [
    "# Model Optimization Stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91f3a40e",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_net.save_model_state_dict(optim=optim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94a9a91a",
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_img, _, __ = valid_set.generate_one()\n",
    "generated_img = torch.tensor(generated_img.transpose(2,0,1)).to(device).unsqueeze(0).float()\n",
    "for _ in range(1000):\n",
    "    small_net.forward(generated_img, profile=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd852210",
   "metadata": {},
   "outputs": [],
   "source": [
    "total = sum(stats.values())  # --> gave 3x speed! (Fast and Accurate Model scaling?)\n",
    "for k,v in stats.items():    # --> the 3x speedup caused underfitting though, so switched to 2x\n",
    "    print(k,(100.*v/total))\n"
   ]
  }
 ],
 "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}