{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-34-11702066605e>:15: MatplotlibDeprecationWarning: \n",
      "The mpl_toolkits.axes_grid module was deprecated in Matplotlib 2.1 and will be removed two minor releases later. Use mpl_toolkits.axes_grid1 and mpl_toolkits.axisartist, which provide the same functionality instead.\n",
      "  from mpl_toolkits.axes_grid.inset_locator import (inset_axes, InsetPosition,\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import math\n",
    "import argparse\n",
    "from augerino import datasets, models, losses\n",
    "import glob\n",
    "import re\n",
    "import pandas as pd\n",
    "import sys\n",
    "sys.path.append(\"../learning-invariance/\")\n",
    "from data.generate_data import *\n",
    "from loss_getters import compute_loss\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid.inset_locator import (inset_axes, InsetPosition,\n",
    "                                                  mark_inset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "softplus = torch.nn.Softplus()\n",
    "savedir = \"./saved-outputs/\"\n",
    "\n",
    "epochs = 20\n",
    "reg = 0.05\n",
    "## stanard running procedure ##\n",
    "net = models.SimpleConv(c=32, num_classes=4)\n",
    "augerino = models.UniformAug()\n",
    "model = models.AugAveragedModel(net, augerino,ncopies=1)\n",
    "\n",
    "ntrain = 10000\n",
    "ntest = 5000\n",
    "\n",
    "trainloader, testloader = generate_mario_data(ntrain=ntrain, ntest=ntest,\n",
    "                                              batch_size=128, dpath=\"../limit-invariance/data/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_widths = torch.ones(6) * -5.\n",
    "\n",
    "model.aug.set_width(start_widths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(),lr=0.01,\n",
    "                            weight_decay=1e-3)\n",
    "use_cuda = torch.cuda.is_available()\n",
    "if use_cuda:\n",
    "    model = model.cuda()\n",
    "\n",
    "logged_aug = []\n",
    "\n",
    "criterion = losses.unif_aug_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch  0\n",
      "epoch  1\n",
      "epoch  2\n",
      "epoch  3\n",
      "epoch  4\n",
      "epoch  5\n",
      "epoch  6\n",
      "epoch  7\n",
      "epoch  8\n",
      "epoch  9\n",
      "epoch  10\n",
      "epoch  11\n",
      "epoch  12\n",
      "epoch  13\n",
      "epoch  14\n",
      "epoch  15\n",
      "epoch  16\n",
      "epoch  17\n",
      "epoch  18\n",
      "epoch  19\n"
     ]
    }
   ],
   "source": [
    "count = 0\n",
    "for epoch in range(epochs):  # loop over the dataset multiple times\n",
    "    for i, data in enumerate(trainloader):\n",
    "        count += 1\n",
    "        if count < 100:\n",
    "            \n",
    "            # get the inputs; data is a list of [inputs, labels]\n",
    "            inputs, labels = data\n",
    "\n",
    "            if use_cuda:\n",
    "                inputs, labels = inputs.cuda(), labels.cuda()\n",
    "\n",
    "            # zero the parameter gradients\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # forward + backward + optimize\n",
    "            # print(inputs.shape)\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels, model,\n",
    "                            reg=reg)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            log = softplus(model.aug.width).tolist()\n",
    "            log += model.aug.width.grad.data.tolist()\n",
    "            log += [loss.item(), count]\n",
    "            logged_aug.append(log)\n",
    "    print(\"epoch \", epoch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now Compute the Gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = pd.DataFrame(logged_aug)\n",
    "cols = ['width' + str(i) for i in range(6)]\n",
    "cols += ['grad' + str(i) for i in range(6)]\n",
    "cols += ['loss', 'iter']\n",
    "logger.columns = cols\n",
    "logger = logger.reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = next(iter(trainloader))\n",
    "x = x.cuda()\n",
    "y = y.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_widths = model.aug.width.data.clone()\n",
    "widths = torch.linspace(-5, 5, 20)\n",
    "grads = torch.zeros_like(widths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "grad_logger = []\n",
    "for batch_idx, batch in enumerate(trainloader):\n",
    "    if batch_idx < 25:\n",
    "        x, y = batch\n",
    "        x, y = x.cuda(), y.cuda()\n",
    "        for idx, ww in enumerate(widths):\n",
    "            optimizer.zero_grad()\n",
    "            temp_widths = train_widths.clone()\n",
    "            temp_widths[2] = ww\n",
    "            model.aug.set_width(temp_widths.cuda())\n",
    "            preds = model(x)\n",
    "\n",
    "            acc = torch.sum(torch.argmax(preds, -1) == y)\n",
    "\n",
    "        #     loss = loss_func(convnet.model(x), y, convnet.model, reg=0.1)\n",
    "            loss = criterion(model(x), y, model,\n",
    "                                reg=reg)\n",
    "\n",
    "            loss.backward()\n",
    "\n",
    "            grad_logger.append([batch_idx, ww.item(), model.aug.width.grad[2].item(), acc.item()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "log = pd.DataFrame(grad_logger)\n",
    "log.columns = ['idx', 'width', 'grad', 'acc']\n",
    "\n",
    "log['sp_width'] = softplus(torch.tensor(log['width']))\n",
    "log['acc'] = log['acc']/128."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27560 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sns.set_style(\"white\")\n",
    "lwd=3.\n",
    "fs=18\n",
    "sns.lineplot(x='sp_width', y='acc', data=log,\n",
    "            linewidth=lwd)\n",
    "plt.xlabel(\"Rotation Distribution Width\", fontsize=fs)\n",
    "plt.ylabel(\"Batch Accuracy\", fontsize=fs)\n",
    "sns.despine()\n",
    "plt.tick_params(\"both\", labelsize=fs-4)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 69876 characters]\n",
      "text/plain": [
       "<Figure size 1000x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(figsize=(10, 5), dpi=100)\n",
    "\n",
    "sns.set_style(\"white\")\n",
    "sns.lineplot(x='sp_width', y='grad', data=log,\n",
    "            linewidth=lwd, ax=ax)\n",
    "ax.set_xlabel(\"Rotation Distribution Width\", fontsize=fs)\n",
    "ax.set_ylabel(\"Gradient\", fontsize=fs)\n",
    "ax.tick_params(\"both\", labelsize=fs-4)\n",
    "sns.despine()\n",
    "\n",
    "ax2 = plt.axes([0.2,0.2, 0.01,0.01])\n",
    "ip = InsetPosition(ax, [0.08,0.45,0.5,0.5])\n",
    "ax2.set_axes_locator(ip)\n",
    "mark_inset(ax, ax2, loc1=2, loc2=4, fc=\"none\", ec='0.5')\n",
    "sns.lineplot(x='sp_width', y='grad', data=log[(log['sp_width']< 1.2) & (log['sp_width'] >0.05)])\n",
    "sns.despine()\n",
    "# ax2.axhline(0)\n",
    "ax2.set_xticks([])\n",
    "ax2.set_yticks([])\n",
    "ax2.set_xlabel(\"\")\n",
    "ax2.set_ylabel(\"\")\n",
    "ax.set_xlim(0, 4.4)\n",
    "plt.savefig(\"./grad_example.png\", bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
