recon trainer

Author: saswatpp

L2_recon.ipynb

@@ -79,7 +79,36 @@
     }
    ],
    "source": [
-    "os.listdir()"
+    "#Modified Decoder\n",
+    "\n",
+    "class Decoder(nn.Module):\n",
+    "    def __init__(self, input_width=28, input_height=28, input_channel=1):\n",
+    "        super(Decoder, self).__init__()\n",
+    "        self.input_width = input_width\n",
+    "        self.input_height = input_height\n",
+    "        self.input_channel = input_channel\n",
+    "        self.reconstraction_layers = nn.Sequential(\n",
+    "            nn.Linear(16 * 10, 512),\n",
+    "            nn.ReLU(inplace=True),\n",
+    "            nn.Linear(512, 1024),\n",
+    "            nn.ReLU(inplace=True),\n",
+    "            nn.Linear(1024, self.input_height * self.input_height * self.input_channel),\n",
+    "            nn.Sigmoid()\n",
+    "        )\n",
+    "\n",
+    "    def forward(self, x, data):\n",
+    "        classes = torch.sqrt((x ** 2).sum(2))\n",
+    "        classes = F.softmax(classes.squeeze(), dim=1)\n",
+    "\n",
+    "        _, max_length_indices = classes.max(dim=1)\n",
+    "        masked = torch.sparse.torch.eye(10)\n",
+    "        if USE_CUDA:\n",
+    "            masked = masked.cuda()\n",
+    "        masked = masked.index_select(dim=0, index=max_length_indices.squeeze().data)\n",
+    "#         t = (x * masked[:, :, None, None]).view(x.size(0), -1)\n",
+    "        reconstructions = self.reconstraction_layers(x)\n",
+    "        reconstructions = reconstructions.view(-1, self.input_channel, self.input_width, self.input_height)\n",
+    "        return reconstructions, None"
    ]
   },
   {
@@ -161,7 +190,8 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Training CapsuleNet"
+    "##  The histogram of L2 distances between the input and the reconstruction using the\n",
+    "## correct capsule or other capsules in CapsNet on the real MNIST images. "
    ]
   },
   {