MNIST_ComputeMI.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os\n",
    "if not os.path.exists('plots/'):\n",
    "    os.mkdir('plots')\n",
    "\n",
    "from six.moves import cPickle\n",
    "from collections import defaultdict, OrderedDict\n",
    "\n",
    "import numpy as np\n",
    "import keras.backend as K\n",
    "\n",
    "import kde\n",
    "import simplebinmi\n",
    "\n",
    "import utils\n",
    "trn, tst = utils.get_mnist()\n",
    "\n",
    "# Which measure to plot\n",
    "infoplane_measure = 'upper'\n",
    "#infoplane_measure = 'bin'\n",
    "\n",
    "DO_SAVE        = False    # Whether to save plots or just show them\n",
    "DO_LOWER       = True    # (infoplane_measure == 'lower')   # Whether to compute lower bounds also\n",
    "DO_BINNED      = True    #(infoplane_measure == 'bin')     # Whether to compute MI estimates based on binning\n",
    "\n",
    "MAX_EPOCHS = 10000      # Max number of epoch for which to compute mutual information measure\n",
    "# MAX_EPOCHS = 1000\n",
    "COLORBAR_MAX_EPOCHS = 10000\n",
    "\n",
    "# Directories from which to load saved layer activity\n",
    "ARCH = '1024-20-20-20'\n",
    "#ARCH = '20-20-20-20-20-20'\n",
    "#ARCH = '32-28-24-20-16-12'\n",
    "#ARCH = '32-28-24-20-16-12-8-8'\n",
    "DIR_TEMPLATE = '%%s_%s'%ARCH\n",
    "\n",
    "# Functions to return upper and lower bounds on entropy of layer activity\n",
    "noise_variance = 1e-1                    # Added Gaussian noise variance\n",
    "Klayer_activity = K.placeholder(ndim=2)  # Keras placeholder \n",
    "entropy_func_upper = K.function([Klayer_activity,], [kde.entropy_estimator_kl(Klayer_activity, noise_variance),])\n",
    "entropy_func_lower = K.function([Klayer_activity,], [kde.entropy_estimator_bd(Klayer_activity, noise_variance),])\n",
    "\n",
    "\n",
    "# nats to bits conversion factor\n",
    "nats2bits = 1.0/np.log(2) \n",
    "\n",
    "\n",
    "# Save indexes of tests data for each of the output classes\n",
    "saved_labelixs = {}\n",
    "for i in range(10):\n",
    "    saved_labelixs[i] = tst.y == i\n",
    "\n",
    "labelprobs = np.mean(tst.Y, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PLOT_LAYERS    = None     # Which layers to plot.  If None, all saved layers are plotted \n",
    "\n",
    "# Data structure used to store results\n",
    "measures = OrderedDict()\n",
    "measures['relu'] = {}\n",
    "measures['tanh'] = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute MI measures\n",
    "-----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for activation in measures.keys():\n",
    "    cur_dir = 'rawdata/' + DIR_TEMPLATE % activation\n",
    "    if not os.path.exists(cur_dir):\n",
    "        print(\"Directory %s not found\" % cur_dir)\n",
    "        continue\n",
    "        \n",
    "    # Load files saved during each epoch, and compute MI measures of the activity in that epoch\n",
    "    print('*** Doing %s ***' % cur_dir)\n",
    "    for epochfile in sorted(os.listdir(cur_dir)):\n",
    "        if not epochfile.startswith('epoch'):\n",
    "            continue\n",
    "            \n",
    "        fname = cur_dir + \"/\" + epochfile\n",
    "        with open(fname, 'rb') as f:\n",
    "            d = cPickle.load(f)\n",
    "\n",
    "        epoch = d['epoch']\n",
    "        if epoch in measures[activation]: # Skip this epoch if its already been processed\n",
    "            continue                      # this is a trick to allow us to rerun this cell multiple times)\n",
    "            \n",
    "        if epoch > MAX_EPOCHS:\n",
    "            continue\n",
    "\n",
    "        print(\"Doing\", fname)\n",
    "        \n",
    "        num_layers = len(d['data']['activity_tst'])\n",
    "\n",
    "        if PLOT_LAYERS is None:\n",
    "            PLOT_LAYERS = []\n",
    "            for lndx in range(num_layers):\n",
    "                #if d['data']['activity_tst'][lndx].shape[1] < 200 and lndx != num_layers - 1:\n",
    "                PLOT_LAYERS.append(lndx)\n",
    "                \n",
    "        cepochdata = defaultdict(list)\n",
    "        for lndx in range(num_layers):\n",
    "            activity = d['data']['activity_tst'][lndx]\n",
    "\n",
    "            # Compute marginal entropies\n",
    "            h_upper = entropy_func_upper([activity,])[0]\n",
    "            if DO_LOWER:\n",
    "                h_lower = entropy_func_lower([activity,])[0]\n",
    "                \n",
    "            # Layer activity given input. This is simply the entropy of the Gaussian noise\n",
    "            hM_given_X = kde.kde_condentropy(activity, noise_variance)\n",
    "\n",
    "            # Compute conditional entropies of layer activity given output\n",
    "            hM_given_Y_upper=0.\n",
    "            for i in range(10):\n",
    "                hcond_upper = entropy_func_upper([activity[saved_labelixs[i],:],])[0]\n",
    "                hM_given_Y_upper += labelprobs[i] * hcond_upper\n",
    "                \n",
    "            if DO_LOWER:\n",
    "                hM_given_Y_lower=0.\n",
    "                for i in range(10):\n",
    "                    hcond_lower = entropy_func_lower([activity[saved_labelixs[i],:],])[0]\n",
    "                    hM_given_Y_lower += labelprobs[i] * hcond_lower\n",
    "                    \n",
    "                    \n",
    "            # # It's also possible to treat the last layer probabilistically. Here is the \n",
    "            # # code to do so. Should only be applied when lndx == num_layers - 1\n",
    "\n",
    "            # ps = activity.mean(axis=0)\n",
    "            # h_lower = h_upper = sum([-p*np.log(p) for p in ps if p != 0])\n",
    "\n",
    "            # x = -activity * np.log(activity)\n",
    "            # x[activity == 0] = 0.\n",
    "            # hM_given_X = np.mean(x.sum(axis=1))\n",
    "\n",
    "            # hM_given_Y=0.\n",
    "            # for i in range(10):\n",
    "            #     ixs = tst.y[::subsample] == i\n",
    "            #     ps = activity[ixs,:].mean(axis=0)\n",
    "            #     hcond = sum([-p*np.log(p) for p in ps if p != 0])\n",
    "            #     prob = np.mean(ixs)\n",
    "            #     hM_given_Y += l * hcond\n",
    "            # hM_given_Y_lower = hM_given_Y_upper = hM_given_Y\n",
    "            # del hM_given_Y\n",
    "                \n",
    "            cepochdata['MI_XM_upper'].append( nats2bits * (h_upper - hM_given_X) )\n",
    "            cepochdata['MI_YM_upper'].append( nats2bits * (h_upper - hM_given_Y_upper) )\n",
    "            cepochdata['H_M_upper'  ].append( nats2bits * h_upper )\n",
    "\n",
    "            pstr = 'upper: MI(X;M)=%0.3f, MI(Y;M)=%0.3f' % (cepochdata['MI_XM_upper'][-1], cepochdata['MI_YM_upper'][-1])\n",
    "            if DO_LOWER:  # Compute lower bounds\n",
    "                cepochdata['MI_XM_lower'].append( nats2bits * (h_lower - hM_given_X) )\n",
    "                cepochdata['MI_YM_lower'].append( nats2bits * (h_lower - hM_given_Y_lower) )\n",
    "                cepochdata['H_M_lower'  ].append( nats2bits * h_lower )\n",
    "                pstr += ' | lower: MI(X;M)=%0.3f, MI(Y;M)=%0.3f' % (cepochdata['MI_XM_lower'][-1], cepochdata['MI_YM_lower'][-1])\n",
    "\n",
    "            if DO_BINNED: # Compute binner estimates\n",
    "                binxm, binym = simplebinmi.bin_calc_information2(saved_labelixs, activity, 0.5)\n",
    "                cepochdata['MI_XM_bin'].append( nats2bits * binxm )\n",
    "                cepochdata['MI_YM_bin'].append( nats2bits * binym )\n",
    "                pstr += ' | bin: MI(X;M)=%0.3f, MI(Y;M)=%0.3f' % (cepochdata['MI_XM_bin'][-1], cepochdata['MI_YM_bin'][-1])\n",
    "            \n",
    "            print('- Layer %d %s' % (lndx, pstr) )\n",
    "\n",
    "        measures[activation][epoch] = cepochdata\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot overall summaries\n",
    "----\n",
    "\n",
    "This is more for diagnostic purposes, not for article\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "import seaborn as sns\n",
    "sns.set_style('darkgrid')\n",
    "\n",
    "\n",
    "#PLOT_LAYERS = [0,1,2,3,4] # [1,2,3]\n",
    "#PLOT_LAYERS = [0,1,2,3]\n",
    "#PLOT_LAYERS = [0,1,2,3]\n",
    "plt.figure(figsize=(8,8))\n",
    "gs = gridspec.GridSpec(4,2)\n",
    "for actndx, (activation, vals) in enumerate(measures.items()):\n",
    "    epochs = sorted(vals.keys())\n",
    "    if not len(epochs):\n",
    "        continue\n",
    "        \n",
    "    plt.subplot(gs[0,actndx])\n",
    "    for lndx, layerid in enumerate(PLOT_LAYERS):\n",
    "        xmvalsU = np.array([vals[epoch]['H_M_upper'][layerid] for epoch in epochs])\n",
    "        if DO_LOWER:\n",
    "            xmvalsL = np.array([vals[epoch]['H_M_lower'][layerid] for epoch in epochs])\n",
    "        plt.plot(epochs, xmvalsU, label='Layer %d'%layerid)\n",
    "        #plt.errorbar(epochs, (xmvalsL + xmvalsU)/2,xmvalsU - xmvalsL, label='Layer %d'%layerid)\n",
    "    plt.xscale('log')\n",
    "    plt.yscale('log')\n",
    "    plt.title(activation)\n",
    "    plt.ylabel('H(M)')\n",
    "    \n",
    "    plt.subplot(gs[1,actndx])\n",
    "    for lndx, layerid in enumerate(PLOT_LAYERS):\n",
    "        #for epoch in epochs:\n",
    "        #    print('her',epoch, measures[activation][epoch]['MI_XM_upper'])\n",
    "        xmvalsU = np.array([vals[epoch]['MI_XM_upper'][layerid] for epoch in epochs])\n",
    "        if DO_LOWER:\n",
    "            xmvalsL = np.array([vals[epoch]['MI_XM_lower'][layerid] for epoch in epochs])\n",
    "        plt.plot(epochs, xmvalsU, label='Layer %d'%layerid)\n",
    "        #plt.errorbar(epochs, (xmvalsL + xmvalsU)/2,xmvalsU - xmvalsL, label='Layer %d'%layerid)\n",
    "    plt.xscale('log')\n",
    "    plt.ylabel('I(X;M)')\n",
    "\n",
    "\n",
    "    plt.subplot(gs[2,actndx])\n",
    "    for lndx, layerid in enumerate(PLOT_LAYERS):\n",
    "        ymvalsU = np.array([vals[epoch]['MI_YM_upper'][layerid] for epoch in epochs])\n",
    "        if DO_LOWER:\n",
    "            ymvalsL = np.array([vals[epoch]['MI_YM_lower'][layerid] for epoch in epochs])\n",
    "        plt.plot(epochs, ymvalsU, label='Layer %d'%layerid)\n",
    "    plt.xscale('log')\n",
    "    plt.ylabel('MI(Y;M)')\n",
    "\n",
    "    if DO_BINNED:\n",
    "        plt.subplot(gs[3,actndx])\n",
    "        for lndx, layerid in enumerate(PLOT_LAYERS):\n",
    "            hbinnedvals = np.array([vals[epoch]['MI_XM_bin'][layerid] for epoch in epochs])\n",
    "            plt.semilogx(epochs, hbinnedvals, label='Layer %d'%layerid)\n",
    "        plt.xlabel('Epoch')\n",
    "        plt.ylabel(\"I(X;M)bin\")\n",
    "    \n",
    "    if actndx == 0:\n",
    "        plt.legend(loc='lower right')\n",
    "        \n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot Infoplane Visualization\n",
    "----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "max_epoch = max( (max(vals.keys()) if len(vals) else 0) for vals in measures.values())\n",
    "sm = plt.cm.ScalarMappable(cmap='gnuplot', norm=plt.Normalize(vmin=0, vmax=COLORBAR_MAX_EPOCHS))\n",
    "sm._A = []\n",
    "\n",
    "fig=plt.figure(figsize=(10,5))\n",
    "for actndx, (activation, vals) in enumerate(measures.items()):\n",
    "    epochs = sorted(vals.keys())\n",
    "    if not len(epochs):\n",
    "        continue\n",
    "    plt.subplot(1,2,actndx+1)    \n",
    "    for epoch in epochs:\n",
    "        c = sm.to_rgba(epoch)\n",
    "        xmvals = np.array(vals[epoch]['MI_XM_'+infoplane_measure])[PLOT_LAYERS]\n",
    "        ymvals = np.array(vals[epoch]['MI_YM_'+infoplane_measure])[PLOT_LAYERS]\n",
    "\n",
    "        plt.plot(xmvals, ymvals, c=c, alpha=0.1, zorder=1)\n",
    "        plt.scatter(xmvals, ymvals, s=20, facecolors=[c for _ in PLOT_LAYERS], edgecolor='none', zorder=2)\n",
    "    \n",
    "    plt.ylim([0, 3.5])\n",
    "    plt.xlim([0, 14])\n",
    "    plt.xlabel('I(X;M)')\n",
    "    plt.ylabel('I(Y;M)')\n",
    "    plt.title(activation)\n",
    "    \n",
    "cbaxes = fig.add_axes([1.0, 0.125, 0.03, 0.8]) \n",
    "plt.colorbar(sm, label='Epoch', cax=cbaxes)\n",
    "plt.tight_layout()\n",
    "\n",
    "if DO_SAVE:\n",
    "    plt.savefig('plots/' + DIR_TEMPLATE % ('infoplane_'+ARCH),bbox_inches='tight')\n",
    "    \n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot SNR curves\n",
    "----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,5))\n",
    "\n",
    "gs = gridspec.GridSpec(len(measures), len(PLOT_LAYERS))\n",
    "saved_data = {}\n",
    "for actndx, activation in enumerate(measures.keys()):\n",
    "    cur_dir = 'rawdata/' + DIR_TEMPLATE % activation\n",
    "    if not os.path.exists(cur_dir):\n",
    "        continue\n",
    "        \n",
    "    epochs = []\n",
    "    means = []\n",
    "    stds = []\n",
    "    wnorms = []\n",
    "    trnloss = []\n",
    "    tstloss = []\n",
    "    for epochfile in sorted(os.listdir(cur_dir)):\n",
    "        if not epochfile.startswith('epoch'):\n",
    "            continue\n",
    "            \n",
    "        with open(cur_dir + \"/\"+epochfile, 'rb') as f:\n",
    "            try:\n",
    "                d = cPickle.load(f)\n",
    "            except:\n",
    "                print('Error loading ', epochfile)\n",
    "                continue\n",
    "            \n",
    "        epoch = d['epoch']\n",
    "        epochs.append(epoch)\n",
    "        wnorms.append(d['data']['weights_norm'])\n",
    "        means.append(d['data']['gradmean'])\n",
    "        stds.append(d['data']['gradstd'])\n",
    "        trnloss.append(d['loss']['trn'])\n",
    "        tstloss.append(d['loss']['tst'])\n",
    "\n",
    "    wnorms, means, stds, trnloss, tstloss = map(np.array, [wnorms, means, stds, trnloss, tstloss])\n",
    "    saved_data[activation] = {'epochs':epochs, 'wnorms':wnorms, 'means': means, 'stds': stds, 'trnloss': trnloss, 'tstloss':tstloss}\n",
    "    \n",
    "    \n",
    "    for lndx,layerid in enumerate(PLOT_LAYERS):\n",
    "        plt.subplot(gs[actndx, lndx])\n",
    "        plt.plot(epochs, means[:,layerid], 'b', label=\"Mean\")\n",
    "        plt.plot(epochs, stds[:,layerid], 'orange', label=\"Std\")\n",
    "        plt.plot(epochs, means[:,layerid]/stds[:,layerid], 'red', label=\"SNR\")\n",
    "        plt.plot(epochs, wnorms[:,layerid], 'g', label=\"||W||\")\n",
    "\n",
    "        plt.title('%s - Layer %d'%(activation, layerid))\n",
    "        plt.xlabel('Epoch')\n",
    "        plt.gca().set_xscale(\"log\", nonposx='clip')\n",
    "        plt.gca().set_yscale(\"log\", nonposy='clip')\n",
    "    \n",
    "\n",
    "plt.legend(loc='lower left', bbox_to_anchor=(1.1, 0.2))\n",
    "plt.tight_layout()\n",
    "\n",
    "if DO_SAVE:\n",
    "    plt.savefig('plots/' + DIR_TEMPLATE % ('snr_'+ARCH), bbox_inches='tight')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GRID_PLOT_LAYERS = [0,1,2,3] # [1,2,3]\n",
    "sns.set_style('whitegrid')\n",
    "max_epoch = max( (max(vals.keys()) if len(vals) else 0) for vals in measures.values())\n",
    "H_X = np.log2(10000)\n",
    "for actndx, (activation, vals) in enumerate(measures.items()):\n",
    "    fig = plt.figure(figsize=(12,11))\n",
    "    gs = gridspec.GridSpec(4, len(GRID_PLOT_LAYERS))\n",
    "    \n",
    "    cur_epochs = np.array(sorted(vals.keys()))\n",
    "    if not len(cur_epochs):\n",
    "        continue\n",
    "    \n",
    "    plt.subplot(gs[0,0])\n",
    "    plt.title('Loss')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.plot(cur_epochs,saved_data[activation]['trnloss'][cur_epochs]/np.log(2), label='Train')\n",
    "    plt.plot(cur_epochs,saved_data[activation]['tstloss'][cur_epochs]/np.log(2), label='Test')\n",
    "    plt.ylabel('Cross entropy loss')\n",
    "    plt.gca().set_xscale(\"log\", nonposx='clip')\n",
    "    \n",
    "    plt.legend(loc='upper right', frameon=True)\n",
    "        \n",
    "    vals_binned = np.array([vals[epoch]['MI_XM_bin'] for epoch in cur_epochs])\n",
    "    vals_lower = np.array([vals[epoch]['MI_XM_lower'] for epoch in cur_epochs])\n",
    "    vals_upper = np.array([vals[epoch]['MI_XM_upper'] for epoch in cur_epochs])\n",
    "    for layerndx, layerid in enumerate(GRID_PLOT_LAYERS):\n",
    "        plt.subplot(gs[1,layerndx])\n",
    "        plt.plot(cur_epochs, cur_epochs*0 + H_X, 'k:', label=r'$H(X)$')\n",
    "        plt.fill_between(cur_epochs, vals_lower[cur_epochs,layerid], vals_upper[:,layerid])\n",
    "        plt.gca().set_xscale(\"log\", nonposx='clip')\n",
    "        plt.ylim([0, 1.1*H_X])\n",
    "        plt.title('Layer %d Mutual Info (KDE)'%(layerid+1))\n",
    "        plt.ylabel(r'$I(X;T)$')\n",
    "        plt.xlabel('Epoch')\n",
    "        if layerndx == len(GRID_PLOT_LAYERS)-1:\n",
    "            plt.legend(loc='lower right', frameon=True)\n",
    "        \n",
    "        plt.subplot(gs[2,layerndx])\n",
    "        plt.plot(cur_epochs, cur_epochs*0 + H_X, 'k:', label=r'$H(X)$')\n",
    "        plt.plot(cur_epochs, vals_binned[cur_epochs,layerid])\n",
    "        plt.gca().set_xscale(\"log\", nonposx='clip')\n",
    "        plt.ylim([0, 1.1*H_X])\n",
    "        plt.ylabel(r'$I(X;T)$')\n",
    "        plt.title('Layer %d Mutual Info (binned)'%(layerid+1))\n",
    "        plt.xlabel('Epoch')\n",
    "        if layerndx == len(GRID_PLOT_LAYERS)-1:\n",
    "            plt.legend(loc='lower right', frameon=True)\n",
    "        \n",
    "        plt.subplot(gs[3,layerndx])\n",
    "        plt.title('Layer %d SNR'%(layerid+1))\n",
    "        plt.plot(cur_epochs, saved_data[activation]['means' ][cur_epochs,layerid], 'b', label=\"Mean\")\n",
    "        plt.plot(cur_epochs, saved_data[activation]['stds'  ][cur_epochs,layerid], 'orange', label=\"Std\")\n",
    "        plt.plot(cur_epochs, saved_data[activation]['means' ][cur_epochs,layerid]/saved_data[activation]['stds'][cur_epochs,layerid], 'red', label=\"SNR\")\n",
    "        plt.plot(cur_epochs, saved_data[activation]['wnorms'][cur_epochs,layerid], 'g', label=\"||W||\")\n",
    "\n",
    "        plt.xlabel('Epoch')\n",
    "        plt.gca().set_xscale(\"log\", nonposx='clip')\n",
    "        plt.gca().set_yscale(\"log\", nonposy='clip')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.legend(loc='lower left', frameon=True)\n",
    "    \n",
    "    if DO_SAVE:\n",
    "        plt.savefig('plots/' + DIR_TEMPLATE % ('gridplot_'+activation) + '.pdf', bbox_inches='tight')\n"
   ]
  },
  {
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}