hypnettorch/tutorials/wrn_chunking_strategies.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chunking strategies for a Wide-ResNet\n",
    "\n",
    "This tutorial shows how to utilize a hypernet container [HContainer](../hnets/hnet_container.py) and class [StructuredHMLP](../hnets/structured_mlp_hnet.py) (a certain kind of hypernetwork that allows *smart* chunking) in combination with a Wide-ResNet [WRN](../mnets/wide_resnet.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ensure code of repository is visible to this tutorial.\n",
    "import sys\n",
    "sys.path.insert(0, '..')\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from hypnettorch.hnets.structured_hmlp_examples import wrn_chunking\n",
    "from hypnettorch.hnets import HContainer, StructuredHMLP\n",
    "from hypnettorch.mnets import WRN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instantiate a WRN-28-10-B(3,3)\n",
    "\n",
    "First, we instantiate a WRN-28-10 (i.e., a WRN containing $28$ convolutional layers (and an additional fully-connected output layer) and a widening factor $k=10$) with no internal weights (`no_weights=True`). Thus, it's weights are expected to originate externally (in our case from a hypernetwork) and to be passed to its `forward` method.\n",
    "\n",
    "In particular, we are interested in instantiating a network that matches the one used in the study [Sacramento et al., \"Economical ensembles with hypernetworks\", 2020](https://arxiv.org/abs/2007.12927) (accessed August 18th, 2020). Therefore, the convolutional layers won't have bias terms (but the final fully-connected layer will)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating a WideResnet \"WRN-28-10-B(3,3)\" with 36479194 weights. The network uses batchnorm.\n"
     ]
    }
   ],
   "source": [
    "net = WRN(in_shape=(32, 32, 3), num_classes=10, n=4, k=10,\n",
    "          num_feature_maps=(16, 16, 32, 64), use_bias=False,\n",
    "          use_fc_bias=True, no_weights=False, use_batch_norm=True,\n",
    "          dropout_rate=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Reproduce the chunking strategy from Sacramento et al.\n",
    "\n",
    "We first design a hypernetwork that matches the chunking strategy described in [Sacramento et al.](https://arxiv.org/abs/2007.12927). Thus, not all parameters are produced by a hypernetwork. Batchnorm weights will be shared among conditions (in their case, each condition represents one ensemble member), while the output layer weights will be condition-specific (ensemble-member-specific). The remaining weight are produced via linear hypernetworks (no bias terms in the hypernets) using a specific chunking strategy, which is described in the paper and in the docstring of function [wrn_chunking](../hnets/structured_hmlp_examples.py). To realize the mixture between shared weights (batchnorm), condition-specific weights (output weights) and hypernetwork-produced weights, we employ the special hypernetwork class [HContainer](../hnets/hnet_container.py).\n",
    "\n",
    "We first create an instance of class [StructuredHMLP](../hnets/structured_mlp_hnet.py) for all hypernetwork-produced weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created Structured Chunked MLP Hypernet.\n",
      "It manages 8 full hypernetworks internally that produce 42 chunks in total.\n",
      "The internal hypernetworks have a combined output size of 2816432 compared to 36454832 weights produced by this network.\n",
      "Hypernetwork with 37462680 weights and 36454832 outputs (compression ratio: 1.03).\n",
      "The network consists of 37457120 unconditional weights (37457120 internally maintained) and 5560 conditional weights (5560 internally maintained).\n"
     ]
    }
   ],
   "source": [
    "# Number of conditions (ensemble members). Arbitrarily chosen!\n",
    "num_conds = 10\n",
    "\n",
    "# Split the network's parameter shapes into shapes corresponding to batchnorm-weights,\n",
    "# hypernet-produced weights and output weights.\n",
    "# Here, we make use of implementation specific knowledge, which could also be retrieved\n",
    "# via the network's \"param_shapes_meta\" attribute, which contains meta information\n",
    "# about all parameters.\n",
    "bn_shapes = net.param_shapes[:2*len(net.batchnorm_layers)] # Batchnorm weight shapes\n",
    "hnet_shapes = net.param_shapes[2*len(net.batchnorm_layers):-2] # Conv layer weight shapes\n",
    "out_shapes = net.param_shapes[-2:] # Output layer weight shapes\n",
    "\n",
    "# This function already defines the network chunking in the same way the paper\n",
    "# specifies it.\n",
    "chunk_shapes, num_per_chunk, assembly_fct = wrn_chunking(net, ignore_bn_weights=True,\n",
    "                                                         ignore_out_weights=True,\n",
    "                                                         gcd_chunking=False)\n",
    "# Taken from table S1 in the paper.\n",
    "chunk_emb_sizes = [10, 7, 14, 14, 14, 7, 7, 7]\n",
    "\n",
    "# Important, the underlying hypernetworks should be linear, i.e., no hidden layers:\n",
    "# ``layers': []``\n",
    "# They also should not use bias vectors -> hence, weights are simply generated via a\n",
    "# matrix vector product (chunk embedding input times hypernet, which is a weight matrix).\n",
    "# Note, we make the chunk embeddings conditional and tell the hypernetwork, that\n",
    "# it doesn't have to expect any other input except those learned condition-specific\n",
    "# embeddings.\n",
    "shnet = StructuredHMLP(hnet_shapes, chunk_shapes, num_per_chunk, chunk_emb_sizes,\n",
    "                       {'layers': [], 'use_bias': False}, assembly_fct,\n",
    "                       cond_chunk_embs=True, uncond_in_size=0,\n",
    "                       cond_in_size=0, num_cond_embs=num_conds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we combine the above produce `shnet` with shared batchnorm weights and condition-specific output weights in an instance of class [HContainer](../hnets/hnet_container.py), which will represent the final hypernetwork."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created Hypernet Container for 1 hypernet(s). Container maintains 50 plain unconditional parameter tensors. Container maintains 2 plain conditional parameter tensors for each of 10 condiditions.\n",
      "Hypernetwork with 37544732 weights and 36479194 outputs (compression ratio: 1.03).\n",
      "The network consists of 37475072 unconditional weights (37475072 internally maintained) and 69660 conditional weights (69660 internally maintained).\n"
     ]
    }
   ],
   "source": [
    "# We first have to create a simple function handle that tells the `HContainer` how to\n",
    "# recombine the batchnorm-weights, hypernet-produced weights and output weights.\n",
    "def simple_assembly_func(list_of_hnet_tensors, uncond_tensors, cond_tensors):\n",
    "    # `list_of_hnet_tensors`: Contains outputs of all linear hypernets (conv \n",
    "    #                         layer weights).\n",
    "    # `uncond_tensors`: Contains the single set of shared batchnorm weights.\n",
    "    # `cond_tensors`: Contains the condition-specific output weights.\n",
    "    return uncond_tensors + list_of_hnet_tensors[0] + cond_tensors\n",
    "\n",
    "hnet = HContainer(net.param_shapes, simple_assembly_func, hnets=[shnet],\n",
    "                  uncond_param_shapes=bn_shapes, cond_param_shapes=out_shapes,\n",
    "                  num_cond_embs=num_conds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create sample predictions for 3 different ensemble members."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction of ensemble member 2: [[0.099, 0.102, 0.102, 0.102, 0.099, 0.104, 0.097, 0.097, 0.097, 0.1  ]]\n",
      "Prediction of ensemble member 3: [[0.1  , 0.095, 0.102, 0.1  , 0.101, 0.101, 0.102, 0.102, 0.101, 0.097]]\n",
      "Prediction of ensemble member 7: [[0.101, 0.098, 0.099, 0.1  , 0.106, 0.098, 0.098, 0.1  , 0.101, 0.099]]\n"
     ]
    }
   ],
   "source": [
    "# Batch of inputs.\n",
    "batch_size = 1\n",
    "x = torch.rand((batch_size, 32*32*3))\n",
    "\n",
    "# Which ensemble members to consider?\n",
    "cond_ids = [2,3,7]\n",
    "\n",
    "# Generate weights for ensemble members defined above.\n",
    "weights = hnet.forward(cond_id=cond_ids)\n",
    "\n",
    "# Compute prediction for each ensemble member.\n",
    "for i in range(len(cond_ids)):\n",
    "    pred = net.forward(x, weights=weights[i])\n",
    "    # Apply softmax.\n",
    "    pred = torch.nn.functional.softmax(pred, dim=1).cpu().detach().numpy()\n",
    "    print('Prediction of ensemble member %d: %s' \\\n",
    "          % (cond_ids[i], np.array2string(pred, precision=3, separator=', ')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a batch-ensemble network\n",
    "\n",
    "Now, we consider the special case where all parameters are shared except for batchnorm weights and output weights. Thus, no \"hypernetwork\" are required. Though, we use the class [HContainer](../hnets/hnet_container.py) for convinience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created Hypernet Container for 0 hypernet(s). Container maintains 28 plain unconditional parameter tensors. Container maintains 52 plain conditional parameter tensors for each of 10 condiditions.\n",
      "Hypernetwork with 36698452 weights and 36479194 outputs (compression ratio: 1.01).\n",
      "The network consists of 36454832 unconditional weights (36454832 internally maintained) and 243620 conditional weights (243620 internally maintained).\n"
     ]
    }
   ],
   "source": [
    "def simple_assembly_func2(list_of_hnet_tensors, uncond_tensors, cond_tensors):\n",
    "    # `list_of_hnet_tensors`: None\n",
    "    # `uncond_tensors`: Contains all conv layer weights.\n",
    "    # `cond_tensors`: Contains the condition-specific batchnorm and output weights.\n",
    "    return cond_tensors[:-2] + uncond_tensors + cond_tensors[-2:]\n",
    "\n",
    "hnet2 = HContainer(net.param_shapes, simple_assembly_func2, hnets=None,\n",
    "                   uncond_param_shapes=hnet_shapes,\n",
    "                   cond_param_shapes=bn_shapes+out_shapes,\n",
    "                   num_cond_embs=num_conds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction of ensemble member 2: [[0.1  , 0.099, 0.1  , 0.096, 0.101, 0.103, 0.098, 0.097, 0.102, 0.103]]\n",
      "Prediction of ensemble member 3: [[0.102, 0.098, 0.098, 0.102, 0.101, 0.1  , 0.098, 0.102, 0.102, 0.097]]\n",
      "Prediction of ensemble member 7: [[0.1  , 0.099, 0.096, 0.096, 0.102, 0.101, 0.102, 0.101, 0.101, 0.103]]\n"
     ]
    }
   ],
   "source": [
    "# Batch of inputs.\n",
    "batch_size = 1\n",
    "x = torch.rand((batch_size, 32*32*3))\n",
    "\n",
    "# Which ensemble members to consider?\n",
    "cond_ids = [2,3,7]\n",
    "\n",
    "# Generate weights for ensemble members defined above.\n",
    "weights = hnet2.forward(cond_id=cond_ids)\n",
    "\n",
    "# Compute prediction for each ensemble member.\n",
    "for i in range(len(cond_ids)):\n",
    "    pred = net.forward(x, weights=weights[i])\n",
    "    # Apply softmax.\n",
    "    pred = torch.nn.functional.softmax(pred, dim=1).cpu().detach().numpy()\n",
    "    print('Prediction of ensemble member %d: %s' \\\n",
    "          % (cond_ids[i], np.array2string(pred, precision=3, separator=', ')))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}