phase3-copy-baseline(nopretrained).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "\n",
    "# pretty print all cell's output and not just the last one\n",
    "InteractiveShell.ast_node_interactivity = \"all\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# External Lib imports\n",
    "import os\n",
    "import re\n",
    "import html\n",
    "import pickle\n",
    "import argparse\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "from functools import partial\n",
    "\n",
    "os.environ['QT_QPA_PLATFORM'] = 'offscreen'\n",
    "\n",
    "# FastAI Imports\n",
    "from fastai import text, core, lm_rnn\n",
    "\n",
    "# Torch imports\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "# Mytorch imports\n",
    "from mytorch import loops, lriters as mtlr, dataiters as mtdi\n",
    "from mytorch.utils.goodies import *\n",
    "\n",
    "# Local imports\n",
    "from options import Phase3 as params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "code_folding": [
     0,
     58,
     66
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7fd197e583d0>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "'\\n    Paths and macros\\n'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "'\\n    Model code\\n'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "'\\n    Prepare data\\n'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device('cuda')\n",
    "np.random.seed(42)\n",
    "torch.manual_seed(42)\n",
    "\n",
    "'''\n",
    "    Paths and macros\n",
    "'''\n",
    "# Path fields\n",
    "BOS = 'xbos'  # beginning-of-sentence tag\n",
    "FLD = 'xfld'  # data field tag\n",
    "\n",
    "DATA_PATH = Path('raw/imdb/aclImdb/')\n",
    "DATA_PATH.mkdir(exist_ok=True)\n",
    "PATH = Path('resources/proc/imdb')\n",
    "DATA_PROC_PATH = PATH / 'data'\n",
    "DATA_LM_PATH = PATH / 'datalm'\n",
    "\n",
    "LM_PATH = Path('resources/models')\n",
    "LM_PATH.mkdir(exist_ok=True)\n",
    "PRE_PATH = LM_PATH / 'wt103'\n",
    "PRE_LM_PATH = PRE_PATH / 'fwd_wt103.h5'\n",
    "CLASSES = ['neg', 'pos', 'unsup']\n",
    "\n",
    "\n",
    "'''\n",
    "    Model code\n",
    "'''\n",
    "class CustomEncoder(lm_rnn.MultiBatchRNN):\n",
    "    @property\n",
    "    def layers(self):\n",
    "        # TODO: ADD ENCODERR!!!!!!!!!!\n",
    "        return torch.nn.ModuleList([torch.nn.ModuleList([self.rnns[0], self.dropouths[0]]),\n",
    "                                    torch.nn.ModuleList([self.rnns[1], self.dropouths[1]]),\n",
    "                                    torch.nn.ModuleList([self.rnns[2], self.dropouths[2]])])\n",
    "\n",
    "\n",
    "class TextClassifier(nn.Module):\n",
    "\n",
    "    # @TODO: inject comments.\n",
    "    def __init__(self,\n",
    "                 _device: torch.device,\n",
    "                 ntoken: int,\n",
    "                 dps: list,\n",
    "                 enc_wgts = None,\n",
    "                 _debug=False):\n",
    "        super(TextClassifier, self).__init__()\n",
    "\n",
    "        self.device = _device\n",
    "\n",
    "        # Load the pre-trained model\n",
    "        args = {'ntoken': ntoken, 'emb_sz': 400, 'n_hid': 1150,\n",
    "                'n_layers': 3, 'pad_token': 0, 'qrnn': False, 'bptt': 70, 'max_seq': 1400,\n",
    "                'dropouti': dps[0], 'wdrop': dps[1], 'dropoute': dps[2], 'dropouth': dps[3]}\n",
    "        self.encoder = CustomEncoder(**args).to(self.device)\n",
    "        if enc_wgts:\n",
    "            self.encoder.load_state_dict(enc_wgts)\n",
    "        '''\n",
    "            Make new classifier.\n",
    "            \n",
    "            Explanation:\n",
    "                400*3 because input is [ h_T, maxpool, meanpool ]\n",
    "                50 is hidden layer dim\n",
    "                2 is n_classes\n",
    "\n",
    "                0.4, 0.1 are drops at various layers\n",
    "        '''\n",
    "        self.linear = text.PoolingLinearClassifier(layers=[400 * 3, 50, 2], drops=[dps[4], 0.1]).to(self.device)\n",
    "        self.encoder.reset()\n",
    "\n",
    "    @property\n",
    "    def layers(self):\n",
    "        layers = [x for x in self.encoder.layers]\n",
    "        layers += [x for x in self.linear.layers]\n",
    "        return torch.nn.ModuleList(layers)\n",
    "\n",
    "    @property\n",
    "    def layers_rev(self):\n",
    "        layers = [x for x in self.encoder.layers]\n",
    "        layers += [x for x in self.linear.layers]\n",
    "        layers.reverse()\n",
    "        return torch.nn.ModuleList(layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # inputs are S*B\n",
    "\n",
    "        # Encoding all the data\n",
    "        op_p = self.encoder(x.transpose(1, 0))\n",
    "        # pos_batch = op_p[1][-1][-1]\n",
    "        score = self.linear(op_p)[0]\n",
    "\n",
    "        return score\n",
    "\n",
    "    def predict(self, x):\n",
    "        with torch.no_grad():\n",
    "            self.eval()\n",
    "            predicted = self.forward(x)\n",
    "            self.train()\n",
    "            return predicted\n",
    "\n",
    "\n",
    "'''\n",
    "    Prepare data\n",
    "'''\n",
    "re1 = re.compile(r'  +')\n",
    "\n",
    "\n",
    "def fixup(x):\n",
    "    x = x.replace('#39;', \"'\").replace('amp;', '&').replace('#146;', \"'\").replace(\n",
    "        'nbsp;', ' ').replace('#36;', '$').replace('\\\\n', \"\\n\").replace('quot;', \"'\").replace(\n",
    "        '<br />', \"\\n\").replace('\\\\\"', '\"').replace('<unk>', 'u_n').replace(' @.@ ', '.').replace(\n",
    "        ' @-@ ', '-').replace('\\\\', ' \\\\ ')\n",
    "    return re1.sub(' ', html.unescape(x))\n",
    "\n",
    "\n",
    "def get_texts(df, n_lbls=1):\n",
    "    labels = df.iloc[:, range(n_lbls)].values.astype(np.int64)\n",
    "    texts = f'\\n{BOS} {FLD} 1 ' + df.iloc[:, 1].astype(str)\n",
    "    texts = list(texts.apply(fixup).values)\n",
    "\n",
    "    tok = text.Tokenizer().proc_all_mp(core.partition_by_cores(texts))\n",
    "    return tok, list(labels)\n",
    "\n",
    "\n",
    "def get_all(df, n_lbls):\n",
    "    tok, labels = get_texts(df)\n",
    "    return tok, labels\n",
    "\n",
    "\n",
    "def get_texts_org(path):\n",
    "    texts, labels = [], []\n",
    "    for idx, label in enumerate(CLASSES):\n",
    "        for fname in (path / label).glob('*.*'):\n",
    "            texts.append(fname.open('r', encoding='utf-8').read())\n",
    "            labels.append(idx)\n",
    "    return np.array(texts), np.array(labels)\n",
    "\n",
    "\n",
    "def epoch_end_hook() -> None:\n",
    "    lr_schedule.reset()\n",
    "\n",
    "\n",
    "def eval(y_pred, y_true):\n",
    "    \"\"\"\n",
    "        Expects a batch of input\n",
    "\n",
    "        :param y_pred: tensor of shape (b, nc)\n",
    "        :param y_true: tensor of shape (b, 1)\n",
    "    \"\"\"\n",
    "    return torch.mean((torch.argmax(y_pred, dim=1) == y_true).float())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n    Make model\\n'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "'\\n    Setup things for training (data, loss, opt, lr schedule etc\\n'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "QUICK, DEBUG, MODEL_NUM = False, True, '1'\n",
    "PRETRAINED = False\n",
    "UNSUP_MODEL_DIR = PATH / 'models' / MODEL_NUM\n",
    "\n",
    "trn_texts, trn_labels = get_texts_org(DATA_PATH / 'train')\n",
    "val_texts, val_labels = get_texts_org(DATA_PATH / 'test')\n",
    "\n",
    "# Lose label 2 from train\n",
    "trn_texts = trn_texts[trn_labels<2]\n",
    "trn_labels = trn_labels[trn_labels<2]\n",
    "\n",
    "# Shuffle data\n",
    "if QUICK:\n",
    "    np.random.seed(42)\n",
    "    trn_idx = np.random.permutation(len(trn_texts))[:1000]\n",
    "    val_idx = np.random.permutation(len(val_texts))[:1000]\n",
    "else:\n",
    "    np.random.seed(42)\n",
    "    trn_idx = np.random.permutation(len(trn_texts))\n",
    "    val_idx = np.random.permutation(len(val_texts))\n",
    "\n",
    "trn_texts, trn_labels = trn_texts[trn_idx], trn_labels[trn_idx]\n",
    "val_texts, val_labels = val_texts[val_idx], val_labels[val_idx]\n",
    "col_names = ['labels', 'text']\n",
    "\n",
    "df_trn = pd.DataFrame({'text': trn_texts, 'labels': trn_labels}, columns=col_names)\n",
    "df_val = pd.DataFrame({'text': val_texts, 'labels': val_labels}, columns=col_names)\n",
    "\n",
    "itos_path = UNSUP_MODEL_DIR / 'itos.pkl'\n",
    "itos2 = pickle.load(itos_path.open('rb'))\n",
    "stoi2 = {v: k for k, v in enumerate(itos2)}\n",
    "\n",
    "trn_clas, trn_labels = get_all(df_trn, 1)\n",
    "val_clas, val_labels = get_all(df_val, 1)\n",
    "\n",
    "trn_clas = np.array([[stoi2.get(w, 0) for w in para] for para in trn_clas])\n",
    "val_clas = np.array([[stoi2.get(w, 0) for w in para] for para in val_clas])\n",
    "trn_labels = [x for y in trn_labels for x in y]\n",
    "val_labels = [x for y in val_labels for x in y]\n",
    "\n",
    "'''\n",
    "    Make model\n",
    "'''\n",
    "dps = list(params.encoder_dropouts)\n",
    "# enc_wgts = torch.load(LM_PATH, map_location=lambda storage, loc: storage)\n",
    "enc_wgts = torch.load(UNSUP_MODEL_DIR / 'unsup_model_enc.torch', map_location=lambda storage, loc: storage)\n",
    "clf = TextClassifier(device, len(itos2), dps, enc_wgts=enc_wgts if PRETRAINED else None)\n",
    "\n",
    "'''\n",
    "    Setup things for training (data, loss, opt, lr schedule etc\n",
    "'''\n",
    "bs = params.bs\n",
    "loss_fn = torch.nn.CrossEntropyLoss()\n",
    "opt_fn = partial(optim.Adam, betas=params.adam_betas)\n",
    "opt = make_opt(clf, opt_fn, lr=0.0)\n",
    "opt.param_groups[-1]['lr'] = 0.01\n",
    "\n",
    "# Make data\n",
    "data_fn = partial(mtdi.SortishSampler, _batchsize=bs, _padidx=1)\n",
    "data = {'train': {'x': trn_clas, 'y': trn_labels}, 'valid': {'x': val_clas, 'y': val_labels}}\n",
    "\n",
    "# Make lr scheduler\n",
    "lr_args = {'iterations': len(data_fn(data['train'])), 'cycles': 1}\n",
    "lr_schedule = mtlr.LearningRateScheduler(optimizer=opt, lr_args=lr_args, lr_iterator=mtlr.CosineAnnealingLR)\n",
    "\n",
    "# Re-write saving location (since this is a diff experiment altogether)\n",
    "UNSUP_MODEL_DIR = mt_save_dir(PATH / 'models', _newdir=True)\n",
    "\n",
    "args = {'epochs': 1, 'data': data, 'device': device,\n",
    "        'opt': opt, 'loss_fn': loss_fn, 'model': clf,\n",
    "        'train_fn': clf, 'predict_fn': clf.predict,\n",
    "        'epoch_end_hook': epoch_end_hook, 'weight_decay': params.weight_decay,\n",
    "        'clip_grads_at': params.clip_grads_at, 'lr_schedule': lr_schedule,\n",
    "        'data_fn': data_fn, 'eval_fn': eval}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PosixPath('resources/proc/imdb/models/9')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "UNSUP_MODEL_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n    Training schedule:\\n    \\n    1. Unfreeze one layer. Train for 1 epoch\\n    2 - 5. Unfreeze one layer, train for 1 epoch\\n    3. Train for 15 epochs (after all layers are unfrozen). Use 15 cycles for cosine annealing.\\n'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1042/1042 [03:59<00:00,  4.40it/s]\n",
      "100%|██████████| 1042/1042 [01:32<00:00, 11.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 000 | Loss: 0.69868 | Tr_c: 0.50402 | Vl_c: 0.50004 | Time: 4.022 min\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1042/1042 [03:59<00:00,  4.36it/s]\n",
      "100%|██████████| 1042/1042 [01:32<00:00, 11.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 000 | Loss: 0.69459 | Tr_c: 0.50856 | Vl_c: 0.50004 | Time: 4.012 min\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1042/1042 [03:59<00:00,  4.04it/s]\n",
      "100%|██████████| 1042/1042 [01:31<00:00, 11.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 000 | Loss: 0.68283 | Tr_c: 0.55696 | Vl_c: 0.52991 | Time: 4.009 min\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1042/1042 [03:59<00:00,  6.81it/s]\n",
      "100%|██████████| 1042/1042 [01:31<00:00, 11.39it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 000 | Loss: 0.69521 | Tr_c: 0.49544 | Vl_c: 0.49368 | Time: 4.008 min\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1042/1042 [03:58<00:00,  3.37it/s]\n",
      "100%|██████████| 1042/1042 [01:32<00:00, 11.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 000 | Loss: 0.69455 | Tr_c: 0.50230 | Vl_c: 0.50092 | Time: 4.002 min\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1042/1042 [03:59<00:00,  4.92it/s]\n",
      "100%|██████████| 1042/1042 [01:31<00:00, 11.36it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 000 | Loss: 0.69457 | Tr_c: 0.49872 | Vl_c: 0.50036 | Time: 4.008 min\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1042/1042 [03:59<00:00,  3.41it/s]\n",
      "100%|██████████| 1042/1042 [01:31<00:00,  9.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001 | Loss: 0.69501 | Tr_c: 0.49902 | Vl_c: 0.49996 | Time: 4.009 min\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1042/1042 [03:59<00:00,  4.18it/s]\n",
      "100%|██████████| 1042/1042 [01:32<00:00, 11.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 002 | Loss: 0.69462 | Tr_c: 0.50054 | Vl_c: 0.50036 | Time: 4.006 min\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 24%|██▍       | 255/1042 [00:58<02:36,  5.02it/s]"
     ]
    }
   ],
   "source": [
    "'''\n",
    "    Training schedule:\n",
    "    \n",
    "    1. Unfreeze one layer. Train for 1 epoch\n",
    "    2 - 5. Unfreeze one layer, train for 1 epoch\n",
    "    3. Train for 15 epochs (after all layers are unfrozen). Use 15 cycles for cosine annealing.\n",
    "'''\n",
    "# opt.param_groups[-1]['lr'] = 0.01\n",
    "traces = loops.generic_loop(**args)\n",
    "\n",
    "opt.param_groups[-1]['lr'] = 0.01\n",
    "opt.param_groups[-2]['lr'] = 0.005\n",
    "lr_schedule = mtlr.LearningRateScheduler(optimizer=opt, lr_args=lr_args, lr_iterator=mtlr.CosineAnnealingLR)\n",
    "args['lr_schedule'] = lr_schedule\n",
    "traces_new = loops.generic_loop(**args)\n",
    "traces = [a+b for a, b in zip(traces, traces_new)]\n",
    "\n",
    "opt.param_groups[-1]['lr'] = 0.01\n",
    "opt.param_groups[-2]['lr'] = 0.005\n",
    "opt.param_groups[-3]['lr'] = 0.001\n",
    "lr_schedule = mtlr.LearningRateScheduler(optimizer=opt, lr_args=lr_args, lr_iterator=mtlr.CosineAnnealingLR)\n",
    "args['lr_schedule'] = lr_schedule\n",
    "traces_new = loops.generic_loop(**args)\n",
    "traces = [a+b for a, b in zip(traces, traces_new)]\n",
    "\n",
    "opt.param_groups[-1]['lr'] = 0.01\n",
    "opt.param_groups[-2]['lr'] = 0.005\n",
    "opt.param_groups[-3]['lr'] = 0.001\n",
    "opt.param_groups[-4]['lr'] = 0.001\n",
    "lr_schedule = mtlr.LearningRateScheduler(optimizer=opt, lr_args=lr_args, lr_iterator=mtlr.CosineAnnealingLR)\n",
    "args['lr_schedule'] = lr_schedule\n",
    "traces_new = loops.generic_loop(**args)\n",
    "traces = [a+b for a, b in zip(traces, traces_new)]\n",
    "\n",
    "opt.param_groups[-1]['lr'] = 0.01\n",
    "opt.param_groups[-2]['lr'] = 0.005\n",
    "opt.param_groups[-3]['lr'] = 0.001\n",
    "opt.param_groups[-4]['lr'] = 0.001\n",
    "opt.param_groups[-5]['lr'] = 0.001\n",
    "lr_schedule = mtlr.LearningRateScheduler(optimizer=opt, lr_args=lr_args, lr_iterator=mtlr.CosineAnnealingLR)\n",
    "args['lr_schedule'] = lr_schedule\n",
    "traces_new = loops.generic_loop(**args)\n",
    "traces = [a+b for a, b in zip(traces, traces_new)]\n",
    "\n",
    "opt.param_groups[-1]['lr'] = 0.01\n",
    "opt.param_groups[-2]['lr'] = 0.005\n",
    "opt.param_groups[-3]['lr'] = 0.001\n",
    "opt.param_groups[-4]['lr'] = 0.001\n",
    "opt.param_groups[-5]['lr'] = 0.001\n",
    "lr_args['cycles'] = 15\n",
    "args['epochs'] = 15\n",
    "lr_schedule = mtlr.LearningRateScheduler(optimizer=opt, lr_args=lr_args, lr_iterator=mtlr.CosineAnnealingLR)\n",
    "args['lr_schedule'] = lr_schedule\n",
    "traces_new = loops.generic_loop(**args)\n",
    "traces = [a+b for a, b in zip(traces, traces_new)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mt_save(UNSUP_MODEL_DIR,\n",
    "            torch_stuff=[tosave('sup_model.torch', clf.state_dict())],\n",
    "            pickle_stuff=[tosave('final_sup_traces.pkl', traces), tosave('unsup_options.pkl', params)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "UNSUP_MODEL_DIR"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "def plot(trcs):\n",
    "    layers = len(trcs[0])\n",
    "    for l in range(layers):\n",
    "        plt.plot(trcs[:,l], label=f\"layer {l}\")\n",
    "    plt.show()\n",
    "\n",
    "opt_fn = partial(optim.Adam, betas=(0.7, 0.99))\n",
    "opt = make_opt(clf, opt_fn, lr=0.01)\n",
    "opt.param_groups[-1]['lr'] = 0.01\n",
    "# lr_args = {'iterations': len(data_fn(data['train'])), 'cycles': 1}\n",
    "lr_args = {'length': len(data_fn(data['train']))}\n",
    "lr_schedule = mtlr.LearningRateScheduler(opt, lr_args, mtlr.ConstantLR)\n",
    "lrss = [lr_schedule.get() for _ in range(lr_args['length'])]\n",
    "plot(np.asarray(lrss))"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "lr_args = {'iterations': len(data_fn(data['train'])), 'cycles': 1}\n",
    "lr_schedule = mtlr.LearningRateScheduler(opt, lr_args, mtlr.CosineAnnealingLR)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "def plot(trcs):\n",
    "    layers = len(trcs[0])\n",
    "    for l in range(layers):\n",
    "        plt.plot(trcs[:,l], label=f\"layer {l}\")\n",
    "    plt.show()\n",
    "    \n",
    "plot(np.asarray(traces[-1][100:]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(np.asarray([[x[-1]] for x in traces[-1][:]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(lr_args)\n",
    "lr_args['cycles'] = 5\n",
    "lr_args['iterations'] = 42*15\n",
    "lr_schedule = mtlr.LearningRateScheduler(optimizer=opt, lr_args=lr_args, lr_iterator=mtlr.CosineAnnealingLR)\n",
    "lrs = []\n",
    "while True:\n",
    "    try:\n",
    "        lrs.append(lr_schedule.get())\n",
    "    except CustomError:\n",
    "        break\n",
    "plot(np.asarray(lrs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}