addpretrainedmodel

Author: saswatpp

CapsNetMNIST.pth

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -12,6 +12,7 @@
     "import torch.nn as nn\n",
     "import torch.nn.functional as F\n",
     "from tqdm import tqdm\n",
+    "import os\n",
     "\n",
     "from model.net import *\n",
     "from utils.training import *\n",
@@ -27,10 +28,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
+    "model_path = os.getcwd()\n",
+    "\n",
     "args = {\n",
     "    'USE_CUDA': True if torch.cuda.is_available() else False,\n",
     "    'BATCH_SIZE': 32,\n",
@@ -50,9 +53,51 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 13,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['data',\n",
+       " 'model',\n",
+       " 'utils',\n",
+       " 'LICENSE',\n",
+       " '.gitignore',\n",
+       " '.git',\n",
+       " 'Capsule Network Train.ipynb',\n",
+       " 'Pretrain_Capsule.ipynb',\n",
+       " 'L2_recon.ipynb',\n",
+       " '.ipynb_checkpoints',\n",
+       " 'README.md',\n",
+       " 'CapsNetMNIS.pth ']"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "os.listdir()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<All keys matched successfully>"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "#Config for 49 16d vectors in the Primary Capsule. Set Softmax dimension to 0 in this case\n",
     "class Config:\n",
@@ -93,7 +138,7 @@
     "if args['USE_CUDA']:\n",
     "    net = net.cuda()\n",
     "    \n",
-    "net.load_state_dict(torch.load('./CapsNetMNIST.pth'), map_location='cpu')"
+    "net.load_state_dict(torch.load(os.path.join(model_path, 'CapsNetMNIST.pth'), map_location='cpu'))"
    ]
   },
   {
@@ -134,7 +179,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# torch.save(capsule_net.state_dict(), \"./CapsNetMNIST.pth \")"
+    "# torch.save(capsule_net.state_dict(), \"./CapsNetMNIST.pth\")"
    ]
   },
   {