Capsnet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Capsnet experiment \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Theory"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some text here \n",
    "![Layout](network.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keras implementation\n",
    "\n",
    "I'm trying to rebuild the network from the article. \n",
    "\n",
    "This means we are going to make a network with the following structure: \n",
    "- 1 convolutional layer with 256 filters, 9x9 kernel size, stride of 1. \n",
    "- 1 layer with primary capsules\n",
    "- 1 layer with digit capsules\n",
    "- Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load MNIST data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from keras.datasets import mnist\n",
    "from keras.utils import to_categorical\n",
    "\n",
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "x_train = x_train.reshape(-1, 28, 28, 1).astype('float32') / 255.\n",
    "x_test = x_test.reshape(-1, 28, 28, 1).astype('float32') / 255.\n",
    "y_train = to_categorical(y_train.astype('float32'))\n",
    "y_test = to_categorical(y_test.astype('float32'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABIEAAACBCAYAAABXearSAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMS4wLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvpW3flQAAGyFJREFUeJzt3Xu4TlUewPGlEU4uIUTJ6ZmcqVGe\nKJd0k0vJTCG6l3TXTUq656mUBtNMMUWkOGpKkksuD6fUQY+SQqioPIWISDIzKoYzf5h+/dZy9ts+\n7+3s/a7v56/fbq1375X3rP2+73r2b/0qlJSUGAAAAAAAAOS2A8p7AAAAAAAAAMg8FoEAAAAAAAA8\nwCIQAAAAAACAB1gEAgAAAAAA8ACLQAAAAAAAAB5gEQgAAAAAAMADLAIBAAAAAAB4gEUgAAAAAAAA\nD7AIBAAAAAAA4IGKWb5eSZavh19VSNN5eA/LT7reQ2N4H8sTczH+mIu5gbkYf8zF3MBcjD/mYm5g\nLsZfqPeQJ4EAAAAAAAA8wCIQAAAAAACAB1gEAgAAAAAA8ACLQAAAAAAAAB5gEQgAAAAAAMADLAIB\nAAAAAAB4gEUgAAAAAAAAD7AIBAAAAAAA4AEWgQAAAAAAADxQsbwHAMBfy5cvl3jWrFlW2+rVq0Od\nIy8vT+IRI0akZ2AAAAAAkIN4EggAAAAAAMADLAIBAAAAAAB4gEUgAAAAAAAAD7AnUAwsWrRI4ilT\nplhtgwcPzvZwgP3Mnz/fOi4uLpa4sLDQatu7d6/EmzdvlvjHH39M6toVKlSQOD8/32q7++67kzon\nAAAAAOQingQCAAAAAADwAItAAAAAAAAAHqhQUlKSzetl9WK5okePHhK76WA6teY3VPjtLqHwHpaf\ndL2HxqT5fXT/LidPnpzyOY899liJW7ZsKfG2bdusfr1795b4tNNOs9qmTZuW8jgygLkYf5GdiygT\n5mKS9uzZI/Ho0aOttoEDB0pcv359q02nDlevXj0dQ2Eu5gbmYgYsWLBA4j59+gT2Gzt2rMTNmzdP\n9nLMxdzAXIy/UO8hTwIBAAAAAAB4gEUgAAAAAAAAD5AOFlE65aVZs2YSf/3111Y/0sGSN2PGDInD\nVqZy54uuTOVq06aNxA0bNizj6ErFo7al0I8uN2rUyGojHQwZ4u1c3L59u8S7du2y2ipVqiRxzZo1\nszamFDAXy2D37t0SDxgwQOK//vWvoc8xbtw4iXv27JmOYXk7F3MMczEN9Bw1xpguXbpIXFRUFPi6\nXr16STxy5EirrXLlymEvz1zMDczFBF555RXreNSoURJ37dpV4r59+2ZtTKUgHQwAAAAAAAD7sAgE\nAAAAAADgARaBAAAAAAAAPFCxvAeQTkuXLpU4Pz/faqtdu3a2h5MSXfba3QcIif38888SL1myROJ+\n/fpZ/ZYtWyaxu7dFkLLsCdSqVSuJdcn0Bg0ahLoWAESBLvk9YsQIibds2WL1q1evnsRTp0612lq3\nbp2h0SFT3M/FBx98UGK9D1DVqlWtfv/9738l1p/HxtifyWnaEwjA/82ZM8c6TrQPkFZYWCixe6/W\nv0cAH+jPMGOMmTBhgsRXXXWV1XbEEUdI3L1798wOLM14EggAAAAAAMADLAIBAAAAAAB4INbpYOvW\nrbOO27dvL3GdOnWsNv3ocrdu3TI7sCRs3rzZOtapSlqTJk2yMZxY+eyzz6zj/v37Szxz5kyJE6Vy\n6TQGY4ypUqWKxLrEppvisGnTJondxwfff/99iXUJwYceemj//wkAaePeP3VZc+2YY46xjuvXr5+x\nMcXJAw88YB0/9thjEidKgf32228lPuecc6w2N3UM0aRLTM+ePdtqGzp0qMTVqlWTWH++uf0++ugj\nq839rAWwz4oVKwLbmjZtGuocHTp0sI6/+uorifW8NMZO7dXc+z/pYPDNiy++aB1fc801EleqVMlq\n0/NFp4bFAU8CAQAAAAAAeIBFIAAAAAAAAA+wCAQAAAAAAOCBWO8J5JYv/eGHH0qNjTFm+vTpEuu9\ng2rUqJGh0ZXNxx9/bB2vWbOm1H7nnntuNoYTeXrPDzfPWe8DpJ1xxhnWsd7r56KLLrLagsq4Dxs2\nzDq+9dZbJX766acDx/vII49IzJ5AqXH3mNDHjRo1yvZwEBG6bOcbb7xhtbl7rv3i5JNPto7nzZuX\n/oFFmN7HTOfA6z2AjDGmcuXKEnfq1Eni4cOHW/2OPPJIib/77rt0DROK+91mypQpEk+cOFHiMWPG\nWP0OO+ywUOdfunSpxOedd15gv7/85S8Sn3XWWVZbos+4q6++OtQ4gChYuXKldXzfffdJfOaZZ1pt\nN954o8QVK4b7eaX37NHfSV3FxcXWcX5+fqn98vLyrGO9R0n16tVDjWnbtm2h+gG5RO/J5X5O6b0Q\na9asabXF+TONJ4EAAAAAAAA8wCIQAAAAAACAB2KdDrZw4cLQfZs3by5x2EciM02XF7/llltCveZ3\nv/tdpoYTeR9++KHEHTt2lNh9PF6n+E2bNk3itm3bpn1M+jHASZMmWW36/UX67N271zouKSmRuGvX\nrtkeDjJMp37qtJOtW7da/fTfQaIy5tqCBQtSHF28LV68WGJdAtX999MpYFOnTpVYl4R3X9e4ceO0\njRO/uvvuu63jZ599ttR+o0ePto6DUrR27txpHffo0SPw2p07d5b4kksukfi9996z+n3++ecS//nP\nf7baatWqFXh+IGoKCwutY73dgLv1QLdu3SQOWyr6hRdekHjdunVWm/5M0/2M2b+Mezrp/w8gl61f\nv17i7t27S6znnuull15K+bpr1661jnVq6ZVXXimxTj/NBJ4EAgAAAAAA8ACLQAAAAAAAAB6IdTrY\nO++8E7qvruK0fPlyiQ8++GCrn35kumrVqskPLgT9SNmqVasC+x100EES9+/fP6NjKm/60fTevXtb\nbTNmzJB4x44dEtetW9fqN378eIkzkQKmHXLIIRK7f0tBFYmQmrlz51rHBx54oMRNmjTJ9nCQBrpC\nypAhQ6y2OXPmSPz9999L7KbG7tmzR2JdzcqY8JWRct0rr7xiHV933XWl9rv55putY/c9+cXYsWMD\nr0VqZvro7wfuexjErdgVxK1MunHjRondFLI+ffpIrCukJPo7OP30063jSpUqhRpXXOl7mfvdRH+n\n3L59u9Wmv/twvypfn332mcSvvvpqRq/15ptvhup39tlnZ3QcWth7BxA3AwYMsI51OvWWLVskdlPi\nCwoKJA6b0uxuVTJixIhSr2uMnR6mx0g6GAAAAAAAAFLGIhAAAAAAAIAHWAQCAAAAAADwQOz2BNL7\nPrg51WGNGTMmsE2XIdc5e0cddVRS19JWr15tHT/11FOhXqf3M9J5+LloyZIlErtl+IJKQDdo0MDq\n16JFiwyNbn+6NK7OI3f17ds3G8Pxwtdff20d6z2zTjrppGwPB0nQZcaNsUtzunv9JOOmm26yjv/2\nt7+lfM64WrRokcTufeg///mPxHofiEcffdTql5eXV+q5p02bFnhdXeYUqXnyyScl1vvhuY488kiJ\nGzduHOrcAwcOtI71/bRdu3ZWm/7+sXjxYolnzZoVeH4fyk3reaVLebt7Eubn50v80UcfWW16T4ir\nrrpKYncPC2Terl27JN67d29gv3r16lnHen/CdKtSpUrGzm2Mfb/Ilf3cdu/eLbG7P0syKla0fzLn\n+u+xOJk/f751fP3110usv7+sXLnS6qffU/1b+91337X66c+7fv36WW36t/zWrVslnjhxotVv5MiR\nEh9wgP0MTs+ePSXu0qWLyRaeBAIAAAAAAPAAi0AAAAAAAAAeiF06mH6kb9KkSaFf16xZM4l1mXU3\nNay4uFjiE044QeJx48ZZ/c4777zQ1/7FP/7xD+tYlxJ1nXLKKRLff//9Zb5WXLz//vvW8eWXXx7q\ndX/6058k7tWrl9X2448/pj6wkPSj3646depIfOONN2ZjOF7Qj2obY8wf//jHchoJfou+x82ePVvi\ne+65J6nztW7dWmK3jPnJJ58s8aGHHprU+XORvkfpEqjG2OkquvR4jRo1As/3xRdfSOyWFz/88MMl\nrl27ttWmH9fW9+hOnTpZ/ebMmSPxmjVrrLbTTjtN4qZNmwaOMe6+/PJL6/jFF18M9bpWrVpJ7Kaq\naMOHD5d45syZVpsuT37qqadabfrvR5eLdz9zH374YYl1ilqu0v+eOlU9UbpkIg8++GCpsTHBafGJ\n6Ne4r7vhhhskdj9L9XvskwULFki8YcOGwH7XXHONdVy/fv2MjSlZ27Ztk/iNN94I7KfTQKP4/5GM\nhQsXSnzGGWekfD43/Ut/HoXlptg2atRIYp2alCspedkyaNAg61hvz6Hvd+497umnn5ZYfx967LHH\nrH46HWzevHlWm04909+J3BRE/b1Uv8YY+zMzm3gSCAAAAAAAwAMsAgEAAAAAAHggdulgYem0AWOM\nefPNNyWuWrWqxBdffHFgv+uuu05iXa3BGHsn/USPpS9fvlziyZMn/9awxR133FHqeHPBsmXLJO7R\no4fVtnHjxsDX6cc5X375ZYmrVauWvsGVkVvhQ9N/WwUFBdkYjhf++c9/WsdXX311OY0ELl0ZwRi7\n6teKFSuSOqeulFBYWChxec77ONFpWG5aiK48lCgF7KeffpL4sssuk/hf//qX1U+fw/1s1Y9Qh01j\ncenx63vqfffdZ/XTf3eJ/r+iaujQodZxohRn/V1kyJAhoc6vq8K5dFqlTiUxxk7D1pU83ZSle++9\nV+J0VPuLuuOOO07iTz75JGvXTXYe6deNGjUqsJ/ewuDYY4+12nQ6Q5MmTSTWFdDiSleXTPRvnKhy\nWCI6vVOnnrnc+3UydOU+Xf3YlUxqU9Tp+5Cb4rZp06Yyn8+tSD19+vQynyPsa4455hjr2E2bDqLT\nut3qhBMmTAh1jjjSvwmNse/D+vvL8ccfb/WrW7duqefTv8GNMeb111+XWKeGGbN/JbFfuBU633rr\nLYkbNmxY6muyjSeBAAAAAAAAPMAiEAAAAAAAgAdYBAIAAAAAAPBAzu4JVKtWLes4aF+dAw6w18HO\nOussiXUpv/Hjx1v9dL7gxIkTrTZdRlCXsU+Ug+ruYRQ2/zMudC5yx44dJXZL6CVyzjnnSFxcXFzq\nf8+0v//979bx+vXrJdZ7KRhjzLBhw7IyJh/s2LFDYjcP/8QTT8z2cLz3zjvvSKz36XnuuedCvd7d\nf+D3v/+9xG75+KOPPjqZIeL/Vq5cKbFbJtctAR6kqKhI4g8++EBid78Mvaebu7+b7pvsXiaaLh/v\nlmrWex+4n89R3SNI76/k7jGg9wb5wx/+YLXpss9HHHFEqGvp87n7jsyePVvi9u3bW216j0M9DnfP\nRB/2AdJmzJgh8apVq5I6h95nwt1jqbx88cUXEuv5Zoy9R4beEyjZvd+iJOy9yv39ENbgwYNDnT8Z\nzz//vHWsf6u416pXr57Eeg/UXKFLxLu/vxL9HtN7sOnfLu4cCPLdd99Zx3quuHbt2iWxvhe79xH9\neeq+j/p3VJUqVST+9NNPQ403F9SuXds6DvvdJoi7H96GDRskTrRXl/7uceGFF6Y0hmzgSSAAAAAA\nAAAPsAgEAAAAAADggdilgz311FOh+r399tvWsX4kXpfzTKRFixYSu+lgupS8W/KvZ8+eErul5ILc\ndddd1nFeXl6o18WFft90ak9ZHoXt37+/xIMGDZI40+lguqT9I488YrXpcqjuo/NIH52i4M4NN8UF\n6afvd8YYc+WVV0qsH1NOlAbSoUMHiXUKmTHGHHrooSmOEGFs3brVOv75558l1qVl9f3VGGPWrl1b\n6vncdIg6depI7JaI1yl/blsy9KP5bol4nSJ1//33W2267HWU/Pvf/5bYTalp1KiRxHPmzLHawqaA\naYnSXXT5eHcclStXlnjkyJEpjSGX6O8ByZZI1yWh9ZYC6TB//nzr+LXXXkvr+XU5ZuzPLc0dNq1I\nc1OtH3roIYn1PeG2226z+un57OrSpYvETZs2LfOY4sQtEe8eB2nTpk0mhiP0Z6v+PH7yySetfnor\nikqVKlltU6dOlVinZuba78hMe/zxxyV+9dVXrbZE6Xha2PWFqOBJIAAAAAAAAA+wCAQAAAAAAOAB\nFoEAAAAAAAA8UCFRqbMMSPliunzw0KFDrTa9H8VBBx1ktV100UUSjx49OtS19L+Ne617771XYjc/\nsGLFX7da2r17d+D5W7ZsKbFbEjbZ0pMJpKsOZVLvof7/0f9ebrnegoICiXVpRmOMadCggcRLly6V\nuG7duskMKSGdm6vL37q53frv6uWXX077OBzprCWa1YmfKp277t6zJk+eLPHYsWOttrlz52ZsTJ07\nd7aO9T45v6Fc52KQcePGWcc6V3348OFWm97Xa8+ePRIfcsghVr8+ffpI3LdvX4kPPvjglMYaAbGZ\ni0888YTEel81Y4zp1KmTxO5eM5qecxdccIHEl19+udXv3HPPTXqcqXD3JNH7Zbj7n+i/VxOhuThm\nzBiJe/fubbXpfR70/c6lywpv3rzZatPlkvUeiR988IHVL9F+Bw888IDEDz/8cGC/LIvNXIwKPV8S\nfR8eNmyYxGH3b9y7d2+yw4rMXAz6vupq27atdaz3ZtH3U71/ljHGzJs3L9Q49H03HaXk3f0TZ86c\nKbEuLZ4C5mJuiMxcTDd3j6x+/fpJrL8Du78z9Hdbd29Ffc/T++jpz+1yEOo95EkgAAAAAAAAD7AI\nBAAAAAAA4IHYlYhfvXq1xG5Z0kcffVRiXaY9WfrxyyuuuMJq06kSzzzzjNUWlAKm08SMsUvXZiD9\nK1J69eol8fjx4yV2Sx3eddddEuuSucYY07x5c4nTnQKm07+Msf+WdOlkN/3hpptuSus4fOamL+jU\nlaKiIol1mWJjjGnYsKHE3377bYZGt7+JEydax2VIB4sM/Xc+ePBgq82dE0H0I+Y6XcQYY9q3b5/8\n4JAWhx12WGBbohQwbcCAARLre3TVqlWTH1gauY9d67RQnRoWZYnuXTp9R5cSTwf370Ony7lj0o+6\n65RQN60b0abni1uKWtPpvNifm9aly0N///33Eu/cuTPwHHrunH766Vbb9OnTUx2itTWG+301TSlg\nQKQtW7ZM4iFDhlht7vf4X+jfrMYY8/zzz0vsrj20a9dOYv17JA5ye+UBAAAAAAAAxhgWgQAAAAAA\nALwQu+pgGzdulNh9VPWyyy6TuHv37qleKqENGzZI7KaBzJ8/X2L92K2ubmWMMbNmzcrM4EpXrru9\nL1myROKOHTtK/MMPPwS+xq28oKsLde3aNZlhBNLV3oyxq8HpayWqzJIFsam8oHfL/+abb6y2a6+9\nVuLt27dL/NVXX1n9Nm3aVOq53dRJXRXQfUzz/PPPl7hbt26B4/38888ldqubuJX7gowaNSpUPxOh\nygv631L/O7rcFJQ777xTYjdV1hOxmYu33HKLxCNGjAjsp+eOWw0u3ffbCInMXNTpl3p+GWM/sr5l\ny5ZUL2VOOeUUid2KivozTldjdenKQmeffXbKY0pBbOYiEorMXAxbHSwddOWw66+/3mpLx2+EvLw8\niXXaSoYwF3NDZOZiWG515kGDBkm8bt06id3qYIcffrjEeosWt0KntmbNGuv4qKOOKttgs4PqYAAA\nAAAAANiHRSAAAAAAAAAPsAgEAAAAAADggdjtCXTppZdKPGHCBKtt0qRJEmd6T6BEiouLJdYlH8u5\nDHwkczzd/Vb0XgX7XVj9rSaTp+3+rSc6h94fY+3atWW+VoZEKt967ty5Ei9evNhq+/LLLyUePXp0\nUuc/+uijJdalx1u0aGH1S3fJ5CyIzFzUcyDRfNDvpzHG5Ofnp3rpuIvUXExk0aJFEj/xxBNWm55L\nuiRq3bp1MzmkKInMXExk1apVEm/dujXl87Vp00Zidy+wXbt2STxlyhSrTX//0vELL7yQ8phSEJu5\niIQiMxf174dp06alerr9NGvWrNTzx628dCmYi7khMnNRc/eoe++99yResWKF1bZt2zaJ9Xfbk046\nyeq3cOHCdA4xStgTCAAAAAAAAPuwCAQAAAAAAOCB2KWD6ce6qlevbrUVFRVJ7D7yhWg+3qcfPTfG\nmMcff1xi9/E+/beqy9Pu3Lkz1LUSpYN16NDBarvjjjsk7tSpU6jzZ0GkHrWtX7++xG6Kli7JqB99\nNsaYgoICic8880yJdalGY+z5XbNmzdQGGy2RmYvz5s2TOFE6WKtWrazjKlWqpHrpuIvUXETSIjMX\no0iXrTfGmK5du0qsU7mfeeYZq98ll1yS2YHZmIu5ITJzUX8vveGGG6y2wsLCwNd169ZNYr2lwMCB\nA61+Bx54oMS6hHsOYC7mhsjMxdtvv11i93NGz9PGjRtbbYMHD5b4p59+kvj888+3+um5mGNIBwMA\nAAAAAMA+LAIBAAAAAAB4gEUgAAAAAAAAD8RuTyCdP+vuQzJmzBiJdXlpGGMilOOZDkuWLJFY7yNk\njDHr168v9TXu37req0bnjxpjTLVq1VIdYiZEKt969erVwSdX/9YxLOGeaTk1Fz0VqbmIpDEXy+Ct\nt96SWO/n1rlzZ6vfjBkzsjYmw1zMFZGcizt27LCOly9fHti3ZcuWEleuXDmdw4gL5mJuKNe52L59\ne4n13nPu/j2tW7eW2P0dePzxxydz6VzCnkAAAAAAAADYh0UgAAAAAAAAD1Qs7wGU1T333CNx27Zt\nrTZSwPxxwgknSPzSSy+V40j8xXwDAH/UqlWrvIcAZFWNGjWs41NPPbWcRgL454orrpC4Q4cOVtuF\nF16Y7eHkHJ4EAgAAAAAA8ACLQAAAAAAAAB6IXXUwJC2SlRdQJlReyA3MxfhjLuYG5mIZLF26VOLu\n3btLXFRUZPUrKCjI2pgMczFXMBfjj7mYG5iL8Ud1MAAAAAAAAOzDIhAAAAAAAIAHWAQCAAAAAADw\nAHsC+YMcz/gj3zo3MBfjj7mYG5iL8cdczA3MxfhjLuYG5mL8sScQAAAAAAAA9mERCAAAAAAAwAMs\nAgEAAAAAAHiARSAAAAAAAAAPsAgEAAAAAADgARaBAAAAAAAAPMAiEAAAAAAAgAdYBAIAAAAAAPAA\ni0AAAAAAAAAeqFBSUlLeYwAAAAAAAECG8SQQAAAAAACAB1gEAgAAAAAA8ACLQAAAAAAAAB5gEQgA\nAAAAAMADLAIBAAAAAAB4gEUgAAAAAAAAD7AIBAAAAAAA4AEWgQAAAAAAADzAIhAAAAAAAIAHWAQC\nAAAAAADwAItAAAAAAAAAHmARCAAAAAAAwAMsAgEAAAAAAHiARSAAAAAAAAAPsAgEAAAAAADgARaB\nAAAAAAAAPMAiEAAAAAAAgAdYBAIAAAAAAPAAi0AAAAAAAAAeYBEIAAAAAADAAywCAQAAAAAAeIBF\nIAAAAAAAAA+wCAQAAAAAAOCB/wEBFqCCl2vu4QAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x1e58ab74a20>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_samples = 10\n",
    "\n",
    "plt.figure(figsize=(n_samples * 2, 3))\n",
    "for index in range(n_samples):\n",
    "    plt.subplot(1, n_samples, index + 1)\n",
    "    sample_image = x_train[index].transpose(2,1,0).squeeze()\n",
    "    plt.imshow(sample_image, cmap='binary')\n",
    "    plt.axis(\"off\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First some basic layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(60000, 28, 28, 1)\n"
     ]
    }
   ],
   "source": [
    "import keras\n",
    "print(x_train[0:].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "input = keras.layers.Input(shape=x_train[0,:].shape, name='input')\n",
    "\n",
    "# the first layer is just a 'simple' convolutional layer with 256, 9 x 9 filters and a stride of 1. \n",
    "conv1 = keras.layers.Conv2D(filters=256, \n",
    "                            kernel_size=9, \n",
    "                            strides=1, \n",
    "                            padding='valid', \n",
    "                            activation='relu',\n",
    "                            name='conv1')(input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# now we are going to make the first primary capsule\n",
    "# This is a convolutional layer with 32 channels of convolutional capsules with 8 dimensions. Each primary capsule\n",
    "# contains 8 convolutional units with a 9 x 9 filter and a stride of 2. \n",
    "\n",
    "# following the paper this means the following: \n",
    "primary_capsule_channels = 32 #of 6 x 6 in size\n",
    "primary_capsule_amount = 6 * 6 * 32\n",
    "primary_capsule_capsules = 8\n",
    "\n",
    "primary_capsule = keras.layers.Conv2D(filters=primary_capsule_channels * primary_capsule_capsules,\n",
    "                                      kernel_size=9,\n",
    "                                      strides=2,\n",
    "                                      padding='valid',\n",
    "                                      activation='relu',\n",
    "                                      name='primary_capsules')(conv1)\n",
    "\n",
    "\n",
    "# NOTE: I can't find out any words from the paper whether the\n",
    "# PrimaryCap convolution does a ReLU activation or not before\n",
    "# squashing function, but experiment show that using ReLU get a\n",
    "# higher test accuracy. So, which one to use will be your choice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "primary_capsule_reshape = keras.layers.Reshape(target_shape=[-1,\n",
    "                                                             primary_capsule_capsules], \n",
    "                                               name='primarycap_reshape')(primary_capsule)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: since we used a kernel size of 9 and no padding (for some reason, that's what \"valid\" means), the image shrunk by 9-1=8 pixels after each convolutional layer (28×28 to 20×20, then 20×20 to 12×12), and since we used a stride of 2 in the second convolutional layer, the image size was divided by 2. This is how we end up with 6×6 feature maps.\n",
    "\n",
    "Next, we reshape the output to get a bunch of 8D vectors representing the outputs of the primary capsules. The output of conv2 is an array containing 32×8=256 feature maps for each instance, where each feature map is 6×6. So the shape of this output is (batch size, 6, 6, 256). We want to chop the 256 into 32 vectors of 8 dimensions each. We could do this by reshaping to (batch size, 6, 6, 32, 8). However, since this first capsule layer will be fully connected to the next capsule layer, we can simply flatten the 6×6 grids. This means we just need to reshape to (batch size, 6×6×32, 8).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(?, 20, 20, 256)\n",
      "(?, 6, 6, 256)\n",
      "(?, ?, 8)\n"
     ]
    }
   ],
   "source": [
    "print(conv1.get_shape())\n",
    "print(primary_capsule.get_shape()) \n",
    "print(primary_capsule_reshape.get_shape()) # 32*6*6 = 1152. Dimensions are: [None, 1152, 8]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The squashing function needs to be implemented. This is a copy from another example, were they noted that there is a change that the squared norm becomes a nan. To prevent this from happinging we implement a small factor epsilon. \n",
    "\n",
    "Squash function: \n",
    "\n",
    "$\\mathbf{s}_j = \\frac{\\|{\\mathbf{s}_j}\\|^2}{1 + \\|{\\mathbf{s}_j}\\|^2} \\frac{\\mathbf{s}_j}{\\|{\\mathbf{s}_j}\\|} $ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import keras.backend as K\n",
    "import tensorflow as tf\n",
    "\n",
    "def squash(vectors, axis=-1, epsilon=1e-7, name=None):    \n",
    "    squared_norm = K.sum(K.square(vectors), axis=axis, keepdims=True)\n",
    "    safe_norm = K.sqrt(squared_norm +  epsilon)\n",
    "    squash_factor = squared_norm/ (1.0 + squared_norm)\n",
    "    unit_vector = vectors / safe_norm\n",
    "    squased_vector = squash_factor * unit_vector\n",
    "    \n",
    "    return squased_vector\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "primary_capsule_output = keras.layers.Lambda(squash, name='primary_capsule_squash')(primary_capsule_reshape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Digit capsules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$ \\pmatrix{ \\mathbf{A} & \\mathbf{B} & \\mathbf{C} \\\\ \\mathbf{D} & \\mathbf{E} & \\mathbf{F} } \\times \\pmatrix{ \\mathbf{G} & \\mathbf{H} & \\mathbf{I} \\\\ \\mathbf{J} & \\mathbf{K} & \\mathbf{L} } = \\pmatrix{ \\mathbf{AG} & \\mathbf{BH} & \\mathbf{CI} \\\\ \\mathbf{DJ} & \\mathbf{EK} & \\mathbf{FL} } $\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the code below we are going to implement the next layer. \n",
    "Especially the Weight matrixes and prediction vectors: \n",
    "\n",
    "$\\hat{u}_{j|i} = \\mathbf{W}_{ij}\\mathbf{u}_i$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "number_of_capsules_digit = 10\n",
    "number_of_dimensions_digit = 16 \n",
    "batch_size = 50 # set for this experiment on 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a weight matrix with the dimensions:\n",
    "#AmountOfCapsulesFirstLayer, # AmountOfCapsulesSecondLayer, #DimensionsOfSecondLayer, #DimensionsOfFirstLayer\n",
    "\n",
    "W_init = K.random_normal(shape=(1, \n",
    "                                primary_capsule_amount, \n",
    "                                number_of_capsules_digit, \n",
    "                                number_of_dimensions_digit, \n",
    "                                primary_capsule_capsules), \n",
    "                        stddev=0.01, \n",
    "                        dtype=tf.float32\n",
    "                        )\n",
    "W = K.variable(W_init, name=\"W\")\n",
    "\n",
    "# repeat this for all images in the batch\n",
    "W_tiled = K.tile(W, [batch_size, 1, 1, 1, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [],
   "source": [
    "primary_capsule_output_expand = K.expand_dims(primary_capsule_output, -1)\n",
    "primary_capsule_output_tile = K.expand_dims(primary_capsule_output_expand, 2)\n",
    "primary_capsule_output_tiled = K.tile(primary_capsule_output_tile, [1, 1, number_of_capsules_digit, 1, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor(\"MatMul_2:0\", shape=(50, 1152, 10, 16, 1), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "predictions_of_second_capsules = tf.matmul(W_tiled, primary_capsule_output_tiled)\n",
    "print(predictions_of_second_capsules)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Routing by agreement\n",
    "calculate c:\n",
    "\n",
    "$ c_{ij} = \\frac{exp(b_{ij})}{\\sum\\nolimits_{K} exp(b_{ik})}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor(\"raw_weights_2:0\", shape=(50, 1152, 10, 1, 1), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "raw_weights = tf.zeros([batch_size, primary_capsule_amount, number_of_capsules_digit, 1, 1],\n",
    "                       dtype=np.float32, name=\"raw_weights\")\n",
    "print(raw_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\mathbf{s}_j = \\sum\\limits_{i}{c_{i,j}\\hat{\\mathbf{u}}_{j|i}}$ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor(\"routing_weights_2:0\", shape=(50, 1152, 10, 1, 1), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "routing_weights = tf.nn.softmax(raw_weights, dim=2, name=\"routing_weights\")\n",
    "print(routing_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$ \\pmatrix{1 & 2 & 3 \\\\ 4 & 5 & 6} \\circ \\pmatrix{10 & 100 & 1000} = \\pmatrix{1 & 2 & 3 \\\\ 4 & 5 & 6} \\circ \\pmatrix{10 & 100 & 1000 \\\\ 10 & 100 & 1000} = \\pmatrix{10 & 200 & 3000 \\\\ 40 & 500 & 6000} $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor(\"Mul_3:0\", shape=(50, 1152, 10, 16, 1), dtype=float32)\n",
      "Tensor(\"weighted_sum_1:0\", shape=(50, 1, 10, 16, 1), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "weighted_predictions = tf.multiply(routing_weights, predictions_of_second_capsules)\n",
    "weighted_sum = tf.reduce_sum(weighted_predictions, axis=1, keep_dims=True,\n",
    "                             name=\"weighted_sum\")\n",
    "print(weighted_predictions)\n",
    "print(weighted_sum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "caps2_output_round_1 = squash(weighted_sum, axis=-2,\n",
    "                              name=\"caps2_output_round_1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor 'mul:0' shape=(50, 1, 10, 16, 1) dtype=float32>"
      ]
     },
     "execution_count": 149,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "caps2_output_round_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(14, 4)\n"
     ]
    }
   ],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}