{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def type_info(obj):\n",
    "    print(type(obj))\n",
    "    for type_ in (Preprocessor, FeatureScaler, MinMaxScaler):\n",
    "        print(type_, isinstance(obj, type_))\n",
    "    print(obj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Preprocessor:\n",
    "    \n",
    "    def __init__(self):\n",
    "        pass\n",
    "    \n",
    "    def forward(self):\n",
    "        print('forward')\n",
    "        \n",
    "    def __repr__(self):\n",
    "        return 'Preprocessor()'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "forward\n"
     ]
    }
   ],
   "source": [
    "pre = Preprocessor()\n",
    "pre.forward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class '__main__.Preprocessor'>\n",
      "<class '__main__.Preprocessor'> True\n",
      "<class '__main__.FeatureScaler'> False\n",
      "<class '__main__.MinMaxScaler'> False\n",
      "Preprocessor()\n"
     ]
    }
   ],
   "source": [
    "type_info(pre)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeatureScaler(Preprocessor):\n",
    "    \n",
    "    def __init__(self, x):\n",
    "        super().__init__()\n",
    "        self.x = x\n",
    "        self.processed = False\n",
    "        \n",
    "    def process(self):\n",
    "        print('processing')\n",
    "        self.processed = True\n",
    "        \n",
    "    @classmethod\n",
    "    def from_list(cls, *args):\n",
    "        return cls(*args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "forward\n",
      "processing\n"
     ]
    }
   ],
   "source": [
    "scaler = FeatureScaler([1, 2, 3])\n",
    "scaler.forward()\n",
    "scaler.process()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class '__main__.FeatureScaler'>\n",
      "<class '__main__.Preprocessor'> True\n",
      "<class '__main__.FeatureScaler'> True\n",
      "<class '__main__.MinMaxScaler'> False\n",
      "Preprocessor()\n"
     ]
    }
   ],
   "source": [
    "type_info(scaler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1, 2]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scaler2 = FeatureScaler.from_list([1, 2])\n",
    "scaler2.x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MinMaxScaler(FeatureScaler):\n",
    "    \n",
    "    def __init__(self, x, y):\n",
    "        super().__init__(x)\n",
    "        self.y = y\n",
    "        \n",
    "    def scale(self):\n",
    "        return [arg/2 for arg in self.x]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "forward\n",
      "processing\n",
      "[0.5, 1.0, 1.5]\n"
     ]
    }
   ],
   "source": [
    "minmax = MinMaxScaler([1, 2, 3], 4)\n",
    "minmax.forward()\n",
    "minmax.process()\n",
    "print(minmax.scale())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class '__main__.MinMaxScaler'>\n",
      "<class '__main__.Preprocessor'> True\n",
      "<class '__main__.FeatureScaler'> True\n",
      "<class '__main__.MinMaxScaler'> True\n",
      "Preprocessor()\n"
     ]
    }
   ],
   "source": [
    "type_info(minmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(7, 8)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "minmax2 = MinMaxScaler.from_list(*[7, 8])\n",
    "minmax2.x, minmax2.y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StandardScaler(FeatureScaler):\n",
    "    \n",
    "    def __init__(self, x, y, z):\n",
    "        super().__init__(x)\n",
    "        self.y = y\n",
    "        self.z = z\n",
    "        \n",
    "    def backward(self):\n",
    "        return self.z, self.y, self.x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "forward\n",
      "(33, 22, 11)\n"
     ]
    }
   ],
   "source": [
    "ss = StandardScaler(11, 22, 33)\n",
    "ss.forward()\n",
    "print(ss.backward())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "forward\n",
      "(11, 9, 7)\n"
     ]
    }
   ],
   "source": [
    "ss2 = StandardScaler.from_list(*[7, 9, 11])\n",
    "ss2.forward()\n",
    "print(ss2.backward())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class '__main__.StandardScaler'>\n",
      "<class '__main__.Preprocessor'> True\n",
      "<class '__main__.FeatureScaler'> True\n",
      "<class '__main__.MinMaxScaler'> False\n",
      "Preprocessor()\n"
     ]
    }
   ],
   "source": [
    "type_info(ss2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TmpModule(nn.Module):\n",
    "    \n",
    "    def __init__(self, init_args):\n",
    "        super().__init__()\n",
    "        init_args.pop('self', None)\n",
    "        init_args.pop('__class__', None)\n",
    "        self.init_args = init_args\n",
    "        \n",
    "    def get_args(self):\n",
    "        print('init args:', self.init_args)\n",
    "    \n",
    "    @classmethod\n",
    "    def from_list(cls, arr):\n",
    "        return cls(*arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImgModel(TmpModule):\n",
    "    \n",
    "    def __init__(self, dim, hidden, bn):\n",
    "        super().__init__(locals())\n",
    "        self.dim = dim\n",
    "        self.hidden = hidden\n",
    "        self.bn = bn\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return 2 * x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_backend': <torch.nn.backends.thnn.THNNFunctionBackend at 0x118f9a320>,\n",
       " '_parameters': OrderedDict(),\n",
       " '_buffers': OrderedDict(),\n",
       " '_backward_hooks': OrderedDict(),\n",
       " '_forward_hooks': OrderedDict(),\n",
       " '_forward_pre_hooks': OrderedDict(),\n",
       " '_state_dict_hooks': OrderedDict(),\n",
       " '_load_state_dict_pre_hooks': OrderedDict(),\n",
       " '_modules': OrderedDict(),\n",
       " 'training': True,\n",
       " 'init_args': {'dim': 12, 'hidden': 4, 'bn': True},\n",
       " 'dim': 12,\n",
       " 'hidden': 4,\n",
       " 'bn': True}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img = ImgModel(12, 4, True)\n",
    "img.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_backend': <torch.nn.backends.thnn.THNNFunctionBackend at 0x118f9a320>,\n",
       " '_parameters': OrderedDict(),\n",
       " '_buffers': OrderedDict(),\n",
       " '_backward_hooks': OrderedDict(),\n",
       " '_forward_hooks': OrderedDict(),\n",
       " '_forward_pre_hooks': OrderedDict(),\n",
       " '_state_dict_hooks': OrderedDict(),\n",
       " '_load_state_dict_pre_hooks': OrderedDict(),\n",
       " '_modules': OrderedDict(),\n",
       " 'training': True,\n",
       " 'init_args': {'dim': 3, 'hidden': 6, 'bn': False},\n",
       " 'dim': 3,\n",
       " 'hidden': 6,\n",
       " 'bn': False}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img2 = ImgModel.from_list([3, 6, False])\n",
    "img2.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "init args: {'dim': 3, 'hidden': 6, 'bn': False}\n"
     ]
    }
   ],
   "source": [
    "img2.get_args()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing changing BaseModel to mixin (called ModelMixin). Ended up switching back bc regular inheritance works after all (not sure what caused prev issue, where from_path method received cls basemodel instead of child model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from htools.ml import ModelMixin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## super init version 1: probably more confusing, use version 2 below \n",
    "-nn.Module comes first here  \n",
    "-still a bit confused about exactly how order and init statements work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module, ModelMixin):\n",
    "    \n",
    "    def __init__(self, x_dim):\n",
    "        super().__init__()\n",
    "        ModelMixin.__init__(self, locals())\n",
    "\n",
    "        self.fc = nn.Linear(x_dim, 1)\n",
    "        self.act = nn.Sigmoid()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.act(self.fc(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., 3., 0.],\n",
       "        [4., 5., 4.],\n",
       "        [4., 1., 0.],\n",
       "        [1., 2., 3.]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randint(6, (4, 3), dtype=torch.float)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Net(\n",
       "  (fc): Linear(in_features=3, out_features=1, bias=True)\n",
       "  (act): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = Net(3)\n",
    "net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2570],\n",
       "        [0.8122],\n",
       "        [0.8215],\n",
       "        [0.5569]], grad_fn=<SigmoidBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(1, 3), (1,)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net.dims()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[((1, 3), True), ((1,), True)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net.trainable()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(0.178, 0.373), (-0.471, nan)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net.weight_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10356 characters]\n",
      "text/plain": [
       "<Figure size 576x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "net.plot_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weights saved from epoch 12.\n"
     ]
    }
   ],
   "source": [
    "net.save(12, '../data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 12 weights loaded from ../data/model_12.pth.\n",
      "Model parameters: {'x_dim': 3}\n",
      "Currently in eval mode.\n"
     ]
    }
   ],
   "source": [
    "net2 = Net.from_path('../data/model_12.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Net(\n",
       "  (fc): Linear(in_features=3, out_features=1, bias=True)\n",
       "  (act): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2570],\n",
       "        [0.8122],\n",
       "        [0.8215],\n",
       "        [0.5569]], grad_fn=<SigmoidBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net2(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Version 2 - simpler way to super init  \n",
    "-ModelMixin must come before nn.Module  \n",
    "-locals() must be passed to super init so mixin can store args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NetB(ModelMixin, nn.Module):\n",
    "    \n",
    "    def __init__(self, x_dim):\n",
    "        super().__init__(locals())\n",
    "\n",
    "        self.fc = nn.Linear(x_dim, 1)\n",
    "        self.act = nn.Sigmoid()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.act(self.fc(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NetB(\n",
       "  (fc): Linear(in_features=3, out_features=1, bias=True)\n",
       "  (act): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "netb = NetB(3)\n",
    "netb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4054],\n",
       "        [0.1432],\n",
       "        [0.1389],\n",
       "        [0.4869]], grad_fn=<SigmoidBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "netb(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[((1, 3), True), ((1,), True)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "netb.trainable()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(1, 3), (1,)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "netb.dims()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Weights saved from epoch 4.\n"
     ]
    }
   ],
   "source": [
    "netb.save(4, '../data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4 weights loaded from ../data/model_4.pth.\n",
      "Model parameters: {'x_dim': 3}\n",
      "Currently in eval mode.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "NetB(\n",
       "  (fc): Linear(in_features=3, out_features=1, bias=True)\n",
       "  (act): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "netb2 = NetB.from_path('../data/model_4.pth')\n",
    "netb2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4054],\n",
       "        [0.1432],\n",
       "        [0.1389],\n",
       "        [0.4869]], grad_fn=<SigmoidBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "netb2(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['weights', 'epoch', 'params'])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = torch.load('../data/model_4.pth')\n",
    "data.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing inheritance again (non mixin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from ml_htools.torch_utils import BaseModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SubModel(BaseModel):\n",
    "    \n",
    "    def __init__(self, dim, norm):\n",
    "        super().__init__(locals())\n",
    "        layers = [nn.Linear(dim, 1)]\n",
    "        if norm:\n",
    "            layers.append(nn.BatchNorm1d(1))\n",
    "        layers.append(nn.Sigmoid())\n",
    "        \n",
    "        self.layers = nn.Sequential(*layers)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.layers(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[5., 3., 1., 3., 4.],\n",
       "        [5., 2., 5., 1., 1.],\n",
       "        [0., 1., 2., 0., 1.]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randint(6, (3, 5), dtype=torch.float)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SubModel(\n",
       "  (layers): Sequential(\n",
       "    (0): Linear(in_features=5, out_features=1, bias=True)\n",
       "    (1): BatchNorm1d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): Sigmoid()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = SubModel(5, True)\n",
    "m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(-0.017, 0.101), (-0.348, nan), (1.0, nan), (0.0, nan)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.weight_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4756],\n",
       "        [0.7806],\n",
       "        [0.2366]], grad_fn=<SigmoidBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 weights saved to ../data/model_e1.pth.\n"
     ]
    }
   ],
   "source": [
    "m.save(1, '../data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileExistsError",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileExistsError\u001b[0m                           Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-17-c681115a8506>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'../data'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/pythonhm/ml_htools/ml_htools/torch_utils.py\u001b[0m in \u001b[0;36msave\u001b[0;34m(self, epoch, dir_, file, overwrite, verbose, **kwargs)\u001b[0m\n\u001b[1;32m     93\u001b[0m         \u001b[0mpath\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight_dir\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34mf'{file}_e{epoch}.pth'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     94\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexists\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0moverwrite\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 95\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0mFileExistsError\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     96\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     97\u001b[0m         data = dict(weights=self.state_dict(),\n",
      "\u001b[0;31mFileExistsError\u001b[0m: "
     ]
    }
   ],
   "source": [
    "m.save(1, '../data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.0804,  0.0819, -0.0600, -0.1556, -0.0313]])\n",
      "tensor([[ 0.1015,  0.1286, -0.2173, -0.0664,  0.0962]]) \n",
      "\n",
      "tensor([-0.3481])\n",
      "tensor([-0.3481]) \n",
      "\n",
      "tensor([1.])\n",
      "tensor([-0.0752]) \n",
      "\n",
      "tensor([0.])\n",
      "tensor([0.0381]) \n",
      "\n"
     ]
    }
   ],
   "source": [
    "penalty = m(x).std()\n",
    "penalty.backward()\n",
    "\n",
    "for p in m.parameters():\n",
    "    print(p.data)\n",
    "    p.data -= 5 * p.grad\n",
    "    print(p.data, '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(0.009, 0.148), (-0.348, nan), (-0.075, nan), (0.038, nan)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.weight_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4830],\n",
       "        [0.5242],\n",
       "        [0.5213]], grad_fn=<SigmoidBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'../data'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.weight_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2 weights saved to ../data/model_e2.pth.\n"
     ]
    }
   ],
   "source": [
    "m.save(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 weights loaded from ../data/model_e1.pth.\n",
      "Model parameters: {'dim': 5, 'norm': True}\n",
      "Currently in eval mode.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "SubModel(\n",
       "  (layers): Sequential(\n",
       "    (0): Linear(in_features=5, out_features=1, bias=True)\n",
       "    (1): BatchNorm1d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): Sigmoid()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1 = SubModel.from_path('../data/model_e1.pth')\n",
    "m1.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4756],\n",
       "        [0.7806],\n",
       "        [0.2366]], grad_fn=<SigmoidBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(-0.017, 0.101), (-0.348, nan), (1.0, nan), (0.0, nan)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m1.weight_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2 weights loaded from ../data/model_e2.pth. \n",
      "Model parameters: {'dim': 5, 'norm': True}\n",
      "Currently in train mode.\n"
     ]
    }
   ],
   "source": [
    "m.load_epoch(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(0.009, 0.148), (-0.348, nan), (-0.075, nan), (0.038, nan)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.weight_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4830],\n",
       "        [0.5242],\n",
       "        [0.5213]], grad_fn=<SigmoidBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 weights loaded from ../data/model_e1.pth. \n",
      "Model parameters: {'dim': 5, 'norm': True}\n",
      "Currently in train mode.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[(-0.017, 0.101), (-0.348, nan), (1.0, nan), (0.0, nan)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.load_epoch(1)\n",
    "m.weight_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.4756],\n",
       "        [0.7806],\n",
       "        [0.2366]], grad_fn=<SigmoidBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scratch - load_epoch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['model_e2.pth']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num = 2\n",
    "[file for file in os.listdir('../data') if file.endswith(f'_e{num}.pth')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['model_e1.pth']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(filter(re.compile(f'.*_e{1}.pth').match, os.listdir('../data')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['lines5000.txt', 'model_e1.pth', 'model_e2.pth', 'lines100.txt']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.listdir('../data')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
