jupyter/alexnet_insights.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Alexnet insights: visualizing the pruning process\n",
    "\n",
    "This notebook examines the results of pruning Alexnet using sensitivity pruning, through a few chosen visualizations created from checkpoints created during pruning.  We also compare the results of an element-wise pruning session, with 2D (kernel) regularization.\n",
    "\n",
    "For the notebook, we pruned Alexnet using sensitivity pruning and captured the checkpoints after the first epoch ends (epoch 0) and after the last epoch ends (epoch 89).\n",
    "\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "1. [Load the training checkpoints](#Load-the-training-checkpoints)\n",
    "2. [Let's see some statistics](#Let's-see-some-statistics)\n",
    "3. [Compare weights distributions](#Compare-weights-distributions)\n",
    "4. [Visualize the weights](#Visualize-the-weights)<br>\n",
    "    4.1. [Fully-connected layers](#Fully-connected-layers)<br>\n",
    "    4.2. [Convolutional layers](#Convolutional-layers)<br>\n",
    "    4.3. [Kernel pruning](#Kernel-pruning)<br>\n",
    "    4.4. [Let's isolate just the 2D kernels](#Let's-isolate-just-the-2D-kernels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib \n",
    "\n",
    "# Load some common jupyter code\n",
    "%run './distiller_jupyter_helpers.ipynb'\n",
    "from distiller.models import create_model\n",
    "from distiller.apputils import *\n",
    "import qgrid\n",
    "\n",
    "from ipywidgets import *\n",
    "from bqplot import *\n",
    "import bqplot.pyplot as bqplt\n",
    "from functools import partial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the training checkpoints\n",
    "Load the checkpoint captured after one pruning event, and fine-tuning for one epoch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch0_model = create_model(False, 'imagenet', 'alexnet', parallel=True)\n",
    "checkpoint_file = \"../examples/classifier_compression/alexnet.checkpoint.0.pth.tar\"\n",
    "try:\n",
    "    load_checkpoint(epoch0_model, checkpoint_file)\n",
    "except NameError as e:\n",
    "    print(\"Did you forget to download the checkpoint file?\")\n",
    "    raise e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the checkpoint captured at the end of the pruning and fine-tuning process:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch89_model = create_model(False, 'imagenet', 'alexnet', parallel=True)\n",
    "checkpoint_file = \"../examples/classifier_compression/alexnet.checkpoint.89.pth.tar\"\n",
    "try:\n",
    "    load_checkpoint(epoch89_model, checkpoint_file)\n",
    "except Exception as e:\n",
    "    print(\"Did you forget to download the checkpoint file?\")\n",
    "    raise e   \n",
    "    \n",
    "sparse_model = epoch89_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load a pre-trained dense Alexnet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_model = create_model(True, 'imagenet', 'alexnet', parallel=True)\n",
    "dense_model = pretrained_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to compare the output of element-wise pruning, to a similar schedule but which also adds 2D (kernel-wise) Lasso regularization, so we load the last checkpoint of that pruning session.\n",
    "The schedule is available at: ```distiller/examples/hybrid/alexnet.schedule_sensitivity_2D-reg.yaml```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reg2D_model = create_model(False, 'imagenet', 'alexnet', parallel=True)\n",
    "checkpoint_file = \"../examples/classifier_compression/checkpoint.alexnet.schedule_sensitivity_2D-reg.pth.tar\"\n",
    "try:\n",
    "    load_checkpoint(reg2D_model, checkpoint_file);\n",
    "except Exception as e:\n",
    "    print(\"Did you forget to download the checkpoint file?\")\n",
    "    raise e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a dictionary of the models, with name as key, so that we can refer to it later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "models_dict = {'Dense': dense_model, 'Sparse': sparse_model, 'Epoch 0': epoch0_model, '2D-Sparse': reg2D_model}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's see some statistics\n",
    "\n",
    "You can use the dropwdown widget to choose which model to display.\n",
    "\n",
    "You can also choose to display the sparsity or the density of the tensors.  These are reported for various granularities (structures) of sparsities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def view_data(what, model_choice):\n",
    "    df_sparsity = distiller.weights_sparsity_summary(models_dict[model_choice])\n",
    "    if what == 'Density':\n",
    "        for granularity in ['Fine (%)', 'Ch (%)', '2D (%)', '3D (%)']:\n",
    "            df_sparsity[granularity] = 100 - df_sparsity[granularity]\n",
    "    display(df_sparsity)\n",
    "\n",
    "model_dropdown = Dropdown(description='Model:', options=models_dict.keys())\n",
    "display_radio = widgets.RadioButtons(options=['Sparsity', 'Density'], value='Sparsity', description='Display:')\n",
    "interact(view_data, what=display_radio, model_choice=model_dropdown);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare weights distributions\n",
    "\n",
    "Compare the distributions of the weight tensors in the sparse and dense models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "nbins = 100\n",
    "\n",
    "def get_hist2(model, nbins, param_name, remove_zeros):\n",
    "    tensor = flatten(model.state_dict()[param_name])\n",
    "    if remove_zeros:\n",
    "        tensor = tensor[tensor != 0]\n",
    "    hist, edges = np.histogram(tensor, bins=nbins, density=False)\n",
    "    return hist, edges\n",
    "\n",
    "def graph_setup(models, titles, param_name, remove_zeros):\n",
    "    #xs, ys = [LinearScale(), LinearScale()], [LinearScale(), LinearScale()]\n",
    "    xs = [LinearScale() for i in range(len(models))]\n",
    "    ys = [LinearScale() for i in range(len(models))]\n",
    "    xax = [Axis(scale=xs[0])] * len(models)\n",
    "    yax = [Axis(scale=ys[0], orientation='vertical',  grid_lines='solid')] * len(models)\n",
    "\n",
    "    bars = []\n",
    "    funcs = []\n",
    "    for i in range(len(models)):\n",
    "        hist, edges = get_hist2(models[i], nbins, param_name, remove_zeros)\n",
    "        bars.append(Bars(x=edges, y=[hist], scales={'x': xs[i], 'y': ys[i]}, padding=0.2, type='grouped'))\n",
    "        funcs.append(Figure(marks=[bars[i]], axes=[xax[i], yax[i]], animation_duration=1000, title=titles[i]))\n",
    "\n",
    "    shape =  distiller.size2str(next (iter (models[0].state_dict().values())).size())\n",
    "    param_info = widgets.Text(value=shape, description='shape:', disabled=True)\n",
    "    \n",
    "    return bars, funcs, param_info\n",
    "\n",
    "\n",
    "params_names = conv_fc_param_names(sparse_model)\n",
    "weights_dropdown = Dropdown(description='weights', options=params_names)\n",
    "\n",
    "\n",
    "def update_models(stam, bars, funcs, param_shape_desc, models):\n",
    "    param_name = weights_dropdown.value\n",
    "    \n",
    "    for i in range(len(models)):\n",
    "        bars[i].y, bars[i].x = get_hist2(models[i], nbins, param_name, remove_zeros.value)\n",
    "\n",
    "    shape =  distiller.size2str(models[0].state_dict()[param_name].size())\n",
    "    param_shape_desc.value = shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "titles = ['Dense', 'Epoch 0', 'Sparse', '2D-Sparse']\n",
    "models = [models_dict[title] for title in titles]\n",
    "\n",
    "bars, funcs, param_shape_desc = graph_setup(models, titles, param_name=weights_dropdown.value, remove_zeros=False)\n",
    "\n",
    "update1 = partial(update_models, bars=bars, funcs=funcs, param_shape_desc = param_shape_desc, models=models)\n",
    "weights_dropdown.observe(update1)\n",
    "remove_zeros = widgets.Checkbox(value=False, description='Remove zeros')\n",
    "remove_zeros.observe(update1)\n",
    "\n",
    "def draw_graph(models):\n",
    "    if len(models) > 2:\n",
    "        return (VBox([\n",
    "          HBox([weights_dropdown, param_shape_desc, remove_zeros] ),\n",
    "            VBox([\n",
    "                HBox([funcs[i] for i in range(len(models)//2)]),\n",
    "                HBox([funcs[i+2] for i in range(len(models)//2)])\n",
    "            ])\n",
    "         ]))\n",
    "    else:\n",
    "        return (VBox([\n",
    "              HBox([weights_dropdown, param_shape_desc, remove_zeros] ),\n",
    "              HBox([funcs[i] for i in range(len(models))])\n",
    "             ]))\n",
    "    \n",
    "draw_graph(models)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize the weights \n",
    "\n",
    "### Fully-connected layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def view_weights(pname, model_choice):\n",
    "    model = models_dict[model_choice]\n",
    "    weights = model.state_dict()[pname]\n",
    "    \n",
    "#     # Color normalization - we want all parameters to share the same color ranges in the kernel plots,\n",
    "#     # so we need to find the min and max across all weight tensors in the model.\n",
    "#     # As a last step, we also center the colorbar so that 0 is white - this makes it easier to see the sparsity.\n",
    "#     extreme_vals = [list((p.max(), p.min())) for param_name, p in model.state_dict().items()  \n",
    "#                     if (p.dim()>1) and (\"weight\" in param_name)]\n",
    "\n",
    "#     flat = [item for sublist in extreme_vals for item in sublist]\n",
    "#     center = (max(flat) + min(flat)) / 2\n",
    "#     model_max = max(flat) - center\n",
    "#     model_min = min(flat) - center\n",
    "\n",
    "    params_names = fc_param_names(model)\n",
    "    \n",
    "    aspect_ratio = weights.size(0) / weights.size(1)\n",
    "    size = 100\n",
    "    plot_params2d([weights], figsize=(int(size*aspect_ratio),size), binary_mask=True);\n",
    "    \n",
    "model_dropdown = Dropdown(description='Model:', options=models_dict.keys(), value='Sparse')\n",
    "params_names = fc_param_names(sparse_model)\n",
    "params_dropdown = widgets.Dropdown(description='Weights:', options=params_names)\n",
    "interact(view_weights, pname=params_dropdown, model_choice=model_dropdown);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.gridspec as gridspec\n",
    "\n",
    "def color_min_max(model, weights, color_normalization, nrow=-1, ncol=-1):\n",
    "    gmin, gmax = None, None\n",
    "    if color_normalization=='Model':        \n",
    "        # Color normalization - we want all parameters to share the same color ranges in the kernel plots,\n",
    "        # so we need to find the min and max across all weight tensors in the model.\n",
    "        # As a last step, we also center the colorbar so that 0 is white - this makes it easier to see the sparsity.\n",
    "        extreme_vals = [list((p.max(), p.min())) for param_name, p in model.state_dict().items()  \n",
    "                        if (p.dim()>1) and (\"weight\" in param_name)]\n",
    "\n",
    "        flat = [item for sublist in extreme_vals for item in sublist]\n",
    "        center = (max(flat) + min(flat)) / 2\n",
    "        gmax = model_max = max(flat) - center\n",
    "        gmin = model_min = min(flat) - center  \n",
    "    elif color_normalization=='Tensor':\n",
    "        # We want to normalize the grayscale brightness levels for all of the images we display (group),\n",
    "        # otherwise, each image is normalized separately and this causes distortion between the different\n",
    "        # filters images we ddisplay.\n",
    "        # We don't normalize across all of the filters images, because the outliers cause the image of each \n",
    "        # filter to be very muted.  This is because each group of filters we display usually has low variance\n",
    "        # between the element values of that group.\n",
    "        gmin = weights.min()\n",
    "        gmax = weights.max()\n",
    "    elif color_normalization=='Group':\n",
    "        gmin = weights[0:nrow, 0:ncol].min()\n",
    "        gmax = weights[0:nrow, 0:ncol].max()\n",
    "        \n",
    "    if isinstance(gmin, torch.Tensor):\n",
    "        gmin = gmin.item()\n",
    "        gmax = gmax.item()\n",
    "    return gmin, gmax\n",
    "\n",
    "def plot_param_kernels(model, weights, layout, size_ctrl, binary_mask=False, color_normalization='Model', \n",
    "                       interpolation=None, first_kernel=0):\n",
    "    ofms, ifms = weights.size()[0], weights.size()[1]\n",
    "    kw, kh = weights.size()[2], weights.size()[3]\n",
    "    \n",
    "    print(\"min=%.4f\\tmax=%.4f\" % (weights.min(), weights.max()))\n",
    "    shape_str = distiller.size2str(weights.size())\n",
    "    volume = distiller.volume(weights)\n",
    "    print(\"size=%s = %d elements\" % (shape_str, volume))\n",
    "    \n",
    "    # Clone because we are going to change the tensor values\n",
    "    weights = weights.clone()\n",
    "    if binary_mask:\n",
    "        weights[weights!=0] = 1\n",
    "    \n",
    "    kernels = weights.view(ofms * ifms, kh, kw)\n",
    "    nrow, ncol = layout[0], layout[1]\n",
    "    \n",
    "    gmin, gmax = color_min_max(model, weights, color_normalization, nrow, ncol)\n",
    "    print(\"gmin=%.4f\\tgmax=%.4f\" % (gmin, gmax))    \n",
    "\n",
    "    fig = plt.figure(figsize=(size_ctrl*8, size_ctrl*8))\n",
    "\n",
    "    # gridspec inside gridspec\n",
    "    outer_grid = gridspec.GridSpec(4, 4, wspace=0.05, hspace=0.05)\n",
    "\n",
    "    for i in range(4*4):\n",
    "        inner_grid = gridspec.GridSpecFromSubplotSpec(3, 3, subplot_spec=outer_grid[i], wspace=0.0, hspace=0.0)\n",
    "        for j in range(3*3):\n",
    "            ax = plt.Subplot(fig, inner_grid[j])\n",
    "            if binary_mask:\n",
    "                ax.matshow(kernels[first_kernel+i*4*3+j].cpu(), cmap='binary', vmin=0, vmax=1);\n",
    "            else:\n",
    "                # Use siesmic so that colors around the center are lighter.  Red and blue are used\n",
    "                # to represent (and visually separate) negative and positive weights \n",
    "                ax.matshow(kernels[first_kernel+i*4*3+j].cpu(), cmap='seismic', vmin=gmin, vmax=gmax, interpolation=interpolation);\n",
    "            ax.set_xticks([])\n",
    "            ax.set_yticks([])\n",
    "            fig.add_subplot(ax);\n",
    "\n",
    "    all_axes = fig.get_axes()\n",
    "\n",
    "    #show only the outside spines\n",
    "    for ax in all_axes:\n",
    "        for sp in ax.spines.values():\n",
    "            sp.set_visible(False)\n",
    "        if ax.is_first_row():\n",
    "            ax.spines['top'].set_visible(True)\n",
    "        if ax.is_last_row():\n",
    "            ax.spines['bottom'].set_visible(True)\n",
    "        if ax.is_first_col():\n",
    "            ax.spines['left'].set_visible(True)\n",
    "        if ax.is_last_col():\n",
    "            ax.spines['right'].set_visible(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convolutional layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some models have long node names and require longer lines\n",
    "from IPython.core.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))\n",
    "\n",
    "import math \n",
    "params_names = conv_param_names(sparse_model)\n",
    "\n",
    "def view_weights(pname, model_choice, apply_mask, color_normalization, interpolation):\n",
    "    weights = models_dict[model_choice].state_dict()[pname]\n",
    "    \n",
    "    num_kernels = weights.size(0) * weights.size(1)\n",
    "    first_kernel = 0\n",
    "    width = 15\n",
    "    size = int(min((num_kernels-first_kernel)//width, width))\n",
    "    layout=(size,width)\n",
    "    \n",
    "    plot_param_kernels(model=models_dict[model_choice], weights=weights, layout=layout, size_ctrl=2, \n",
    "                       binary_mask=apply_mask, color_normalization=color_normalization,\n",
    "                       interpolation=interpolation, first_kernel=first_kernel);\n",
    "\n",
    "interpolations = [None, 'none', 'nearest', 'bilinear', 'bicubic', 'spline16',\n",
    "                 'spline36', 'hanning', 'hamming', 'hermite', 'kaiser', 'quadric',\n",
    "                  'catrom', 'gaussian', 'bessel', 'mitchell', 'sinc', 'lanczos']\n",
    "\n",
    "#model_radio = widgets.RadioButtons(options=['Sparse', 'Dense'], value='Sparse', description='Model:')\n",
    "model_dropdown = Dropdown(description='Model:', options=models_dict.keys())\n",
    "normalize_radio = widgets.RadioButtons(options=['Group', 'Tensor', 'Model'], value='Model', description='Normalize:')\n",
    "params_dropdown = widgets.Dropdown(description='Weights:', options=params_names)\n",
    "interpolation_dropdown = widgets.Dropdown(description='Interploation:', options=interpolations, value='bilinear') \n",
    "mask_choice = widgets.Checkbox(value=False, description='Binary mask')\n",
    "\n",
    "interact(view_weights, pname=params_dropdown, \n",
    "         model_choice=model_dropdown, apply_mask=mask_choice,\n",
    "         color_normalization=normalize_radio,\n",
    "         interpolation=interpolation_dropdown);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Kernel pruning\n",
    "\n",
    "Look how 2D (kernel) pruning removes kernels.\n",
    "\n",
    "Each row is a flattened view of the kernels that generate one OFM (output feature map)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The default font size is too small, so let's increase it\n",
    "matplotlib.rcParams.update({'font.size': 32})\n",
    "\n",
    "params_names = conv_param_names(sparse_model)\n",
    "\n",
    "def view_weights(pname, unused, binary_mask, model_choice):\n",
    "    model = models_dict[model_choice]\n",
    "    weights = model.state_dict()[pname]\n",
    "    weights = weights.view(weights.size(0), -1)\n",
    "    \n",
    "    gmin, gmax = color_min_max(model, weights, color_normalization=\"Model\")\n",
    "    print(\"gmin=%.4f\\tgmax=%.4f\" % (gmin, gmax))\n",
    "    \n",
    "    plot_params2d([weights], figsize=(50,50), binary_mask=binary_mask, xlabel=\"#channels * k * k\", ylabel=\"OFM\", gmin=gmin, gmax=gmax);\n",
    "    shape = distiller.size2str(model.state_dict()[pname].size())\n",
    "    param_info.value = shape\n",
    "\n",
    "shape = distiller.size2str(next (iter (dense_model.state_dict().values())).size())\n",
    "param_info = widgets.Text(value=shape, description='shape:', disabled=True)\n",
    "\n",
    "mask_choice = widgets.Checkbox(value=True, description='Binary mask')\n",
    "params_dropdown = widgets.Dropdown(description='Weights:', options=params_names, value='features.module.6.weight')\n",
    "model_dropdown = Dropdown(description='Model:', options=models_dict.keys(), value='2D-Sparse')\n",
    "interact(view_weights, pname=params_dropdown, unused=param_info, binary_mask=mask_choice, model_choice=model_dropdown);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's isolate just the 2D kernels\n",
    "\n",
    "Now let's try something slightly different: in the diagram below, we fold each kernel (k * k) into a single pixel.  If the value of all of the elements in the kernel is zero, then the 2D kernel is colored white (100% sparse); otherwise, it is colored black (has at least one non-zero element in it)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "matplotlib.rcParams.update({'font.size': 22})\n",
    "params_names = conv_param_names(sparse_model)\n",
    "\n",
    "def view_weights(pname, unused, model_choice):\n",
    "    model = models_dict[model_choice]\n",
    "    weights = model.state_dict()[pname]\n",
    "    \n",
    "    k_view = weights.view(weights.size(0) * weights.size(1), -1).abs().sum(dim=1)\n",
    "    weights = k_view.view(weights.size(0), weights.size(1))\n",
    "    \n",
    "    #gmin, gmax = color_min_max(model, weights, color_normalization=\"Model\")\n",
    "    \n",
    "    plot_params2d([weights], figsize=(10,10), binary_mask=True, xlabel=\"#channels\", ylabel=\"OFM\");\n",
    "    shape = distiller.size2str(sparse_model.state_dict()[pname].size())\n",
    "    param_info.value = shape\n",
    "\n",
    "shape =  distiller.size2str(next (iter (sparse_model.state_dict().values())).size())\n",
    "param_info = widgets.Text(value=shape, description='shape:', disabled=True)\n",
    "\n",
    "params_dropdown = widgets.Dropdown(description='Weights:', options=params_names, value='features.module.6.weight')\n",
    "model_dropdown = Dropdown(description='Model:', options=models_dict.keys(), value='2D-Sparse')\n",
    "interact(view_weights, pname=params_dropdown, unused=param_info, model_choice=model_dropdown);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}