ELI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P7aXbQlylnyx",
    "tags": []
   },
   "source": [
    "### Energy-based Latent Aligner for Incremental Learning [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/JosephKJ/ELI/blob/main/ELI.ipynb)\n",
    "\n",
    "#### CVPR 2022\n",
    "\n",
    "This notebook contains code to replicate the MNIST experiment described in Section 3.3. The size of the latent dimention is **32**. The key logic of our proposed methodology is encapsulated inside `EBMAligner` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "sWSZmaW-ef-q"
   },
   "outputs": [],
   "source": [
    "# some initial imports\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from sklearn import manifold\n",
    "from time import time\n",
    "from matplotlib.ticker import NullFormatter\n",
    "import copy\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision.datasets as datasets\n",
    "import torchvision.transforms as transforms\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import TensorDataset, DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "bWXs3Fy9b-Pu"
   },
   "outputs": [],
   "source": [
    "# Seeding for reproducibility\n",
    "seed = 120\n",
    "np.random.seed(seed=seed)\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "# Setting up Device.\n",
    "use_cuda = True\n",
    "use_cuda = use_cuda and torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kyob1mpgui_j",
    "outputId": "a6ffb822-6213-4cf9-8267-487bccab0756"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fatal: destination path 'continualai/colab' already exists and is not an empty directory.\n",
      "Files already downloaded!\n",
      "Size of training data is: 60000, size of test data: 10000\n"
     ]
    }
   ],
   "source": [
    "# Getting MNIST data\n",
    "!git clone https://github.com/ContinualAI/colab.git continualai/colab\n",
    "from continualai.colab.scripts import mnist\n",
    "mnist.init()\n",
    "\n",
    "x_train, t_train, x_test, t_test = mnist.load()\n",
    "\n",
    "print(f'Size of training data is: {t_train.shape[0]}, size of test data: {t_test.shape[0]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 147
    },
    "id": "3PYIvkASeFfn",
    "outputId": "c879a380-97bf-4a96-bd38-53e1f6090e8b"
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAV4AAAB/CAYAAACnrKo9AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/Il7ecAAAACXBIWXMAAAsTAAALEwEAmpwYAAAkcElEQVR4nO3de1xVdbo/8M93X7iDXOUilz1CSkrpESrHmGy00fzJ0RgDdULTHDVnrDNmTXI005yy8mRjpU4F3jiY02hN0ilrsEY9piWJgmiCCEjcBJWLym1vPr8/FI+ZXMS9F9g879fr+/LlYu/1PDxr7Wd/12VvFEkIIYTQjq67ExBCiH810niFEEJj0niFEEJj0niFEEJj0niFEEJj0niFEEJj7TZepZSXUurQ5VGulCq56v92Ha1cKTVNKfVWB49RSqk3lFInlFJZSqkhN/pLCCHErcTQ3g9JngEwGACUUksAnCf5X1bOYQyA2y6PewCsvfyvEEL8JN3wqQal1Eyl1AGl1GGl1DallNPl5XFKqSOXl+++zvPGKqX2KaW8r/nReACbeMl+AO5KKf8u/TZCCHEL6Mo53g9I3kVyEIBjAGZcXr4YwOjLy8dd/QSlVCyABQD+H8mqa9bXB0DxVf///vIyIYT4SWr3VEMbIpRSfwLgDsAFwGeXl+8FsEEp9T6AD656/AgAUQBGkay9iVyFEOInoSsz3g0A5pK8A8BSAA4AQPJxAIsABAH4Vinldfnx+QBcAfRrY30ll5/TKvDyMiGE+EnqSuN1BVCmlDICeKR1oVIqlOTXJBcDqMT/NdMiABMAbFJKDbzO+rYDmHr57oahAGpIlnUhLyGEuCV0pfE+B+BrXDq18N1Vy1copbKVUkcAfAXgcOsPSH6HS036b0qp0GvW9wmAkwBOAHgXwO+6kJMQQtwylHwtpBBCaKsrF9esSilltc5PUnVnfABVJH26MwepQddq0BNy6O74PSGH7o6vVQ7d3ng7y93dHQMGDMCFCxdw/PhxNDQ0dHdK11PU3Qn0ADapQXh4OPz9/ZGZmYnq6mpbhBA3wWAwYPDgwcjLy0NNTY1NYtjZ2SE4OBheXl44c+YMzp49i7Nnz9oklq3dEo3X3d0dL774Ih5++GFUV1fjT3/6E1JSUrotH09PT0ycOBGbN2+22U7WnlGjRuHhhx9Gamoqdu3apUlMpRQcHR0RFBSEgIAA7N69GxaLRZPYTk5OWLVqFaKjozF58mRs375dk7hXc3Nzg7e3N2JiYhAREXFleWVlJf7yl7+guLi4nWdbh6OjI0aPHo29e/eisrLS5vFuxNixY5GYmIhZs2YhKyvL6ut3dXXF7NmzMWnSJPTp0welpaVIS0vDK6+8gvr6eqvHAwBfX1+EhYUhMDAQ/fv3R21tLc6dOwcAaGxsREZGBk6ePImWlpYbXznJbh0A2NZQSrFPnz5MSkrixYsXWVtby+zsbH7++eccNGjQjx5/M/G9vb0ZFxdHe3v7NvNpHXfffTc3bdpEf3//a3+WYe0aXDsiIiJ4+PBhms1mzp4926o1uN4wGo0cMGAAFy1axP3797O4uJglJSVMSEigTqe73nOsXgOj0cjU1FRaLBYmJCR0WCNr74sTJkxgeno6T548yfr6elosliujqamJH374Id3d3W26HQDwnnvuYVVVFe+9917Na9DeiI6O5rFjx5iamkovLy+b1MBgMPCPf/wjL168SLPZzJaWFtbV1V33NWCNGhgMBm7evJk1NTWsq6vjuXPnWF5eztLSUhYUFLCsrIxlZWX8+9//zqFDh95wDj16xhsZGYnExESMGTMG9vb2MJvNSEtLw2OPPYZ58+Zh+vTprYW6aSEhIfj1r3+NHTt2oLGxsc3HKaUwaNAguLq6oq6uziqxO8toNGLEiBHo06cP0tPTkZGRYdN4fn5+mDFjBuLj4zFgwAAUFBRg5cqVeOKJJ/DAAw/g73//O86fP2/THFp1aVZhJffffz+GDx+Oqqoq7Nq1CyTRq1cvDBo0CPb29oiOjoa3t7fNT4EYDAZ4eXnB09PTpnHaM3DgQPTq1Qv79u0DSRiNRtx///3Q6XRISkrCmTNnbBLXbDZj06ZNmDhxIkJCQuDq6goXFxcYDLZpYSRx6NAhFBcX4/z58zh16hRqay99/uv8+fPw8vLC888/j3HjxiEnJwf79++/8QDdOdDGO4W/vz9TU1PZ0NDAr776iitXrmRmZibHjh3Lt99+m59//jkdHBysNssYPXo0P//8c7q6urb77u7l5cWDBw/y1Vdf5eWT8JrNeGNjY1leXs6CggJGRkZeL75VZxkLFy5kfX096+rqmJmZyejoaPr6+jIrK4tbtmyhm5ubJjNeJycnfvrpp7RYLFy9ejV9fX01ne15e3tz0KBBDAsLo5ubG93c3BgcHMy9e/fSbDYzLS3NqvtiW+Pee+8lSU6ePLnbZrxpaWlcu3YtHR0dCYDBwcEsKiri888/T4PBYPMaxMXFMTMzkxaLhSQ5Y8YMm9XAYDDQaDRSr9f/4LVmZ2fHMWPGsLi4mGVlZRw3btwN59Ajv4/XxcUFc+bMwbhx45CXl4eFCxfijTfewIQJE/Dpp58iJycHkZGRiI6Otko8nU6Hn//853B0dOzwsa6urggICEBmZmbrRtKE0WjEY489Bnd3d7z33nv49ttvbR7/7Nmz+OCDD/Dcc89hypQp+Oqrr+Dl5dWpOlnTxYsXsW7dOgDAhAkT8LOf/UzT+FVVVTh8+DBOnDiB2tpaNDQ04Be/+AWCg4MBAOnp6Zpc7LVYLGhubsbtt99u81jX069fPwQFBaGmpgYtLS3Q6XQYO3YsXFxccOTIEZjNZpvnkJeXh8OHD1/Z9wcPHmyzWGazGc3NzbBYLFfieXl5Ydq0aUhKSoKrqyuWLl2KTz/99IbX3eNONej1ekyZMgXTp0/HO++8g/Xr16O4uPhHF7GcnJwwaNAgpKen33RMnU6HO+64AyUlJW3uPHq9Ht7e3viP//gPNDc3o6hIuxsY7O3tMXfuXPziF7/AP//5T7z99tuaxE1OTkZKSgqampquXFwaOHAgXF1dcfHiRU0P/48fPw6lunR30E1zcHCAj48PjEYjbrvtNjz55JO4++674ezsjMOHD2t2gbOyshLFxcVwcXHRJN7VPD09MXfuXDQ0NCApKQlmsxnjx4/HvHnz8NZbb+Hjjz+2aXxHR0f84Q9/wJw5c+Dj4wO9Xg8AOHjwoE3jXs3V1RXLli1DQkICXFxcUF5ejtLSUnh5eaGqquqG3nh6XOMNCgrC1KlTkZOTg6VLl145r3I1pRSUUleKbw0GgwHFxcVXrtT7+PjAx8cHQUFBCAwMhMlkQr9+/TB8+HDs2bMHOTk5VovdkcjISEybNg3FxcV44YUXcOrUKU3iNjc3w8nJCXFxcZg0aRIMBgM8PT1RX1+PlJQUzc7vApfO8Wp5hNHK1dUVTz/9NMaNGwd3d3f06tULvXr1AkkcOHAACQkJKCws1CQXi8UCs9kMo9GoSbxWQUFBWLp0KSZNmoRt27bBbDYjKCgIiYmJqKysRFJSUrvXRayBJJqamuDu7g4HB4cry++//36sX7/eprFbOTg4wMXF5coEwM/PD//93/+NzMxMbN++HampqSgvL+/Uunpc4x01ahQiIiLw9NNPX7fp6nQ6ODk54fz588jPz7dKTJIoKyvDqFGjEBAQAIvFgp/97GdoaWlBXV0d8vPzkZWVhbS0NISEhKC4uPi6udmCr68vVq1ahdtvvx2rV6/G/v37NWlAjo6OmDBhAubMmQN7e3tUVFRgwIABCAgIwPfffw9nZ2cYDAZNDi+vppSCnV2Hf/zEahwdHTFs2DD07t0bTU1NqK2tRV1dHby9veHr6wsAmt1W19LSgqamJkRFRWkSDwBMJhNWrFiB0aNHIz8/HyNGjED//v1hsVjQv39/PP744ygtLbV5Hg0NDVi7di1OnToFb29veHt74+mnn4abmxv0er0m2+DMmTOYP38+3nnnHQQEBGDEiBGIj49HdHQ07rrrLgwbNgyPPvpo5yYkXT0Bba2Bq05EK6W4YMECnjlzhiaT6UcnqpVSHDlyJLOzs/nmm2/SaDRa7WS+p6cnY2NjOXXqVE6dOpVxcXEMDw//wQWD4cOHMz8/n7NmzWrrZLpVLywZDAY+8cQTrK2tZUVFBSMjI21yUeXadXh6enLBggWsrKzka6+9RpPJxGHDhjE7O5uFhYU8duwYjx49ytmzZzMwMJA+Pj42vbgGgHfeeSdbWlpYX1/PBQsW/OhCzs3WoL0c/P39GRERQT8/P/r6+tJkMjElJYXl5eUcMGCAzbbDtcPDw4MfffQRs7OzO7zt0VqvxzVr1vDAgQN84oknGBQUxMWLF/PcuXO0WCw0m838xz/+QWdnZ81q0DqCgoJ44sQJZmRk0MPDQ5P94Nphb2/PmJgYZmRkXLm9LTY2tlM59LgZLwAUFhb+aEbp4eGBX//615g2bRo2btyIjRs3orm52Woxz549iw8//LDdx7i7u6O5uRn/+7//a7W47enXrx+eeuoptLS0YOXKlThy5IjNYzo5OSExMRETJ07EypUrsW3bNtxxxx34z//8T1RVVWHu3LmoqanBQw89hOnTp2PmzJk4fvw4Hn30UZvOfqurq3HkyBEMGDAAo0ePxsaNG1FWZtsvsdPpdNDpdCgrK/tBLG9vb6ue5uqs5uZmVFdXw8HBAV5eXprMNPft24dVq1YhLy8PLS0tyM3NhVIKL774Ig4ePIgTJ07gwoULNotvZ2cHBwcH1NXVtTZFKKWQkJCA0NBQvPbaa5ocfSql4OzsjIaGhiv7eWNjIz7++GNER0djyJAhINn5D3N09V3BWgPXmfHW1tZyzJgxdHd3Z3h4OCdNmsS//vWvLCws5MKFC20622tvjB8/nt99912bsxxYebY3c+ZMWiwW5ubmMiQkpFM53mz8GTNmsLa2ljt27OAf//hHpqWlsby8nC+99BL9/Px+EGvgwIGcOXPmtfWwyYzX2dmZa9asYVNTE4uKijhkyBCbzvY8PDz46KOPcuzYsVduJVJK0c/Pj2vXrmVjYyNLSkoYHh6u6b6YkJDAo0eP8s4777T5jPfaYTKZuGPHDu7cubPDW/qsUQOlFOfOnctt27axX79+tLe3Z+/evfnggw+ysLCQ5eXlHDZsmM1roNPp+Mtf/pJvvPHGDz40pZRi3759uXPnTra0tPDEiRPs3bv3rTnjbW5uhoODA9566y0cPXoUISEh8Pf3x9GjR7F48WKkpaV1a356vR5ubm42j9O7d2/MmjULFosFH3/8sc1nd618fHzQ1NSEIUOGICQkBPv27cPvf/977Ny580cfEMjJydHsIuOFCxewZ88e/OY3v0FgYCACAgJsekV70qRJeOGFF/DNN9+gX79+IAlHR0c88MADiIqKQmFhIZKTkzW9uwW4dGtbRUUFjh07pmlcBwcHTJ06FaGhoXjyySdx+vRpTeKazWb827/9G9577z0cOXIEJpMJt912G3r16oUVK1YgMzPT5jm0Xlz08PDAmjVr0NDQAL1ej7i4OMyaNevKd8i88847nf8od1ffFaw1cM27S0REBLdu3Xrlo5kZGRlMTExs8zwOrPgO29EYP348i4qKGB8fb/MZb1hYGMvKylhVVdXp2a41amBnZ8fo6GiOHz+e3t7e1Ov1nY5t7RpcOwIDA6/cPB8TE2PTmc7kyZOZm5tLs9nM5uZmNjc38+zZszx58iSTkpIYFRXV1kembbovTp48mdnZ2ezbt6/V94P2cujXrx9zc3PbPeK0RQ18fHw4duxYfv7558zKymJTUxPNZjPfeeed636AyBY18PT0ZHJyMpuamlhVVcVjx44xNzeXDQ0NbGlpYXZ2NufNm0cnJ6dO59DjZrw5OTn4/e9/j9WrV8NoNKKwsBAFBQVWPZ/bVSUlJSgvL9ckl+rqauTk5MBkMml21RwAmpqaNDuHfaPKy8vx29/+Fl5eXjaf6Xz44YeorKzEggULEBYWhi+//BLbt29HQUEBCgoKuuXLkQCAJPLz8zX/Vq6ioiL85je/QUmJtn+Vq7KyEp988gmysrLg7OyMPn36QKfT4dChQ61N0uaqq6uxbNkyfP/99wgODobJZEJTUxMyMjJw8OBBpKen47vvvruhD9F0+xehd/f3b95IfKUUDAYDLBZLWx8e+JbkDd/r01YOBoMBSqkbavS2rkEnWLUGXdGVGrSVQ+s2aGlpuaE3QFtth9YLfh1dyLRmDbqqu/dFa9dAp9Nd+QzB5fV3eH95Wzn0uBlvT0ZS05m31vfIih/radugpaWlW78w6F+ZNeveExpvFazz5dkh3Ry/J+TQ3fF7Qg5djd8Tcuju+D0hh+6Or0kO3X6qQQgh/tX0yG8nE0KInzJpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIoTFpvEIIobF2G69SykspdejyKFdKlVz1f7uOVq6UmqaUequDx4QrpfYppRqVUk/f6C8ghBC3GkN7PyR5BsBgAFBKLQFwnuR/WTmHswCeBPCQldcrhBA90g2falBKzVRKHVBKHVZKbVNKOV1eHqeUOnJ5+e7rPG/s5Zmt99XLSZ4meQBAc5d/CyGEuIV05RzvByTvIjkIwDEAMy4vXwxg9OXl465+glIqFsACAP+PZNXNJCyEELe6dk81tCFCKfUnAO4AXAB8dnn5XgAblFLvA/jgqsePABAFYBTJ2pvIVQghfhK6MuPdAGAuyTsALAXgAAAkHwewCEAQgG+VUl6XH58PwBVAv5vOVgghfgK60nhdAZQppYwAHmldqJQKJfk1ycUAKnGpAQNAEYAJADYppQbebMJCCHGr60rjfQ7A17h0auG7q5avUEplK6WOAPgKwOHWH5D8Dpea9N+UUqFXr0wp5aeU+h7AUwAWKaW+V0q5dSEvIYS4JSiS3Z2DEEL8S+nKxTWrUkpZrfOTVN0ZH0AVSZ/uzMGWNVBKwd3dHQ4ODigrK2vrYbdkDXpCDjcaXymFtiZO/yo1sHZ8rXKQjwxbV1F3J2ALSimEhIQgOTkZ3333HebPn9/ew3+SNehJ/Pz88MILL+Cjjz6Cp6dnd6cjuqDbZ7ztUUrBYDDA0dERjo6OGDZsGJydnZGXl4f4+HhUVlYiOTkZlZWV3Z3qT1poaChWrFiB++67D5s3b8Ybb7zR3Sn1CEop/PznP8fIkSPx+uuv4/z58zaP6efnhyVLlqCurg4vv/wyampqbB6zp9LpdFBKoaWlpc2Zf0/V4xqvXq/HoEGDcMcddyAqKgr9+/fH4MGDoZSCg4MDMjMz8f333yMoKAhpaWlobGy0aT7u7u4IDAyEu7s7amtrcfr0aeh0OlRUVMBisdg0dnv0ej0cHBwQEhICNzc3FBUVtXf432VGoxExMTEYO3YsDhw4gLfeegvFxcVWj9NZfn5+cHZ2hq+vL5RSOHz4sCYNz87ODsHBwXBycoK/vz98fX1RWlqKWbNmwWQyYeXKlTbPwc3NDQsXLsT58+exaNEim+/7benVqxduu+02ODs7o6ioCIWFhZrGd3Fxwa9+9StMnDgRffr0QUpKCpKSktDS0qJpHjeFZLcOALx62Nvb8/XXX2dTUxMtFgstFgvPnDnDvLw81tfXMzY2lqGhoQwICKCTk9MPnmuN+FePPn36cO3atczPz2d+fj6zs7P5j3/8g3/9618ZGBh4vedkWDuH642wsDDOmzePKSkpPH78OGtqarhw4UKb1GDw4ME8evQo33vvPUZERHQmP6vXQCnFsLAwPvTQQ/zoo4+YkZHBCxcu8MKFC1yxYgUvn5O7qf2goxwiIiK4Z88eFhYWsri4mLm5uczOzmZlZSXXrl1LOzs7m+6LADh9+nRu2rSJPj4+HW4HW9SgtQ5JSUk8ffo0zWYz9+zZw8jISKvl0FH8Xr16cfHixdyyZQtHjRrFqVOnMikpifb29prV4Or90tnZmb1796abm9uP9sP2cuhxM97GxkYkJyfj9OnTGDJkCEaOHInFixfjww8/xC9/+Ut88803KCkpsXkeoaGheOGFFzB06FCsWbMGH330ESZOnIhnn30WX331Fcxms81zuJpOp0NERASef/553H777WhoaMD58+fx7rvvIj09HQUFBVaPGRUVhU2bNqGwsBDz589HaWmp1WO0Ra/XIzw8HC+++CJCQkLg4eGBxsZG5Obmwmg0orm5GS4uLhgxYkS7F5mswd/fH8uXL0dISAg2bNiA7du3w2Aw4N1330VLSwtSU1PR1NRks/gA4OHhgcmTJ2Pt2rXddmotODgYf/7zn+Ht7Y23334blZWVeOqppxATE4NDhw7Z/AhQp9Phd7/7HUJDQzF//nxUVFRg/vz5KCsrs+rsX6/Xw93dHWfOnPnBcnt7e3h4eKBv376YOHEi/P394eHhgfDwcOTl5eGxxx7r/Oy/q+8K1hpo493EYDBw3LhxzMvL44MPPtipdyBrxQ8PD+eOHTt48uRJPvLII9Tr9ezduzcPHDjA6upqTpw4kTqdTpPZXusYOHAgP/vsM65bt44JCQkMCgqiq6urzWpgNBq5bNkyWiwWLly48MqMLjAwkA8++CAHDBhg0xmvyWTirl272NzczKamJqalpTE2NpbOzs6Mjo7m0aNHeeHCBSYmJtp8xuvt7c0tW7YwJyeHkZGRDAgI4MaNG1laWspnn32Wzs7ONp/tTZw4kX/729+uG8tar4WOcpg9ezYLCwv5q1/9igDo4eHB3bt3880336SDg4PNazBy5Eh+9tlnV468hg4dyszMTM6aNcvqNbje7HXhwoXcuXMni4uLWVFRwWeeeYYzZsxgeXk5S0pKOGjQoE7n0KXErDnaK7SbmxtTUlKYmprKsLCwH51asMWGNplM/OKLL5ibm8v4+Hja2dnRz8+PTzzxBKurq7l8+fIfHVbasvG6u7tzyZIlXLJkCf/93/+dbm5uNq8BcOmQbufOnWxubmZMTAwNBgOnTp3K1NRUlpaW8tSpU5w8eTL1er3Va6CU4rx583jx4kWePn2ay5Yto4+PD/V6PQ0GA6dNm8YzZ85w7969DA0NtXnT0el0TEhIYGVlJTdt2sQvvviC1dXVXLx4MR0dHW26HVpfB+vXr+eoUaPo4ODAiIgIxsbGtnW6y2aNd86cOczLy+PYsWM5btw4rlq1iuvWrWNYWJjNa+Dt7c1t27ZxwoQJVErRw8OD77//PjMzM+nt7a1JDbZu3cr09HQmJCQwPDycffv25datW3n27Fk+99xz190X2ozR1eSsNdrb0MCl2efevXt56NAhLlq0qN3zW9aIP2bMGNbX13PGjBkMCwvjzJkzuWPHDjY0NPDLL79kUFBQe/lavfHOmTOHJ0+e5H333ddunaxZA6UUhw4dyoqKCh48eJDh4eEcPnw49+/fz1OnTjE5OZnp6ek8cODA9XZ6qzTeN954g4WFhXz88cfZq1evKz978MEHmZmZyRMnTjAmJsZqM62OtsPIkSNZVVVFi8XCxsZGrl+/vt19wZrxBw4cyK1bt9LX15exsbFMSkri6tWruX79+jZfD7aowbPPPsuamhoWFxezvr6eGzZsuO5Rly1qMGjQIP7zn/+kv78/9Xo9//CHP7CgoICPPPKIVeO3l0NgYOCV/S0oKIhJSUlXrjPc6BtwlxKz5uiokeh0Og4cOJApKSksLS3lmjVrGBAQYLMNfd9997G8vJzHjx9nVlYWs7KymJOTw7q6Oi5atKitUwxWb7wGg4FKKc6YMYMFBQV8/fXX6enpqUnj9fDw4IYNG1hQUMBRo0Zx+PDh/Prrr5mVlcWpU6cyNjaWhw8f5ptvvkmDwWCTGnh6etJkMv3g6CIiIoIHDx5kTU0NExMT29wW1t4XjUYj58yZw7q6Ora0tDAtLY1eXl7XPRy1RdOJjY3lunXr6OHhwSlTpjA0NJQODg587rnnbHKY3dbw8vLi+PHj+T//8z88d+4c4+PjNavB5MmTuWnTJg4ePJivvvoqS0tL+fzzz7fZ8GxVAwD09fXlunXrePbsWb755pvs3bv3DefQpcSsOTr6JVuHi4sL//znP7OxsZHLly+32lXMa9dhb2/PCRMmcOnSpUxISKDJZOKyZctYUlLCkSNHdpSn1c5vJiYm8s4776TRaOSMGTPY1NTEmJgYTRpvSEgI9+3bx23bttHFxYW7du1iXl4ehw4dyjFjxjA/P5+vvPJKW28ENjnP7ezszPXr19NisXDLli0/mAXb8gVnb2/PuLg4VlRUsKSkhGVlZdy6dasm26F1/Pa3v+WLL75Io9FIBweHK81u8ODBXLlypaZNx9nZmRs3buTJkyfbPMS3RQ2ioqL45Zdfcvfu3fz666954MABmkwmq8fvqAZeXl5MTk5mXV0dly9f3uXTnz3uroZrtd4/GRMTg6FDh4Ik6uvrbXZXQWNjI7Zt24Zt27YBAJydndGnTx80Njbi8OHDHTz75ul0OqxatQrOzs744osv0Lt3b5hMJhQUFODUqVM2j9+KJCorK+Ht7Q13d3fU1dVh4cKFcHV1xeLFi/HJJ5/g3LlzmuVjMplw7733orS0FOvXr0dtre2/2lkphcjISCxduhQZGRn4y1/+gieffFLT37tVQ0MDzGYzmpv/7w+1BAcHa3IP89UcHR3h4uKCI0eO/CAXWzt06BAeffRRBAcHY8mSJVi/fr3m9w97enpiyZIliI+Px+bNm/Hqq6/i4sWLXVtZV98VrDXQxjuFTqdjSEgIExMTefz4cTY3N7Ouro7bt2/nsGHDbP4O2zqMRiNff/115uXldeaKslXOb06ZMoUZGRmsqalhRUUFDx06xPj4eBqNRk1mWq6urkxNTWVxcTGTk5NZVVVFs9l85X7h651XtWYNrq1HREQE33//febn53P69OnXO71hk5mOr68v09LSuHPnTvr5+TE+Pp7V1dVctGiRJtuhdcTFxfG11177wVFe3759mZyczKioKM1me3q9nomJidyzZw/vv/9+TWvQOl5++eUOzy3bogZGo5FLlixhZWUl16xZQ19f3w7jt5dDlxKz5rg2UYPBwICAAMbFxXHv3r2sq6tjXV0dv/zySz788MPt/sK22NDu7u7ctm2bZo0XuHR4GxoayrvvvptDhw5l//79O9V0rVUDvV7PBx54gOnp6SwrK2NeXh7Xrl3L6dOn08vLS5MatA5nZ2euWrWK1dXVfOaZZzpsutZ8wZlMJn7zzTfcsGEDo6KiuH//fu7atYt33nmnpk0nJCSEmzdv5rRp0zhw4EBOnjyZW7Zs4cMPP6zZeW6lFEePHs3jx49z1qxZNtsO7a0vKCio003f2jUIDg5mUVERU1JSOt1028uhS4lZc7QmaGdnx8jISC5atIh5eXlsampic3Mz9+zZw7i4OJu9w3W0TqPRyOXLlzMvL68zF7c0+eSaVjXQ6XRXRnsXUWxZg1mzZrGoqIjvvvtuh+fTrP2C69u3Lw8cOMDs7GyWlZXxk08+afPCrq33xdDQUC5fvpyrV6/mSy+9xIiICKtf2Govh7CwMO7evZspKSn08/PTvAZOTk587bXX+PLLL1/vFkab18Dd3Z0vvfRSm5/Su9EcupSYNUdrgsHBwdy6dStLS0uZn5/PlJQUzp8/n2FhYZ1+0Vt7Z28dw4cPZ05ODmfOnKnpYXZXRnfHt2YNfHx8mJaWxkOHDjE6OtqmNbheDk5OTpw9ezaPHj3KzZs3d9jsetJ2sFYNgEtHYK+88gqzsrIYHh7eLfti//79uX//fvbv379bagBcmhzewASk3Ry6lJg1x9VJGo1G2tnZ0Wg0Uq/XW+2XvNmd3Wg08plnnuG3337Le+65R5Omo+ULzprxrVmD4cOHMzc3l1OnTu3Uoa0tXnBKqSv74620HaxVA4PBwClTpvCzzz7jAw88cEOvSWvWICQkhHPmzGn39jFb7gfW3g7d/hcouvuLjzsb39HREXfddRfy8vLa+xawb0lG2SqHzrBlDTrJajVwcXHBkCFD8O233+LChQudXldXatBWDl3V3dvBWjUwGAyIiIjA+fPnUVhYeEN3E/1UamCLHHpC462Edb48O4Rd+8sH1orfE3Lo7vg9IYcuxe8JOXR3/J6QQ3fH1yqHbm+8Qgjxr0b+9I8QQmhMGq8QQmhMGq8QQmhMGq8QQmhMGq8QQmhMGq8QQmhMGq8QQmhMGq8QQmhMGq8QQmjs/wM1rTOLgFTPYQAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 360x216 with 20 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Creating two tasks\n",
    "task_classes_arr = [(0, 1, 2, 3, 4), (5, 6, 7, 8, 9)]\n",
    "tasks_num = len(task_classes_arr) \n",
    "\n",
    "task_data = []\n",
    "for i, task_classes in enumerate(task_classes_arr):\n",
    "  train_mask = np.isin(t_train, task_classes)\n",
    "  test_mask = np.isin(t_test, task_classes)\n",
    "  x_train_task, t_train_task = x_train[train_mask], t_train[train_mask] - task_classes[0]\n",
    "  x_test_task, t_test_task = x_test[test_mask], t_test[test_mask] - task_classes[0]\n",
    "  task_data.append((x_train_task, t_train_task, x_test_task, t_test_task))\n",
    "\n",
    "# Display tasks\n",
    "def plot_task(axs, data, samples_num):\n",
    "  for sample in range(samples_num):\n",
    "    axs[sample].imshow(data[sample][0], cmap=\"gray\")\n",
    "    axs[sample].xaxis.set_ticks([])\n",
    "    axs[sample].yaxis.set_ticks([])\n",
    "\n",
    "n_tasks, samples = 2, 10\n",
    "_, axs = plt.subplots(n_tasks, samples, figsize=(5, 3))\n",
    "for task in range(n_tasks):\n",
    "  axs[task, 0].set_ylabel(f'Task {task}', rotation=0)\n",
    "  axs[task, 0].yaxis.set_label_coords(-0.5,1)\n",
    "  plot_task(axs[task], task_data[task][0], samples)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "awLQn1AqkqgZ"
   },
   "outputs": [],
   "source": [
    "class Base(nn.Module):\n",
    "  \"\"\"The base model.\n",
    "  \"\"\"\n",
    "  def __init__(self, hsize=32):\n",
    "    super(Base, self).__init__()\n",
    "    self.l1 = nn.Linear(784, hsize)\n",
    "    \n",
    "  def forward(self, x, return_z_also=False):\n",
    "    x = x.view(x.size(0), -1)\n",
    "    z = self.l1(x)\n",
    "    if return_z_also:\n",
    "        return z\n",
    "    else:\n",
    "        return z\n",
    "\n",
    "class Head(nn.Module):\n",
    "  \"\"\"The classifier head.\n",
    "  \"\"\"\n",
    "  def __init__(self, fe, hsize=32):\n",
    "    super(Head, self).__init__()\n",
    "    self.fe = fe\n",
    "    self.l2 = nn.Linear(hsize, 5)\n",
    "\n",
    "  def clf(self, z):\n",
    "    x = self.l2(F.relu(z))\n",
    "    return x\n",
    "\n",
    "  def forward(self, x, return_z_also=False):\n",
    "    z = self.fe.forward(x)\n",
    "    x = self.clf(z)\n",
    "    if return_z_also:\n",
    "        return x, z\n",
    "    else:\n",
    "        return x, z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "CxkFzfnlriSA"
   },
   "outputs": [],
   "source": [
    "def train(model, device, train_loader, optimizer, epochs=10, log_training=False):\n",
    "    for epoch in range(epochs):\n",
    "      model.train()\n",
    "      \n",
    "      for x, y in train_loader:\n",
    "        x, y = x.to(device), y.to(device).long()\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        output, _ = model(x)\n",
    "        loss = F.cross_entropy(output, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if log_training:\n",
    "          print('Train Epoch: {} \\tLoss: {:.6f}'.format(epoch, loss.item()))\n",
    "\n",
    "def test(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    accuracy_metric = Metrics()\n",
    "\n",
    "    for x, y in test_loader:\n",
    "      with torch.no_grad():\n",
    "        x, y = x.to(device), y.to(device).long()\n",
    "        output, _ = model(x)\n",
    "        test_loss += F.cross_entropy(output, y).item() # sum up batch loss\n",
    "\n",
    "        accuracy = compute_accuracy(output, y)[0].item()\n",
    "        accuracy_metric.update(accuracy)\n",
    "\n",
    "    test_loss /= len(test_loader)\n",
    "    return test_loss, accuracy_metric.avg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Metrics:\n",
    "    def __init__(self):\n",
    "        self.val = 0\n",
    "        self.sum = 0\n",
    "        self.count = 0\n",
    "        self.avg = 0\n",
    "\n",
    "    def update(self, val, n=1):\n",
    "        self.val = val\n",
    "        self.sum += val*n\n",
    "        self.count += n\n",
    "        self.avg = self.sum / self.count\n",
    "        \n",
    "def compute_accuracy(output, target, topk=(1,)):\n",
    "    maxk = max(topk)\n",
    "    batch_size = target.size(0)\n",
    "\n",
    "    _, pred = output.topk(maxk, 1)\n",
    "    pred = pred.t()\n",
    "    correct = pred.eq(target.view(1, -1).expand_as(pred))\n",
    "\n",
    "    result = []\n",
    "    for k in topk:\n",
    "        correct_k = correct[:k].view(-1).float().sum(0)\n",
    "        result.append(correct_k.mul_(100.0 / batch_size))\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 303
    },
    "id": "tvpaBTeDuHQq",
    "outputId": "ed90cbec-41f5-4c0e-dc17-a46aa36c1f44"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training on Task 1\n",
      "Testing on Task 1: accuracy -> 99.16158536585365\n"
     ]
    }
   ],
   "source": [
    "z_size = 32\n",
    "batch_size = 128\n",
    "\n",
    "# Creating the models\n",
    "base_model = Base(hsize=z_size).to(device)\n",
    "model_t1 = Head(base_model, hsize=z_size).to(device)\n",
    "model_t2 = Head(base_model, hsize=z_size).to(device)\n",
    "\n",
    "# Setting up the optimizers\n",
    "optimizer = optim.SGD(model_t1.parameters(), lr=0.1, momentum=0.9)\n",
    "x_train_1, t_train_1, x_test_1, t_test_1 = task_data[0]\n",
    "x_train_2, t_train_2, x_test_2, t_test_2 = task_data[1]\n",
    "\n",
    "# Setting up data\n",
    "t1_train_ds = TensorDataset(torch.Tensor(x_train_1), torch.Tensor(t_train_1))\n",
    "t1_test_ds = TensorDataset(torch.Tensor(x_test_1), torch.Tensor(t_test_1))\n",
    "t2_train_ds = TensorDataset(torch.Tensor(x_train_2), torch.Tensor(t_train_2))\n",
    "t2_test_ds = TensorDataset(torch.Tensor(x_test_2), torch.Tensor(t_test_2))\n",
    "t1_train_loader = DataLoader(t1_train_ds, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=True)\n",
    "t1_test_loader = DataLoader(t1_test_ds, batch_size=batch_size, shuffle=False, num_workers=0, drop_last=False)\n",
    "t2_train_loader = DataLoader(t2_train_ds, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=True)\n",
    "t2_test_loader = DataLoader(t2_test_ds, batch_size=batch_size, shuffle=False, num_workers=0, drop_last=False)\n",
    "\n",
    "print('Training on Task 1')\n",
    "train(model_t1, device, t1_train_loader, optimizer, epochs=10)\n",
    "\n",
    "loss, acc = test(model_t1, device, t1_test_loader)\n",
    "print(f'Testing on Task 1: accuracy -> {acc}')\n",
    "\n",
    "model_t1_backup = copy.deepcopy(model_t1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 405
    },
    "id": "QGkNEHxUUWCt",
    "outputId": "cd6ccc4d-69a6-4edc-9f2c-6e53123a5e83"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training on Task 2\n",
      "Testing on Task 2 using model_t2: accuracy -> 95.35164481715152\n",
      "Testing on Task 1 using the backbone features adapted after learning task 2: accuracy -> 20.88314344824814\n"
     ]
    }
   ],
   "source": [
    "print('Training on Task 2')\n",
    "train(model_t2, device, t2_train_loader, optimizer, epochs=10)\n",
    "\n",
    "loss, acc = test(model_t2, device, t2_test_loader)\n",
    "print(f'Testing on Task 2 using model_t2: accuracy -> {acc}')\n",
    "\n",
    "loss, acc = test(model_t1, device, t1_test_loader)\n",
    "print(f'Testing on Task 1 using the backbone features adapted after learning task 2: accuracy -> {acc}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "TCXJBaj9twY6"
   },
   "outputs": [],
   "source": [
    "class EBMAligner:\n",
    "    \"\"\"Manages the lifecycle of the proposed Energy Based Latent Alignment.\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        self.enabled = False\n",
    "\n",
    "        # Configs of the EBM model\n",
    "        self.ebm_latent_dim = 32\n",
    "        self.ebm_n_layers = 1\n",
    "        self.ebm_n_hidden_layers = 64\n",
    "        self.ebm_ema = None\n",
    "\n",
    "        # EBM Learning configs\n",
    "        self.max_iter = 17100\n",
    "        self.ebm_lr = 0.0001\n",
    "        self.n_langevin_steps = 30\n",
    "        self.langevin_lr = 0.1\n",
    "        self.ema_decay = 0.99\n",
    "\n",
    "        # EBM Loss config\n",
    "        self.alpha = 1.0\n",
    "\n",
    "    def ema(self, model1, model2, decay=0.999):\n",
    "        par1 = dict(model1.named_parameters())\n",
    "        par2 = dict(model2.named_parameters())\n",
    "        for k in par1.keys():\n",
    "            par1[k].data.mul_(decay).add_(par2[k].data, alpha=1 - decay)\n",
    "\n",
    "    def requires_grad(self, model, flag=True):\n",
    "        for p in model.parameters():\n",
    "            p.requires_grad = flag\n",
    "\n",
    "    def sampler(self, ebm_model, x, langevin_steps=30, lr=0.1):\n",
    "        \"\"\"The langevin sampler to sample from the ebm_model\n",
    "\n",
    "        :param ebm_model: The source EBM model\n",
    "        :param x: The data which is updated to minimize energy from EBM\n",
    "        :param langevin_steps: The number of langevin steps\n",
    "        :param lr: The langevin learning rate\n",
    "        :return: Samples from EBM\n",
    "        \"\"\"\n",
    "        x = x.clone().detach()\n",
    "        x.requires_grad_(True)\n",
    "        sgd = torch.optim.SGD([x], lr=lr)\n",
    "\n",
    "        for k in range(langevin_steps):\n",
    "            ebm_model.zero_grad()\n",
    "            sgd.zero_grad()\n",
    "            energy = ebm_model(x).sum()\n",
    "\n",
    "            (-energy).backward()\n",
    "            sgd.step()\n",
    "\n",
    "        return x.clone().detach()\n",
    "\n",
    "    def learn_ebm(self, prev_model, current_model, current_task_data, validation_data=None):\n",
    "        \"\"\"Learn the EBM.\n",
    "\n",
    "        current_task_data + prev_model acts as in-distribution data, and\n",
    "        current_task_data + current_model acts as out-of-distribution data.\n",
    "        This is used for learning the energy manifold.\n",
    "\n",
    "        :param prev_model: Model trained till previous task.\n",
    "        :param current_model: Model trained on current task.\n",
    "        :param current_task_data: Datapoints from the current incremental task.\n",
    "        :param validation_data: OPTIONAL, if passed, used for evaluation.\n",
    "        :return: None.\n",
    "        \"\"\"\n",
    "        ebm = EBM(latent_dim=self.ebm_latent_dim, n_layer=self.ebm_n_layers,\n",
    "                       n_hidden=self.ebm_n_hidden_layers).cuda()\n",
    "        if self.ebm_ema is None:\n",
    "            self.ebm_ema = EBM(latent_dim=self.ebm_latent_dim, n_layer=self.ebm_n_layers,\n",
    "                               n_hidden=self.ebm_n_hidden_layers).cuda()\n",
    "            # Initialize the exponential moving average of the EBM.\n",
    "            self.ema(self.ebm_ema, ebm, decay=0.)\n",
    "\n",
    "        ebm_optimizer = torch.optim.RMSprop(ebm.parameters(), lr=self.ebm_lr)\n",
    "\n",
    "        iterations = 0\n",
    "        prev_model.eval()\n",
    "        current_model.eval()\n",
    "        data_iter = iter(current_task_data)\n",
    "\n",
    "        print('Starting to learn the EBM')\n",
    "        while iterations < self.max_iter:\n",
    "            ebm.zero_grad()\n",
    "            ebm_optimizer.zero_grad()\n",
    "\n",
    "            try:\n",
    "                inputs, _ = next(data_iter)\n",
    "            except (OSError, StopIteration):\n",
    "                data_iter = iter(current_task_data)\n",
    "                inputs, _ = next(data_iter)\n",
    "\n",
    "            inputs = inputs.cuda()\n",
    "            _, prev_z = prev_model(inputs, return_z_also=True)\n",
    "            _, current_z = current_model(inputs, return_z_also=True)\n",
    "\n",
    "            self.requires_grad(ebm, False)\n",
    "            sampled_z = self.sampler(ebm, current_z.clone().detach(), langevin_steps=self.n_langevin_steps, lr=self.langevin_lr)\n",
    "            self.requires_grad(ebm, True)\n",
    "\n",
    "            indistribution_energy = ebm(prev_z)\n",
    "            oodistribution_energy = ebm(sampled_z)\n",
    "\n",
    "            loss = -(indistribution_energy - oodistribution_energy).mean()\n",
    "\n",
    "            loss.backward()\n",
    "            ebm_optimizer.step()\n",
    "            self.ema(self.ebm_ema, ebm, decay=self.ema_decay)\n",
    "\n",
    "            if iterations == 0 or iterations % 1000 == 0:\n",
    "                if validation_data is not None:\n",
    "                    accuracy = self.evaluate(prev_model, current_model, validation_data)\n",
    "                    print(\"Iteration: {:5d}, accuracy: {:5.2f}\".format(iterations, accuracy))\n",
    "                else:\n",
    "                    print(\"Iter: {:5d}\".format(iterations))\n",
    "\n",
    "            iterations += 1\n",
    "\n",
    "        self.enabled = True\n",
    "\n",
    "    def evaluate(self, previous_model, current_model, validation_data):\n",
    "        previous_model.eval()\n",
    "        current_model.eval()\n",
    "        accuracy_metric = Metrics()\n",
    "\n",
    "        for inputs, labels in validation_data:\n",
    "            inputs = inputs.cuda()\n",
    "            labels = labels.cuda()\n",
    "            _, current_z = current_model(inputs, return_z_also=True)\n",
    "            aligned_z = self.align_latents(current_z)\n",
    "\n",
    "            output = previous_model.clf(aligned_z)\n",
    "            accuracy = self.compute_accuracy(output, labels)[0].item()\n",
    "            accuracy_metric.update(accuracy)\n",
    "\n",
    "        return accuracy_metric.avg\n",
    "\n",
    "    def compute_accuracy(self, output, target, topk=(1,)):\n",
    "        maxk = max(topk)\n",
    "        batch_size = target.size(0)\n",
    "\n",
    "        _, pred = output.topk(maxk, 1)\n",
    "        pred = pred.t()\n",
    "        correct = pred.eq(target.view(1, -1).expand_as(pred))\n",
    "\n",
    "        result = []\n",
    "        for k in topk:\n",
    "            correct_k = correct[:k].view(-1).float().sum(0)\n",
    "            result.append(correct_k.mul_(100.0 / batch_size))\n",
    "        return result\n",
    "\n",
    "    def align_latents(self, z):\n",
    "        self.requires_grad(self.ebm_ema, False)\n",
    "        aligned_z = self.sampler(self.ebm_ema, z.clone().detach(), langevin_steps=self.n_langevin_steps, lr=self.langevin_lr)\n",
    "        self.requires_grad(self.ebm_ema, True)\n",
    "        return aligned_z\n",
    "\n",
    "\n",
    "class EBM(nn.Module):\n",
    "    \"\"\"Defining the Energy Based Model.\n",
    "    \"\"\"\n",
    "    def __init__(self, latent_dim=32, n_layer=1, n_hidden=64):\n",
    "        super().__init__()\n",
    "\n",
    "        mlp = nn.ModuleList()\n",
    "        if n_layer == 0:\n",
    "            mlp.append(nn.Linear(latent_dim, 1))\n",
    "        else:\n",
    "            mlp.append(nn.Linear(latent_dim, n_hidden))\n",
    "\n",
    "            for _ in range(n_layer-1):\n",
    "                mlp.append(nn.LeakyReLU(0.2))\n",
    "                mlp.append(nn.Linear(n_hidden, n_hidden))\n",
    "\n",
    "            mlp.append(nn.LeakyReLU(0.2))\n",
    "            mlp.append(nn.Linear(n_hidden, 1))\n",
    "\n",
    "        self.mlp = nn.Sequential(*mlp)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.mlp(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "H_mGonJaukqH",
    "outputId": "072d21ef-5f0b-4603-c7ab-0f5bd69a2838",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting to learn the EBM\n",
      "Iteration:     0, accuracy: 20.60\n",
      "Iteration:  1000, accuracy: 20.60\n",
      "Iteration:  2000, accuracy: 21.66\n",
      "Iteration:  3000, accuracy: 28.51\n",
      "Iteration:  4000, accuracy: 43.23\n",
      "Iteration:  5000, accuracy: 49.75\n",
      "Iteration:  6000, accuracy: 53.25\n",
      "Iteration:  7000, accuracy: 55.88\n",
      "Iteration:  8000, accuracy: 65.16\n",
      "Iteration:  9000, accuracy: 70.75\n",
      "Iteration: 10000, accuracy: 73.72\n",
      "Iteration: 11000, accuracy: 75.88\n",
      "Iteration: 12000, accuracy: 77.02\n",
      "Iteration: 13000, accuracy: 78.99\n",
      "Iteration: 14000, accuracy: 80.40\n",
      "Iteration: 15000, accuracy: 81.73\n",
      "Iteration: 16000, accuracy: 83.33\n",
      "Iteration: 17000, accuracy: 83.70\n"
     ]
    }
   ],
   "source": [
    "aligner = EBMAligner()\n",
    "aligner.learn_ebm(model_t1_backup, model_t2, t2_train_loader, t1_test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy from un-modified latents : 20.88314344824814\n",
      "Accuracy from aligned latents (using ELI): 83.44231377578363\n"
     ]
    }
   ],
   "source": [
    "def align(model, aligner, device, data_loader):\n",
    "  model.eval()\n",
    "  model_t1.eval()\n",
    "  latents = []\n",
    "  labels = []\n",
    "  accuracy_metric = Metrics()\n",
    "  accuracy_metric_adap = Metrics()\n",
    "    \n",
    "  for x, y in data_loader:\n",
    "    x, y = x.to(device), y.to(device).long()\n",
    "    _, z = model(x)\n",
    "\n",
    "    output = model_t1.clf(z)\n",
    "    accuracy = compute_accuracy(output, y)[0].item()\n",
    "    accuracy_metric.update(accuracy)\n",
    "\n",
    "    z_samples = aligner.align_latents(z.clone().detach())\n",
    "\n",
    "    output = model_t1.clf(z_samples)\n",
    "    accuracy = compute_accuracy(output, y)[0].item()\n",
    "    accuracy_metric_adap.update(accuracy)\n",
    "\n",
    "    latents.extend(z_samples.cpu().numpy())\n",
    "    labels.extend(y.cpu().numpy())\n",
    "\n",
    "  print('Accuracy from un-modified latents : {}'.format(accuracy_metric.avg))\n",
    "  print('Accuracy from aligned latents (using ELI): {}'.format(accuracy_metric_adap.avg))\n",
    "\n",
    "align(model_t2, aligner, device, t1_test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "EBM Anchoring.ipynb",
   "provenance": []
  },
  "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}