2024-01-04-Stitching-Pretrained-Imagenet-b.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision.datasets as thd\n",
    "import timm\n",
    "import recomb.layers as ly\n",
    "import igraph as ig\n",
    "import matplotlib.pyplot as plt\n",
    "import polars as pl\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "import recomb.cx as cx\n",
    "import recomb.layers as ly\n",
    "import recomb.problems as problems\n",
    "from recomb.cx import forward_get_all_feature_maps, construct_trained_cx_network_stitching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: imagenet networks are often image-size invariant.\n",
    "# We need to account for this, somehow, when stitching.\n",
    "# (Because this means the image shape does not need to be constrained)\n",
    "imgnet_in_shape = (1, 3, 256, 256)\n",
    "# imgnet_train = thd.ImageNet(\"<add-dataset-folder>\")\n",
    "dataset_path = Path(\"<add-dataset-folder>\")\n",
    "# Some allowable batch size (just in case)\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# resnet50 and resnext50_32x4d ?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timm.list_models(\"resnet*\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_a = timm.create_model(\"resnet50\", pretrained=True)\n",
    "model_a.eval()\n",
    "gca = ly.trace_network(model_a, imgnet_in_shape).to_neural_net_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import SVG\n",
    "with open(\"g.dot\", \"w\") as f:\n",
    "    gca.to_dot(f, include_ord_label=True)\n",
    "! dot g.dot -Tsvg -og.png\n",
    "SVG(\"g.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_b = timm.create_model(\"vgg13\", pretrained=True)\n",
    "model_b = timm.create_model(\"resnext50_32x4d\", pretrained=True)\n",
    "model_b.eval()\n",
    "gcb = ly.trace_network(model_b, imgnet_in_shape).to_neural_net_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import SVG\n",
    "with open(\"g.dot\", \"w\") as f:\n",
    "    gcb.to_dot(f, include_ord_label=True)\n",
    "! dot g.dot -Tsvg -og.png\n",
    "SVG(\"g.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.transforms import v2 as transforms\n",
    "from timm.data import resolve_data_config\n",
    "from timm.data.transforms_factory import create_transform\n",
    "\n",
    "def get_transform_for_timm_model(model):\n",
    "    transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))\n",
    "    return transform\n",
    "\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])\n",
    "tf = transforms.Compose([\n",
    "    transforms.ToImage(),\n",
    "    transforms.Resize(235, antialias=True),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToDtype(torch.float32, scale=True),\n",
    "    normalize,\n",
    "])\n",
    "# tf = get_transform_for_timm_model(model_a)\n",
    "\n",
    "imagenet_train = problems.ImageNetKG(\"<add-dataset-folder>\", \"train\", transform=tf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "dl = DataLoader(imagenet_train)\n",
    "dli = iter(dl)\n",
    "for _ in range(5):\n",
    "    X_o, y_o = next(dli)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(X_o[0, :, :, :].permute(1, 2, 0))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outpt = torch.topk(model_a(X_o), k=5)\n",
    "list(imagenet_train.class_name(outpt.indices.ravel().numpy())), list(imagenet_train.class_name(y_o.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outpt = torch.topk(gca(X_o), k=5)\n",
    "list(imagenet_train.class_name(outpt.indices.ravel().numpy())), list(imagenet_train.class_name(y_o.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outpt = torch.topk(model_b(X_o), k=5)\n",
    "list(imagenet_train.class_name(outpt.indices.ravel().numpy())), list(imagenet_train.class_name(y_o.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outpt = torch.topk(gcb(X_o), k=5)\n",
    "list(imagenet_train.class_name(outpt.indices.ravel().numpy())), list(imagenet_train.class_name(y_o.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recomb.problems import NeuralNetIndividual, ImageNetProblem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "neti_a = NeuralNetIndividual(gca)\n",
    "neti_b = NeuralNetIndividual(gcb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set validation sample limit to 1000.\n",
    "# There are 23000-ish samples in the full validation set.\n",
    "# As evaluating 1000 samples takes ~15s per # network, \n",
    "# 23000-ish samples should take approximately 345s each.\n",
    "problem = ImageNetProblem(\"<add-dataset-folder>\", validation_sample_limit=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "problem.evaluate_network(dev, neti_a, batch_size=16, objective=\"both\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "problem.evaluate_network(dev, neti_b, batch_size=16, objective=\"both\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Does ensembling provide identical results too?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_be = ly.LinearEnsemble([model_a, model_b], [0.5, 0.5])\n",
    "neti_be = NeuralNetIndividual(net_be)\n",
    "torch.manual_seed(42)\n",
    "problem.evaluate_network(dev, neti_be, batch_size=16, objective=\"both\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_be = ly.LinearEnsemble([gca, gcb], [0.5, 0.5]).to_graph()\n",
    "neti_be = NeuralNetIndividual(net_be)\n",
    "torch.manual_seed(42)\n",
    "problem.evaluate_network(dev, neti_be, batch_size=16, objective=\"both\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_be = ly.LinearEnsemble([model_a, model_b], [1.0, 1.0])\n",
    "neti_be = NeuralNetIndividual(net_be)\n",
    "torch.manual_seed(42)\n",
    "problem.evaluate_network(dev, neti_be, batch_size=16, objective=\"both\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_be = ly.LinearEnsemble([gca, gcb], [1.0, 1.0]).to_graph()\n",
    "neti_be = NeuralNetIndividual(net_be)\n",
    "torch.manual_seed(42)\n",
    "problem.evaluate_network(dev, neti_be, batch_size=16, objective=\"both\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import importlib \n",
    "import recomb.cx as cx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We use the training dataset to avoid training on the validation data\n",
    "dataset = problem.get_dataset_train()\n",
    "\n",
    "# Grab an item from a dataloader for use in the forward pass\n",
    "from torch.utils.data import DataLoader\n",
    "dl = DataLoader(dataset)\n",
    "dli = iter(dl)\n",
    "X, _y = next(dli)\n",
    "\n",
    "for p in gca.parameters():\n",
    "    p.requires_grad_(False)\n",
    "for p in gcb.parameters():\n",
    "    p.requires_grad_(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "importlib.reload(cx)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_S(net_a, net_b, stitching_library, compute_similarity):\n",
    "    net_a.store_state_eval()\n",
    "    net_b.store_state_eval()\n",
    "    use_gpu = True \n",
    "    if use_gpu:\n",
    "        net_b.to(dev)\n",
    "        net_a.to(dev)\n",
    "        X_in_many = X.to(dev)\n",
    "    # Set up reference networks.\n",
    "    fms_a, points_a = forward_get_all_feature_maps(net_a, X_in_many, return_points=True)\n",
    "    fms_b, points_b = forward_get_all_feature_maps(net_b, X_in_many, return_points=True)\n",
    "\n",
    "    net_a.train_restore()\n",
    "    net_b.train_restore()\n",
    "\n",
    "    points_a_v = [p for (p, fm) in zip(points_a, fms_a) if fm is not None]\n",
    "    fms_a_v = [fm for fm in fms_a if fm is not None]\n",
    "    points_b_v = [p for (p, fm) in zip(points_b, fms_b) if fm is not None]\n",
    "    fms_b_v = [fm for fm in fms_b if fm is not None]\n",
    "\n",
    "    # Characterize graphs\n",
    "    stitching_library.characterize_graph(net_a.graph)\n",
    "    stitching_library.characterize_graph(net_b.graph)\n",
    "\n",
    "    # Store feature map shapes\n",
    "    for fmidx, (p, fm) in enumerate(zip(points_a, fms_a)):\n",
    "        if fm is None:\n",
    "            continue\n",
    "        # net_a.graph.vs[p[1]][\"sh\"] = list(fm.shape)\n",
    "        stitching_library.characterize_fm(net_a.graph.vs[p[1]], fm, net_a.graph)\n",
    "        net_a.graph.vs[p[1]][\"fmidx\"] = fmidx\n",
    "    for fmidx, (p, fm) in enumerate(zip(points_b, fms_b)):\n",
    "        if fm is None:\n",
    "            continue\n",
    "        # net_b.graph.vs[p[1]][\"sh\"] = list(fm.shape)\n",
    "        stitching_library.characterize_fm(net_b.graph.vs[p[1]], fm, net_b.graph)\n",
    "        net_b.graph.vs[p[1]][\"fmidx\"] = fmidx\n",
    "\n",
    "\n",
    "    S = cx.compute_pairwise_similarities(\n",
    "        net_a.graph,\n",
    "        fms_a_v,\n",
    "        points_a_v,\n",
    "        net_b.graph,\n",
    "        fms_b_v,\n",
    "        points_b_v,\n",
    "        compute_similarity=compute_similarity,\n",
    "        stitching_library=stitching_library,\n",
    "    )\n",
    "    # Input & output override\n",
    "    S[0, 0] = 1\n",
    "    S[1, 1] = 1\n",
    "\n",
    "    return S\n",
    "\n",
    "S = compute_S(\n",
    "    gca, gcb,\n",
    "    cx.CVStitchingLib(True, False),\n",
    "    # cx.BalancingCVStitchingLib(True, False),\n",
    "    cx.compute_mock_similarity\n",
    ")\n",
    "plt.imshow(S)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "np.nanmin(S)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save((gca.graph, gcb.graph, S), \"stitching-problem-2.pickle\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gca.graph.write_graphmlz(\"stitching-problem-2-graph_a.graphml.gz\")\n",
    "gcb.graph.write_graphmlz(\"stitching-problem-2-graph_b.graphml.gz\")\n",
    "np.savetxt(\"stitching-problem-2-similarity.txt.gz\", S)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# stitched = cx.construct_trained_cx_network_stitching(\n",
    "#     dataset=dataset,\n",
    "#     dev=dev,\n",
    "#     net_a=gca,\n",
    "#     net_b=gcb,\n",
    "#     X_in_many=X,\n",
    "#     ensemblers=[ly.LinearCombine([0.5, 0.5])],\n",
    "#     compute_similarity=cx.compute_mock_similarity,\n",
    "#     feature_shape_should_match=False,\n",
    "#     batch_size=batch_size,\n",
    "\n",
    "#     pretrain_cx_network = False\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from IPython.display import SVG\n",
    "# with open(\"g.dot\", \"w\") as f:\n",
    "#     stitched[0].to_dot(f, include_ord_label=True)\n",
    "# ! dot g.dot -Tsvg -og.png\n",
    "# SVG(\"g.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Alternative that initializes the layers a bit more intelligently.\n",
    "class CVInitStitchingLib(cx.StitchingLib):\n",
    "\n",
    "    def __init__(self, image_shape_should_match, feature_shape_should_match):\n",
    "        self.image_shape_should_match = image_shape_should_match\n",
    "        self.feature_shape_should_match = feature_shape_should_match\n",
    "\n",
    "    def characterize_fm(self, v, fm, graph=None):\n",
    "        if isinstance(fm, torch.Tensor):\n",
    "            v[\"ty\"] = \"tensor\"\n",
    "            v[\"sh\"] = list(fm.shape)\n",
    "            with torch.no_grad():\n",
    "                v[\"std\"], v[\"mean\"] = torch.std_mean(fm)\n",
    "        else:\n",
    "            v[\"ty\"] = \"unk\"\n",
    "\n",
    "    def can_stitch(self, a, b):\n",
    "        # If types are unknown - do not allow stitching at these points.\n",
    "        if a[\"ty\"] == \"unk\": return False\n",
    "        if b[\"ty\"] == \"unk\": return False\n",
    "\n",
    "        if a[\"ty\"] == \"tensor\" and b[\"ty\"] == \"tensor\":\n",
    "            sh_a = a[\"sh\"]\n",
    "            sh_b = b[\"sh\"]\n",
    "\n",
    "            # if fm_a.shape[0] != fm_b.shape[0]:\n",
    "            #     continue\n",
    "            if self.image_shape_should_match and not (sh_a[2:] == sh_b[2:]):\n",
    "                return False\n",
    "            if self.feature_shape_should_match and not (sh_a[1] == sh_b[1]):\n",
    "                return False\n",
    "            return True\n",
    "        \n",
    "        # Cascade through\n",
    "        return False\n",
    "\n",
    "    def create_stitch(self, a, b):\n",
    "        if a[\"ty\"] == \"tensor\" and b[\"ty\"] == \"tensor\":\n",
    "            sh_a = a[\"sh\"]\n",
    "            sh_b = b[\"sh\"]\n",
    "\n",
    "            # normally we preserve mean & variance with the chosen init\n",
    "            # This strategy sets the weights & biases accordingly\n",
    "            # albeit under an uniform assumption.\n",
    "            offset = b[\"mean\"] - a[\"mean\"]\n",
    "            scale = b[\"std\"] / a[\"std\"]\n",
    "\n",
    "            num_features_in = sh_a[1]\n",
    "            num_features_out = sh_b[1]\n",
    "            if len(sh_a) == 4 and len(sh_b) == 4 and sh_a[2:] == sh_b[2:]:\n",
    "                stitch = ly.Conv2d(\n",
    "                    num_features_in, num_features_out, kernel_size=(1, 1)\n",
    "                )\n",
    "                with torch.no_grad():\n",
    "                    stitch.layer.bias += offset\n",
    "                    stitch.layer.weight *= scale\n",
    "                return stitch\n",
    "            elif len(sh_a) == 2 and len(sh_b) == 2:\n",
    "                stitch = ly.Linear(num_features_in, num_features_out)\n",
    "                with torch.no_grad():\n",
    "                    stitch.layer.bias += offset\n",
    "                    stitch.layer.weight *= scale\n",
    "                return stitch\n",
    "            else:\n",
    "                raise Exception(\n",
    "                    f\"cannot join items. No merging layer defined for shapes a: {sh_a} b: {sh_b}\"\n",
    "                )\n",
    "        raise Exception(\n",
    "                    f\"cannot join items. No stitching layer defined between layers from type {a['ty']} to {b['ty']}\"\n",
    "                )\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "summarywriter = SummaryWriter(\"./logs/train-stitch-imagenet-3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raise ValueError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "import torch\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size=16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.autograd.set_detect_anomaly(False)\n",
    "# dev = torch.device(\"cuda:1\")\n",
    "stitched = cx.construct_trained_cx_network_stitching(\n",
    "    dataset=dataset,\n",
    "    dev=dev,\n",
    "    net_a=gca,\n",
    "    net_b=gcb,\n",
    "    X_in_many=X,\n",
    "    ensemblers=[ly.LinearCombine([0.5, 0.5])],\n",
    "    stitching_library=cx.CVStitchingLib(True, False),\n",
    "    # stitching_library=CVInitStitchingLib(True, False),\n",
    "    # stitching_library=cx.BalancingCVStitchingLib(True, False),\n",
    "    # compute_similarity=cx.compute_mock_similarity,\n",
    "    compute_similarity=cx.compute_mock_similarity,\n",
    "    feature_shape_should_match=False,\n",
    "    batch_size=batch_size,\n",
    "\n",
    "    lr_pretrain=1e-3,\n",
    "    weight_decay_pretrain=1e-5,\n",
    "\n",
    "    num_epochs_pretrain = 3,\n",
    "    pretrain_cx_network = True,\n",
    "    summarywriter=summarywriter,\n",
    "    immediately_backprop = False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(stitched[1].joiners)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(stitched[1].joiners) * 2 + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for es in stitched[1].joiners:\n",
    "    for e in es:\n",
    "        e.agg = None\n",
    "stitched[1].output_switch.agg = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stitched_simpl = (stitched[0], cx.SitchingInfo(stitched[1].joiners, stitched[1].output_switch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(stitched_simpl, \"stitched-imagenet-b-a-resnet50-b-resnext50_32x4d.th\")\n",
    "# torch.save(stitched_simpl, \"stitched-imagenet-a-resnet152-b-efficientnet-b4--r.th\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# stitched = torch.load(\"stitched-imagenet-a-resnet152-b-efficientnet-b4.th\")\n",
    "stitched = stitched_simpl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stitchnet, stitchinfo = stitched"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note: Only the output above is used during the tests.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train supernetwork by randomly sampling layers\n",
    "for m in stitchnet.submodules:\n",
    "    if not isinstance(m, cx.CXN): continue\n",
    "    m.active = [0, 1]\n",
    "    m.p = [0.9, 0.1]\n",
    "    m.determine_p()\n",
    "    m.randomize_per_sample = True\n",
    "stitchinfo.output_switch.active = [0, 1, 2]\n",
    "stitchinfo.output_switch.p = None\n",
    "stitchinfo.output_switch.determine_p()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "for m in stitchnet.submodules:\n",
    "    if not isinstance(m, cx.CXN): continue\n",
    "    if m.p is None: continue\n",
    "    if m.p[-1] is None:\n",
    "        m.p = None\n",
    "        continue\n",
    "    m.p = np.cumsum(m.p)\n",
    "    m.p /= m.p[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.tensorboard import SummaryWriter\n",
    "summarywriter = SummaryWriter(\"./logs/refine-imagenet\")\n",
    "stitchneti = NeuralNetIndividual(stitchnet)\n",
    "problem.train_network(dev, stitchneti, lr=1e-4, weight_decay=1e-5, num_epochs=5, minout_nan=True, batch_size=batch_size, raise_on_nan_loss=False, summarywriter=summarywriter)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stitchneti.net.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save((stitchnet, stitchinfo), \"stitch-train-test.th\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embed computational cost info\n",
    "import torchinfo\n",
    "import recomb.eval_costs as ec\n",
    "cost_summary = torchinfo.summary(stitchnet, input_data=[X_o])\n",
    "ec.embed_cost_stats_in_model(cost_summary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print((\"The stitched neural network has \"\n",
    "       f\"{len(stitchinfo.joiners)} matches\"\n",
    "       ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate predetermined neighborhood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate reference networks\n",
    "reference_q = []\n",
    "stitchinfo.output_switch.active = 2\n",
    "stitchinfo.output_switch.simplify = True\n",
    "\n",
    "for j in stitchinfo.joiners:\n",
    "        j[0].active = 0\n",
    "        j[1].active = 0\n",
    "        j[0].simplify = True\n",
    "        j[1].simplify = True\n",
    "    \n",
    "stitchnet_pruned = stitchnet.to_graph()\n",
    "stitchnet_pruned.prune_unused()\n",
    "total_mult_adds, total_bytes = ec.evaluate_compute_cost(stitchnet_pruned)\n",
    "\n",
    "neti_os = NeuralNetIndividual(stitchnet_pruned)\n",
    "accuracy, loss = problem.evaluate_network(dev, neti_os, objective=\"both\")\n",
    "reference_q.append((accuracy, loss, total_bytes, total_mult_adds, cx.convert_stitcher_to_genotype(stitchinfo, stringify=False)))\n",
    "\n",
    "stitchinfo.output_switch.active = 1\n",
    "stitchnet_pruned = stitchnet.to_graph()\n",
    "stitchnet_pruned.prune_unused()\n",
    "total_mult_adds, total_bytes = ec.evaluate_compute_cost(stitchnet_pruned)\n",
    "\n",
    "neti_os = NeuralNetIndividual(stitchnet_pruned)\n",
    "accuracy, loss = problem.evaluate_network(dev, neti_os, objective=\"both\")\n",
    "reference_q.append((accuracy, loss, total_bytes, total_mult_adds, cx.convert_stitcher_to_genotype(stitchinfo, stringify=False)))\n",
    "\n",
    "stitchinfo.output_switch.active = 0\n",
    "stitchnet_pruned = stitchnet.to_graph()\n",
    "stitchnet_pruned.prune_unused()\n",
    "total_mult_adds, total_bytes = ec.evaluate_compute_cost(stitchnet_pruned)\n",
    "\n",
    "neti_os = NeuralNetIndividual(stitchnet_pruned)\n",
    "accuracy, loss = problem.evaluate_network(dev, neti_os, objective=\"both\")\n",
    "reference_q.append((accuracy, loss, total_bytes, total_mult_adds, cx.convert_stitcher_to_genotype(stitchinfo, stringify=False)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ensembles_a = []\n",
    "ensembles_b = []\n",
    "start_a_end_b = []\n",
    "start_b_end_a = []\n",
    "\n",
    "# note - usually 1, but due to the large amount of matches, this has been\n",
    "# increased so that we can evaluate blocks of solutions instead.\n",
    "step = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate neighborhood of networks\n",
    "offset = 9\n",
    "stitchinfo.output_switch.active = 2\n",
    "stitchinfo.output_switch.simplify = True\n",
    "\n",
    "for j in stitchinfo.joiners:\n",
    "        j[0].active = 0\n",
    "        j[1].active = 0\n",
    "        j[0].simplify = True\n",
    "        j[1].simplify = True\n",
    "\n",
    "for i in range(offset, len(stitchinfo.joiners), step):\n",
    "    j = stitchinfo.joiners[i]\n",
    "    j[0].active = 0\n",
    "    j[1].active = 1\n",
    "\n",
    "    stitchnet_pruned = stitchnet.to_graph()\n",
    "    stitchnet_pruned.prune_unused()\n",
    "\n",
    "    # Get compute & memory requirements\n",
    "    # s = torchinfo.summary(stitchnet_pruned, input_data=[X])\n",
    "    total_mult_adds, total_bytes = ec.evaluate_compute_cost(stitchnet_pruned)\n",
    "\n",
    "    neti_os = NeuralNetIndividual(stitchnet_pruned)\n",
    "    accuracy, loss = problem.evaluate_network(dev, neti_os, objective=\"both\")\n",
    "    ensembles_a.append((accuracy, loss, total_bytes, total_mult_adds, cx.convert_stitcher_to_genotype(stitchinfo, stringify=False)))\n",
    "    \n",
    "    j[0].active = 0\n",
    "    j[1].active = 0\n",
    "\n",
    "for i in range(offset, len(stitchinfo.joiners), step):\n",
    "    j = stitchinfo.joiners[i]\n",
    "    j[0].active = 1\n",
    "    j[1].active = 0\n",
    "\n",
    "    stitchnet_pruned = stitchnet.to_graph()\n",
    "    stitchnet_pruned.prune_unused()\n",
    "\n",
    "    # Get compute & memory requirements\n",
    "    # s = torchinfo.summary(stitchnet_pruned, input_data=[X])\n",
    "    total_mult_adds, total_bytes = ec.evaluate_compute_cost(stitchnet_pruned)\n",
    "\n",
    "    neti_os = NeuralNetIndividual(stitchnet_pruned)\n",
    "    accuracy, loss = problem.evaluate_network(dev, neti_os, objective=\"both\")\n",
    "    ensembles_b.append((accuracy, loss, total_bytes, total_mult_adds, cx.convert_stitcher_to_genotype(stitchinfo, stringify=False)))\n",
    "    \n",
    "    j[0].active = 0\n",
    "    j[1].active = 0\n",
    "\n",
    "stitchinfo.output_switch.active = 1\n",
    "for i in range(offset, len(stitchinfo.joiners), step):\n",
    "    j = stitchinfo.joiners[i]\n",
    "    j[0].active = 0\n",
    "    j[1].active = 1\n",
    "\n",
    "    stitchnet_pruned = stitchnet.to_graph()\n",
    "    stitchnet_pruned.prune_unused()\n",
    "\n",
    "    # Get compute & memory requirements\n",
    "    # s = torchinfo.summary(stitchnet_pruned, input_data=[X])\n",
    "    total_mult_adds, total_bytes = ec.evaluate_compute_cost(stitchnet_pruned)\n",
    "\n",
    "    neti_os = NeuralNetIndividual(stitchnet_pruned)\n",
    "    accuracy, loss = problem.evaluate_network(dev, neti_os, objective=\"both\")\n",
    "    start_a_end_b.append((accuracy, loss, total_bytes, total_mult_adds, cx.convert_stitcher_to_genotype(stitchinfo, stringify=False)))\n",
    "    \n",
    "    j[0].active = 0\n",
    "    j[1].active = 0\n",
    "\n",
    "stitchinfo.output_switch.active = 0\n",
    "for i in range(offset, len(stitchinfo.joiners), step):\n",
    "    j = stitchinfo.joiners[i]\n",
    "    j[0].active = 1\n",
    "    j[1].active = 0\n",
    "\n",
    "    stitchnet_pruned = stitchnet.to_graph()\n",
    "    stitchnet_pruned.prune_unused()\n",
    "\n",
    "    # Get compute & memory requirements\n",
    "    # s = torchinfo.summary(stitchnet_pruned, input_data=[X])\n",
    "    total_mult_adds, total_bytes = ec.evaluate_compute_cost(stitchnet_pruned)\n",
    "\n",
    "    neti_os = NeuralNetIndividual(stitchnet_pruned)\n",
    "    accuracy, loss = problem.evaluate_network(dev, neti_os, objective=\"both\")\n",
    "    start_b_end_a.append((accuracy, loss, total_bytes, total_mult_adds, cx.convert_stitcher_to_genotype(stitchinfo, stringify=False)))\n",
    "    \n",
    "    j[0].active = 0\n",
    "    j[1].active = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_schema = [\"accuracy\", \"loss\", \"total bytes\", \"multiply-adds\", \"genotype\"]\n",
    "samples_reference = pl.DataFrame(reference_q, schema=df_schema).\\\n",
    "    with_columns([\n",
    "        pl.lit(pl.Series([\"ensemble\", \"b\", \"a\"])).alias(\"set\"),\n",
    "        pl.lit(False).alias(\"contains stitch\"),\n",
    "    ])\n",
    "samples_ensemble_a = pl.DataFrame(ensembles_a, schema=df_schema).\\\n",
    "    with_columns([\n",
    "        pl.lit(\"ensemble-major-a\").alias(\"set\"),\n",
    "        pl.lit(True).alias(\"contains stitch\"),\n",
    "    ])\n",
    "samples_ensemble_b = pl.DataFrame(ensembles_b, schema=df_schema).\\\n",
    "    with_columns([\n",
    "        pl.lit(\"ensemble-major-b\").alias(\"set\"),\n",
    "        pl.lit(True).alias(\"contains stitch\"),\n",
    "    ])\n",
    "samples_ab = pl.DataFrame(start_a_end_b, schema=df_schema).\\\n",
    "    with_columns([\n",
    "        pl.lit(\"stitch-a-to-b\").alias(\"set\"),\n",
    "        pl.lit(True).alias(\"contains stitch\"),\n",
    "    ])\n",
    "samples_ba = pl.DataFrame(start_b_end_a, schema=df_schema).\\\n",
    "    with_columns([\n",
    "        pl.lit(\"stitch-b-to-a\").alias(\"set\"),\n",
    "        pl.lit(True).alias(\"contains stitch\"),\n",
    "    ])\n",
    "\n",
    "samples = pl.concat([\n",
    "    samples_reference,\n",
    "    samples_ensemble_a,\n",
    "    samples_ensemble_b,\n",
    "    samples_ab,\n",
    "    samples_ba,\n",
    "]).with_columns(\n",
    "    pl.col(\"loss\").clip(0.0, 4.0).alias(\"loss-clip\")\n",
    ")\n",
    "samples.write_ipc(\"resnet-efficientnet-stitch-samples.arrow\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot approximation front"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = pl.read_ipc(\"resnet-efficientnet-stitch-samples.arrow\")\n",
    "\n",
    "# Extract some rows of reference interest\n",
    "dfcna = samples[2]\n",
    "dfcnb = samples[1]\n",
    "dfcnens = samples[0]\n",
    "\n",
    "# \n",
    "improvement_direction = {\n",
    "    \"accuracy\": 1,\n",
    "    \"loss\": -1,\n",
    "    \"loss-clip\": -1,\n",
    "    \"total bytes\": -1,\n",
    "    \"multiply-adds\": -1,\n",
    "    # \"genotype\": 0, # -- not a criterion\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(stitchinfo.joiners) * 4 * 15 / 60 / 60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# How many seconds per evaluated sample?\n",
    "number_of_minutes = 16 * 8\n",
    "number_of_seconds = 0\n",
    "number_of_samples = len(samples) - 3\n",
    "seconds_total = number_of_minutes * 60 + number_of_seconds\n",
    "seconds_per_sample = seconds_total / number_of_samples\n",
    "\n",
    "print(f\"spent {number_of_minutes}m{number_of_seconds}s \"\n",
    "      f\"to evaluate {number_of_samples} samples.\\n\"\n",
    "      f\"Resulting in a cost of {seconds_per_sample}s per sample.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute pareto set from these points, with respect to these\n",
    "# two criteria / objectives\n",
    "c0 = \"accuracy\"\n",
    "c1 = \"multiply-adds\"\n",
    "\n",
    "samples_pareto = (samples.lazy()\n",
    "    .sort(c0, descending=improvement_direction[c0] > 0)\n",
    "    .with_columns((pl.col(c1) * -improvement_direction[c1]).alias(\"c1-min\"))\n",
    "    .with_columns((pl.col(\"c1-min\")).cummin().alias(\"mv\"))\n",
    "    .with_columns((pl.col(\"c1-min\") < pl.col(\"mv\").shift(1)).alias(\"is pareto\")).fill_null(True)\n",
    "    .filter(pl.col(\"is pareto\"))\n",
    ").collect()\n",
    "\n",
    "samples_pareto_stitch_only = (samples.lazy()\n",
    "    .filter(pl.col(\"contains stitch\"))\n",
    "    .sort(c0, descending=improvement_direction[c0] > 0)\n",
    "    .with_columns((pl.col(c1) * -improvement_direction[c1]).alias(\"c1-min\"))\n",
    "    .with_columns((pl.col(\"c1-min\")).cummin().alias(\"mv\"))\n",
    "    .with_columns((pl.col(\"c1-min\") < pl.col(\"mv\").shift(1)).alias(\"is pareto\")).fill_null(True)\n",
    "    .filter(pl.col(\"is pareto\"))\n",
    ").collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for sn, df in samples.filter(pl.col(\"contains stitch\")).group_by(\"set\", maintain_order=True):\n",
    "    plt.scatter(df[c0], df[c1], label=sn, s=1.0)\n",
    "\n",
    "plt.scatter(samples_pareto[c0], samples_pareto[c1], alpha=0.4, marker=\"s\", color=\"grey\")\n",
    "plt.scatter(samples_pareto_stitch_only[c0], samples_pareto_stitch_only[c1], s=20.0, alpha=0.5, color=\"grey\")\n",
    "\n",
    "plt.scatter(dfcna[c0], dfcna[c1], label=\"a\", marker='x')\n",
    "plt.scatter(dfcnb[c0], dfcnb[c1], label=\"b\", marker='x')\n",
    "plt.scatter(dfcnens[c0], dfcnens[c1], label=\"ensemble\", marker='x')\n",
    "\n",
    "def get_direction_arrow(c):\n",
    "    return '->' if improvement_direction[c] > 0 else '<-'\n",
    "\n",
    "plt.xlabel(f\"{c0} ({get_direction_arrow(c0)})\")\n",
    "plt.ylabel(f\"{c1} ({get_direction_arrow(c1)})\")\n",
    "plt.legend(loc='upper left',\n",
    "           bbox_to_anchor=(1.0, 1.0),\n",
    "           fancybox=False,\n",
    "           shadow=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Potential Points of Improvement?**\n",
    "1. Pretrain for longer? (e.g. specific stopping condition?)\n",
    "2. Train using actual loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist2d(samples[c0], samples[c1])\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Some test evaluations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev2 = torch.device(\"cuda:1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate original networks\n",
    "neti_a = NeuralNetIndividual(gca)\n",
    "neti_b = NeuralNetIndividual(gcb)\n",
    "problem.evaluate_network(dev2, neti_a, objective=\"both\"),\\\n",
    "    problem.evaluate_network(dev, neti_b, objective=\"both\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stitchneti = NeuralNetIndividual(stitchnet)\n",
    "for j in stitchinfo.joiners:\n",
    "    j[0].active = 0\n",
    "    j[1].active = 0\n",
    "\n",
    "stitchinfo.output_switch.active = 0\n",
    "roa = problem.evaluate_network(dev, stitchneti, objective=\"both\")\n",
    "stitchinfo.output_switch.active = 1\n",
    "rob = problem.evaluate_network(dev, stitchneti, objective=\"both\")\n",
    "roa, rob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stitchneti = NeuralNetIndividual(stitchnet)\n",
    "for j in stitchinfo.joiners:\n",
    "    j[0].active = 0\n",
    "    j[1].active = 0\n",
    "stitchinfo.output_switch.active = 2\n",
    "j = stitchinfo.joiners[18]\n",
    "# j[0].active = 0\n",
    "# j[1].active = 1\n",
    "j[0].active = 1\n",
    "j[1].active = 0\n",
    "\n",
    "problem.evaluate_network(dev, stitchneti, objective=\"both\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "\n",
    "\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "recombnet",
   "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": 2
}