1.0--tfmr.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prelims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from fastai import *\n",
    "from fastai.vision import *\n",
    "from fastai.text import *\n",
    "\n",
    "import pdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "PATH = Path('data/IAM_handwriting')\n",
    "TMP_PATH = PATH/'tmp'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# # Quick test of export.pkl\n",
    "# from fastai.tfmr_extensions import *\n",
    "\n",
    "# learn = load_learner(PATH)\n",
    "\n",
    "# fname = PATH/'test/edit_test4.png'\n",
    "# im = open_image(fname)\n",
    "\n",
    "# seq,res,preds = learn.predict(im)\n",
    "# print(seq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss and Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def loss_prep(input, target):\n",
    "    \"equalize input/target sl; combine bs/sl dimensions\"\n",
    "    bs,tsl = target.shape\n",
    "    _ ,sl,vocab = input.shape\n",
    "        \n",
    "    # F.pad( front,back for dimensions: 1,0,2 )\n",
    "    if sl>tsl: target = F.pad(target, (0,sl-tsl))\n",
    "        \n",
    "    # this should only be used when testing for small seq_lens\n",
    "    # if tsl>sl: target = target[:,:sl]\n",
    "    \n",
    "    if tsl>sl: input = F.pad(input, (0,0,0,tsl-sl))\n",
    "    # not ideal => adds 82 logits all 0s...\n",
    "        \n",
    "    targ = target.contiguous().view(-1).long()\n",
    "    pred = input.contiguous().view(-1, vocab)\n",
    "    return pred, targ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LabelSmoothing(nn.Module):\n",
    "    def __init__(self, smoothing=0.1):\n",
    "        super(LabelSmoothing, self).__init__()\n",
    "        self.confidence = 1.0 - smoothing\n",
    "        self.smoothing = smoothing\n",
    "        \n",
    "    def forward(self, pred, target):\n",
    "        pred,targ = loss_prep(pred, target)\n",
    "        pred = F.log_softmax(pred, dim=-1)  # need this for KLDivLoss\n",
    "        true_dist = pred.data.clone()\n",
    "        true_dist.fill_(self.smoothing / pred.size(1))                  # fill with 0.0012\n",
    "        true_dist.scatter_(1, targ.data.unsqueeze(1), self.confidence)  # [0.0012, 0.0012, 0.90, 0.0012]\n",
    "        return F.kl_div(pred, true_dist, reduction='sum')/bs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import Levenshtein as Lev\n",
    "\n",
    "class CER(Callback):\n",
    "    def __init__(self, itos):\n",
    "        super().__init__()\n",
    "        self.name = 'cer'\n",
    "        self.itos = itos\n",
    "\n",
    "    def on_epoch_begin(self, **kwargs):\n",
    "        self.errors, self.total = 0, 0\n",
    "    \n",
    "    def on_batch_end(self, last_output, last_target, **kwargs):\n",
    "        error,size = self._cer(last_output, last_target)\n",
    "        self.errors += error\n",
    "        self.total += size\n",
    "    \n",
    "    def on_epoch_end(self, last_metrics, **kwargs):\n",
    "        return add_metrics(last_metrics, self.errors/self.total)\n",
    "\n",
    "    def _cer(self, preds, targs):\n",
    "        bs,sl = targs.size()\n",
    "        \n",
    "        res = torch.argmax(preds, dim=2)\n",
    "        error = 0\n",
    "        for i in range(bs):\n",
    "            p = self._char_label_text(res[i])   #.replace(' ', '')\n",
    "            t = self._char_label_text(targs[i]) #.replace(' ', '')\n",
    "            error += Lev.distance(t, p)/len(t)\n",
    "        return error, bs\n",
    "\n",
    "    def _char_label_text(self, pred):\n",
    "#         return self.sp.DecodeIds(pred.tolist())\n",
    "        ints = to_np(pred).astype(int)\n",
    "        nonzero = ints[np.nonzero(ints)]\n",
    "        return ''.join([self.itos[i] for i in nonzero])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def rshift(tgt, bos_token=1):\n",
    "    \"Shift y to the right by prepending token\"\n",
    "    bos = torch.zeros((tgt.size(0),1), device=device).type_as(tgt) + bos_token\n",
    "    return torch.cat((bos, tgt[:,:-1]), dim=-1)\n",
    "\n",
    "def subsequent_mask(size):\n",
    "    return torch.tril(torch.ones((size,size), device=device).byte()).unsqueeze(0)\n",
    "\n",
    "class TeacherForce(LearnerCallback):\n",
    "    def __init__(self, learn:Learner, bos_token:int=1):\n",
    "        super().__init__(learn)\n",
    "        self.bos_token = bos_token\n",
    "        \n",
    "    def on_batch_begin(self, last_input, last_target, **kwargs):\n",
    "        s = rshift(last_target, self.bos_token).long()\n",
    "        mask = subsequent_mask(s.size(-1))\n",
    "        return {'last_input':(last_input, s, mask), 'last_target':last_target}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## synth dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fname = 'edited_sm_synth.csv' #'small_synth_words.csv'\n",
    "CSV = PATH/fname\n",
    "FOLDER = 'edited_sm_synth'\n",
    "\n",
    "df = pd.read_csv(CSV)\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sz,bs = 128,100\n",
    "# sz,bs = 256,100\n",
    "\n",
    "seq_len = 45"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Concat Lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "nums = '3-6'   #'7-10'  #'11-14'\n",
    "CSV = PATH/f'cat_lines_{nums}.csv'\n",
    "FOLDER = 'resized_cat_lines'\n",
    "\n",
    "csv = pd.read_csv(CSV)\n",
    "test = pd.read_csv(PATH/'test_pg.csv')\n",
    "\n",
    "len(csv), len(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lengths = np.array([len(i.split(' ')) for i in csv.char_ids.values])\n",
    "lengths.max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sz,bs = 512,20   #1000,5  #800,8   #512,20\n",
    "seq_len = 600 #600 #450  #300\n",
    "stats = (np.array([0.941, 0.941, 0.941], dtype=np.float32), np.array([0.128, 0.128, 0.128], dtype=np.float32))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Concat All"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "a = pd.read_csv(PATH/'cat_lines_3.csv').sample(1000)\n",
    "b = pd.read_csv(PATH/'cat_lines_4.csv').sample(1000)\n",
    "c = pd.read_csv(PATH/'cat_lines_5.csv').sample(1000)\n",
    "d = pd.read_csv(PATH/'cat_lines_6.csv').sample(1000)\n",
    "e = pd.read_csv(PATH/'cat_lines_7.csv').sample(1000)\n",
    "f = pd.read_csv(PATH/'cat_lines_8.csv').sample(1000)\n",
    "g = pd.read_csv(PATH/'cat_lines_9.csv').sample(1000)\n",
    "h = pd.read_csv(PATH/'cat_lines_10.csv').sample(1000)\n",
    "i = pd.read_csv(PATH/'cat_lines_11.csv').sample(1000)\n",
    "j = pd.read_csv(PATH/'cat_lines_12.csv').sample(1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "k = pd.read_csv(PATH/'paragraph_chars.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "val_idxs = np.array(k.sample(frac=0.15, random_state=42).index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "trn = k[~k.index.isin(val_idxs)]\n",
    "test = k[k.index.isin(val_idxs)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "new = pd.concat([a,b,c,d,e,f,g,h,i,j,trn], ignore_index=True)\n",
    "len(new)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# new.to_csv(PATH/'cat_lines_pg.csv', index=False)\n",
    "# test.to_csv(PATH/'test_pg.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Paragraphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fname = 'edited_pg.csv' #'paragraphs.csv'\n",
    "CSV = PATH/fname\n",
    "FOLDER = 'paragraphs'\n",
    "\n",
    "df = pd.read_csv(CSV)\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# sz,bs = 1000,5  #1024,5  #1000,5   #~2000x1000 full size\n",
    "# sz,bs = 800,8\n",
    "sz,bs = 512,10\n",
    "\n",
    "seq_len = 700   #~400 chars/paragraph - max: 705\n",
    "# stats = (np.array([0.941, 0.941, 0.941], dtype=np.float32), np.array([0.128, 0.128, 0.128], dtype=np.float32))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Downloaded Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "CSV = PATH/'downloaded_images.csv'\n",
    "FOLDER = 'downloaded_images'\n",
    "\n",
    "csv = pd.read_csv(CSV)\n",
    "len(csv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# csv['filename'] = csv['filename'].apply(lambda x: f\"dl_{x}\")\n",
    "# csv.head()\n",
    "# csv.to_csv(CSV, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# sz,bs = 1000,5  #1024,5  #1000,5   #~2000x1000 full size\n",
    "# sz,bs = 800,8\n",
    "sz,bs = 512,20\n",
    "\n",
    "seq_len = 700   #~400 chars/paragraph - max: 705\n",
    "stats = (np.array([0.941, 0.941, 0.941], dtype=np.float32), np.array([0.128, 0.128, 0.128], dtype=np.float32))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fname = 'test_pg.csv'\n",
    "CSV = PATH/fname\n",
    "FOLDER = 'paragraphs'\n",
    "\n",
    "df = pd.read_csv(CSV)\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sz,bs = 512,15\n",
    "\n",
    "seq_len = 700"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Mix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fname = 'mix.csv'\n",
    "CSV = PATH/fname\n",
    "FOLDER = 'mix'\n",
    "\n",
    "df = pd.read_csv(CSV)\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sz,bs = 512,15\n",
    "# sz,bs = 400,20\n",
    "\n",
    "seq_len = 800"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ModelData"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "tfms = get_transforms(do_flip=False, max_zoom=1, max_rotate=2.0, max_warp=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def label_collater(samples:BatchSamples, pad_idx:int=0):\n",
    "    \"Function that collect samples and pads end of labels.\"\n",
    "    data = to_data(samples)\n",
    "    ims, lbls = zip(*data)\n",
    "    imgs = torch.stack(list(ims))\n",
    "    if len(data) is 1:\n",
    "        labels = torch.zeros(1,1).long()\n",
    "        return imgs, labels\n",
    "    max_len = max([len(s) for s in lbls])\n",
    "    labels = torch.zeros(len(data), max_len).long() + pad_idx\n",
    "    for i,lbl in enumerate(lbls):\n",
    "        labels[i,:len(lbl)] = torch.from_numpy(lbl)  #padding end    \n",
    "    return imgs, labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### SentencePiece"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SPMProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None):\n",
    "        self.sp = ds.sp if ds is not None else None\n",
    "\n",
    "    def process_one(self,item): return self.sp.EncodeAsIds(item)\n",
    "    def process(self, ds): super().process(ds)\n",
    "    \n",
    "class SPMList(ItemList):\n",
    "    _processor = [SPMProcessor]\n",
    "\n",
    "    def __init__(self, items:Iterator, sp_processor, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.sp = sp_processor\n",
    "        self.copy_new += ['sp']\n",
    "\n",
    "    def get(self, i):\n",
    "        o = super().get(i)\n",
    "        return Text(o, self.sp.DecodeIds(o))\n",
    "\n",
    "    def reconstruct(self, t:Tensor):\n",
    "        return Text(t, self.sp.DecodeIds(t.tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import sentencepiece as spm\n",
    "\n",
    "sp = spm.SentencePieceProcessor()\n",
    "sp.Load(str(PATH/'spm_train.model'))\n",
    "sp.SetEncodeExtraOptions(\"bos:eos\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageList.from_df(df, path=PATH, folder=FOLDER)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=SPMList, sp_processor=sp)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=label_collater)\n",
    "        .normalize()\n",
    "       )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Chars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "itos = pickle.load(open(PATH/'itos.pkl', 'rb'))  #'char_itos.pkl'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class CharTokenizer():\n",
    "    def __init__(self, n_cpus:int=None):\n",
    "        self.n_cpus = ifnone(n_cpus, defaults.cpus)\n",
    "\n",
    "    def tokenize(self, t:str): return ['xxbos']+list(t)+['xxeos']\n",
    "\n",
    "    def _process_all_1(self, texts:Collection[str]) -> List[List[str]]:\n",
    "        \"Process a list of `texts` in one process.\"\n",
    "        return [self.tokenize(str(t)) for t in texts]\n",
    "\n",
    "    def process_all(self, texts:Collection[str]) -> List[List[str]]:\n",
    "        \"Process a list of `texts`.\"\n",
    "        if self.n_cpus <= 1: return self._process_all_1(texts)\n",
    "        with ProcessPoolExecutor(self.n_cpus) as e:\n",
    "            return sum(e.map(self._process_all_1, partition_by_cores(texts, self.n_cpus)), [])\n",
    "\n",
    "class CharVocab(Vocab):\n",
    "    def __init__(self, itos:Collection[str]):\n",
    "        self.itos = itos\n",
    "        self.stoi = collections.defaultdict(lambda: 3, {v:k for k,v in enumerate(self.itos)})\n",
    "\n",
    "    def textify(self, nums:Collection[int], sep=''):\n",
    "        nums = nums[1:-1]  #remove bos,eos tokens\n",
    "        return sep.join([self.itos[i] for i in nums]) if sep is not None else [self.itos[i] for i in nums]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "vocab = CharVocab(itos)\n",
    "procs = [TokenizeProcessor(tokenizer=CharTokenizer(), include_bos=False),\n",
    "         NumericalizeProcessor(vocab=vocab)]\n",
    "\n",
    "data = (ImageList.from_df(df, path=PATH, folder=FOLDER)\n",
    "        .split_none()\n",
    "        #.split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=TextList, sep='', pad_idx=0, vocab=vocab, processor=procs)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=label_collater)\n",
    "        .normalize()\n",
    "       )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### SequenceList (old)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# df['lens'] = df.label.map(len)\n",
    "# df.sort_values('lens', ascending=False, inplace=True)\n",
    "# df.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "itos = pickle.load(open(PATH/'itos.pkl', 'rb'))  #'char_itos.pkl'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SequenceItem(ItemBase):\n",
    "    def __init__(self,data,vocab): self.data,self.vocab = data,vocab        \n",
    "    def __str__(self): return self.textify(self.data)\n",
    "    def __hash__(self): return hash(str(self))\n",
    "    def textify(self, data): return ''.join([self.vocab[i] for i in data[:-1]])\n",
    "        \n",
    "class ArrayProcessor(PreProcessor):\n",
    "    \"Convert df column (string of ints) into np.array\"\n",
    "    def __init__(self, ds:ItemList=None): None\n",
    "    def process_one(self,item): return np.array(item.split(), dtype=np.int64)\n",
    "    def process(self, ds): super().process(ds)\n",
    "        \n",
    "class ItosProcessor(PreProcessor):\n",
    "    def __init__(self, ds:ItemList=None): self.itos = ds.itos\n",
    "    def process(self, ds:ItemList): ds.itos = self.itos\n",
    "        \n",
    "class SequenceList(ItemList):\n",
    "    _processor = [ItosProcessor, ArrayProcessor]\n",
    "    \n",
    "    def __init__(self, items:Iterator, itos:List[str]=None, **kwargs):\n",
    "        super().__init__(items, **kwargs)\n",
    "        self.itos = itos\n",
    "        self.copy_new += ['itos']\n",
    "        self.c = len(self.items)\n",
    "\n",
    "    def get(self, i):\n",
    "        o = super().get(i)\n",
    "        return SequenceItem(o, self.itos)\n",
    "\n",
    "    def reconstruct(self,t):\n",
    "        # Converting padded tensor back into np.array\n",
    "        o = t.numpy()\n",
    "        o = o[np.nonzero(o)]                  # remove 0 padding\n",
    "        return SequenceItem(o, self.itos)\n",
    "    \n",
    "    def analyze_pred(self,pred):\n",
    "        return torch.argmax(pred, dim=-1)\n",
    "        # method called in learn.predict() or learn.show_results()\n",
    "        # to transform predictions in an output tensor suitable for reconstruct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageList.from_df(df, path=PATH, folder=FOLDER)\n",
    "        .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "        .label_from_df(label_cls=SequenceList, itos=itos)\n",
    "        .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=custom_collater)\n",
    "        .normalize()\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# # This slows down training...\n",
    "# class CustomSampler(Sampler):\n",
    "#     \"sort dataset by longest y sequences\"\n",
    "#     def __init__(self, dataset):\n",
    "#         self.dataset = dataset\n",
    "#         self.lengths = [len(i) for i in dataset.y.items]\n",
    "#         self.sorted_idxs = np.flip(np.argsort(self.lengths))\n",
    "#     def __len__(self): return len(self.dataset)\n",
    "#     def __iter__(self): return iter(self.sorted_idxs)\n",
    "\n",
    "# ds = data.train_ds\n",
    "# tfms = data.train_dl.tfms\n",
    "# sampler = CustomSampler(ds)\n",
    "# dl = DataLoader(ds, bs, shuffle=False, sampler=sampler, num_workers=num_cpus(), collate_fn=custom_collater, drop_last=True)\n",
    "# ddl = DeviceDataLoader(dl, device, tfms, custom_collater)\n",
    "# data.train_dl = ddl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data.batch_stats()\n",
    "# no normalization: [tensor([0.9403, 0.9403, 0.9403]), tensor([0.1604, 0.1604, 0.1604])]\n",
    "# imagenet_stats:   [tensor([1.9973, 2.1714, 2.3839]), tensor([0.6974, 0.7130, 0.7098])]\n",
    "# normalize():      [tensor([0.0225, 0.0225, 0.0225]), tensor([0.9676, 0.9676, 0.9676])]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=2, ds_type=DatasetType.Train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Transformer Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "LayerNorm = partial(nn.LayerNorm, eps=1e-4)  # accomodates mixed precision training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SublayerConnection(nn.Module):\n",
    "    \"A residual connection followed by a layer norm.  Note: (for code simplicity) norm is first.\"\n",
    "    def __init__(self, size, dropout):\n",
    "        super(SublayerConnection, self).__init__()\n",
    "        self.norm = LayerNorm(size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, sublayer):\n",
    "        return x + self.dropout(sublayer(self.norm(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def clones(module, N):\n",
    "    \"Produce N identical layers.\"\n",
    "    return nn.ModuleList([deepcopy(module) for _ in range(N)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, layer, N):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.layers = clones(layer, N)\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    \"Encoder: self-attn and feed forward\"\n",
    "    def __init__(self, size, self_attn, feed_forward, dropout):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.size = size\n",
    "        self.self_attn = self_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 2)\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))\n",
    "        return self.sublayer[1](x, self.feed_forward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    \"Generic N layer decoder with masking.\"\n",
    "    def __init__(self, layer, N):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.layers = clones(layer, N)\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "        \n",
    "    def forward(self, x, src, tgt_mask=None):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, src, tgt_mask)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    \"Decoder: self-attn, src-attn, and feed forward\"\n",
    "    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.size = size\n",
    "        self.self_attn = self_attn\n",
    "        self.src_attn = src_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 3)  # wraps layer in residual,dropout,norm\n",
    " \n",
    "    def forward(self, x, src, tgt_mask=None):\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))\n",
    "        x = self.sublayer[1](x, lambda x: self.src_attn(x, src, src))\n",
    "        return self.sublayer[2](x, self.feed_forward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def attention(query, key, value, mask=None, dropout=None):\n",
    "    \"Compute 'Scaled Dot Product Attention'\"\n",
    "    depth = query.size(-1)\n",
    "    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(depth)\n",
    "    if mask is not None:\n",
    "        scores = scores.masked_fill(mask == 0, -1e4)  #changed from: -1e9 to accomodate mixed precision  \n",
    "    p_attn = F.softmax(scores, dim=-1)\n",
    "    if dropout is not None:\n",
    "        p_attn = dropout(p_attn)\n",
    "    return torch.matmul(p_attn, value), p_attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SingleHeadedAttention(nn.Module):\n",
    "    def __init__(self, d_model, dropout=0.2):\n",
    "        super(SingleHeadedAttention, self).__init__()\n",
    "        self.linears = clones(nn.Linear(d_model, d_model), 4)\n",
    "        self.attn = None\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, query, key, value, mask=None):        \n",
    "        query, key, value = [l(x) for l, x in zip(self.linears, (query, key, value))]\n",
    "        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)\n",
    "        return self.linears[-1](x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiHeadedAttention(nn.Module):\n",
    "    def __init__(self, d_model, h=8, dropout=0.2):\n",
    "        super(MultiHeadedAttention, self).__init__()\n",
    "        assert d_model % h == 0\n",
    "        self.d_k = d_model // h        # assume d_v always equals d_k\n",
    "        self.h = h\n",
    "        self.linears = clones(nn.Linear(d_model, d_model), 4)\n",
    "        self.attn = None\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, q, k, v, mask=None):\n",
    "        if mask is not None: mask = mask.unsqueeze(1)\n",
    "        bs = q.size(0)\n",
    "        \n",
    "        # 1) Do all the linear projections in batch from d_model => h x d_k \n",
    "        q, k, v = [l(x).view(bs, -1, self.h, self.d_k).transpose(1,2) for l, x in zip(self.linears, (q, k, v))]\n",
    "        \n",
    "        # 2) Apply attention on all the projected vectors in batch. \n",
    "        x, self.attn = attention(q, k, v, mask=mask, dropout=self.dropout)\n",
    "        \n",
    "        # 3) \"Concat\" using a view and apply a final linear. \n",
    "        x = x.transpose(1, 2).contiguous().view(bs, -1, self.h * self.d_k)\n",
    "        return self.linears[-1](x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class GeLU(nn.Module):\n",
    "    def forward(self, x): return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))\n",
    "\n",
    "class PositionwiseFeedForward(nn.Module):\n",
    "    def __init__(self, d_model, dropout=0.2):\n",
    "        super(PositionwiseFeedForward, self).__init__()\n",
    "        self.w_1 = nn.Linear(d_model, d_model*4)\n",
    "        self.w_2 = nn.Linear(d_model*4, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.activation = GeLU() #nn.ReLU(inplace=True)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.w_2(self.dropout(self.activation(self.w_1(x))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, dropout=0.1, max_len=2000):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        \n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        position = torch.arange(0.0, max_len).unsqueeze(1)\n",
    "        log_increment = math.log(1e4) / d_model\n",
    "        div_term = torch.exp(torch.arange(0.0, d_model, 2) * -log_increment)  \n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe.unsqueeze_(0)\n",
    "\n",
    "        self.register_buffer('pe', pe)    #(1,max_len,d_model)\n",
    "        # registered buffers are Tensors (not Variables)\n",
    "        # not a parameter but still want in the state_dict\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.pe[:, :x.size(1)]\n",
    "        return self.dropout(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Embeddings(nn.Module):\n",
    "    def __init__(self, d_model, vocab):\n",
    "        super(Embeddings, self).__init__()\n",
    "        self.lut = nn.Embedding(vocab, d_model)\n",
    "        self.d_model = d_model\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.lut(x) * math.sqrt(self.d_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Full Arch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, src, tgt, tgt_mask=None):\n",
    "        return self.decode(self.encode(src), tgt, tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(src)\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(self.tgt_embed(tgt), src, tgt_mask)\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz, d_model):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.linear = nn.Linear(em_sz, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x).flatten(2,3).permute(0,2,1)\n",
    "        x = self.linear(x) * 8\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, N=4, drops=0.2):\n",
    "    c = deepcopy\n",
    "    attn = SingleHeadedAttention(d_model)\n",
    "#     attn = MultiHeadedAttention(d_model, 4)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), PositionalEncoding(d_model, drops, 2000)\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer, seq_len=500):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        self.seq_len = seq_len\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None):\n",
    "        # inference (greedy decode)\n",
    "        if tgt is None:\n",
    "            with torch.no_grad():\n",
    "                feats = self.transformer.encode(self.img_enc(src))\n",
    "                bs = src.size(0)\n",
    "                tgt = torch.ones((bs,1), dtype=torch.long, device=device)\n",
    "\n",
    "                res = []\n",
    "                for i in progress_bar(range(self.seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "                    dec_outs = self.transformer.decode(feats, tgt, mask)\n",
    "                    prob = self.transformer.generate(dec_outs[:,-1])\n",
    "                    res.append(prob)\n",
    "                    pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "                    if (pred==0).all(): break\n",
    "                    tgt = torch.cat([tgt,pred], dim=-1)\n",
    "                out = torch.stack(res).transpose(1,0).contiguous()\n",
    "                \n",
    "        #training\n",
    "        else:\n",
    "            feats = self.img_enc(src)\n",
    "            dec_outs = self.transformer(feats, tgt, tgt_mask)    # ([bs, sl, d_model])\n",
    "            out = self.transformer.generate(dec_outs)            # ([bs, sl, vocab])\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "d_model = 512\n",
    "em_sz = 256\n",
    "img_encoder = ResnetBase(em_sz, d_model)\n",
    "transformer = make_full_model(len(itos), d_model)  #len(itos)\n",
    "net = Img2Seq(img_encoder, transformer, seq_len)\n",
    "\n",
    "# AdamW16 = partial(optim.Adam, betas=(0.9,0.99), eps=1e-4)  ->  #.to_fp16(max_scale=256)\n",
    "# AMSGrad = partial(optim.Adam, amsgrad=True, eps=1e-4)\n",
    "# partial: way to always call a function with a given set of arguments or keywords\n",
    "\n",
    "learn = Learner(data, net, loss_func=LabelSmoothing(smoothing=0.1),\n",
    "                metrics=[CER(itos)], callback_fns=TeacherForce)\n",
    "learn.clip_grad(0.25)\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# MobileNet Arch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, src, tgt, tgt_mask=None):\n",
    "        return self.decode(self.encode(src), tgt, tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(src)\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(self.tgt_embed(tgt), src, tgt_mask)\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from torchvision.models import mobilenet_v2, densenet201, densenet121"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class CNNBase(nn.Module):\n",
    "    def __init__(self, d_model):\n",
    "        super().__init__()\n",
    "\n",
    "        net = densenet201(True)\n",
    "        modules = list(net.children())[0][:-3]\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.linear = nn.Linear(1792, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x).flatten(2,3).permute(0,2,1)\n",
    "        x = self.linear(x) * 2 #8\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, N=4, drops=0.2):\n",
    "    c = deepcopy\n",
    "    attn = SingleHeadedAttention(d_model)\n",
    "#     attn = MultiHeadedAttention(d_model, 4)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), PositionalEncoding(d_model, drops, 2000)\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer, seq_len=500):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        self.seq_len = seq_len\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None):\n",
    "        # inference (greedy decode)\n",
    "        if tgt is None:\n",
    "            with torch.no_grad():\n",
    "                feats = self.transformer.encode(self.img_enc(src))\n",
    "                bs = src.size(0)\n",
    "                tgt = torch.ones((bs,1), dtype=torch.long, device=device)\n",
    "\n",
    "                res = []\n",
    "                for i in progress_bar(range(self.seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "                    dec_outs = self.transformer.decode(feats, tgt, mask)\n",
    "                    prob = self.transformer.generate(dec_outs[:,-1])\n",
    "                    res.append(prob)\n",
    "                    pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "                    if (pred==0).all(): break\n",
    "                    tgt = torch.cat([tgt,pred], dim=-1)\n",
    "                out = torch.stack(res).transpose(1,0).contiguous()\n",
    "                \n",
    "        #training\n",
    "        else:\n",
    "            feats = self.img_enc(src)\n",
    "            dec_outs = self.transformer(feats, tgt, tgt_mask)    # ([bs, sl, d_model])\n",
    "            out = self.transformer.generate(dec_outs)            # ([bs, sl, vocab])\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "d_model = 512\n",
    "\n",
    "img_encoder = CNNBase(d_model)\n",
    "transformer = make_full_model(len(itos), d_model)\n",
    "net = Img2Seq(img_encoder, transformer, seq_len)\n",
    "\n",
    "# AdamW16 = partial(optim.Adam, betas=(0.9,0.99), eps=1e-4)  ->  #.to_fp16(max_scale=256)\n",
    "# AMSGrad = partial(optim.Adam, amsgrad=True, eps=1e-4)\n",
    "# partial: way to always call a function with a given set of arguments or keywords\n",
    "\n",
    "learn = Learner(data, net, loss_func=LabelSmoothing(smoothing=0.1),\n",
    "                metrics=[CER(itos)], callback_fns=TeacherForce)\n",
    "learn.clip_grad(0.25)\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# LM Arch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, src, tgt, tgt_mask=None):\n",
    "        return self.decode(self.encode(src), tgt, tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(src)\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(self.tgt_embed(tgt), src, tgt_mask)\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ResnetBase(nn.Module):\n",
    "    def __init__(self, em_sz, d_model):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:s]\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.linear = nn.Linear(em_sz, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x).flatten(2,3).permute(0,2,1)\n",
    "        x = self.linear(x) * 8\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, N=4, drops=0.2):\n",
    "    c = deepcopy\n",
    "    attn = SingleHeadedAttention(d_model)\n",
    "#     attn = MultiHeadedAttention(d_model, 4)\n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drops), N),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), PositionalEncoding(d_model, drops, 2000)\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class LM(nn.Module):\n",
    "    def __init__(self, vocab, d_model, n_hidden, n_layers):\n",
    "        super(LM, self).__init__()\n",
    "        self.lm = AWD_LSTM(vocab, d_model, n_hidden, n_layers, pad_token=0,\n",
    "                           hidden_p=0.1, input_p=0.25, embed_p=0.05, weight_p=0.2)\n",
    "        self.mixer = nn.Linear(d_model, d_model)\n",
    "        \n",
    "    def forward(self, outs, tgts, decode=False):\n",
    "        res = self.lm(tgts)\n",
    "        lm_outs = res[0][-1]\n",
    "        if decode: lm_outs = lm_outs[:,-1]\n",
    "        return self.mixer(lm_outs+outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer, lm, seq_len=500):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        self.lm = lm\n",
    "        self.seq_len = seq_len\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None):\n",
    "        # inference (greedy decode)\n",
    "        if tgt is None:\n",
    "            with torch.no_grad():\n",
    "                feats = self.transformer.encode(self.img_enc(src))\n",
    "                bs = src.size(0)\n",
    "                tgt = torch.ones((bs,1), dtype=torch.long, device=device)\n",
    "\n",
    "                res = []\n",
    "                for i in progress_bar(range(self.seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "                    dec_outs = self.transformer.decode(feats, tgt, mask)\n",
    "                    outs = self.lm(dec_outs[:,-1], tgt, decode=True)\n",
    "                    prob = self.transformer.generate(outs)\n",
    "                    res.append(prob)\n",
    "                    pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "                    if (pred==0).all(): break\n",
    "                    tgt = torch.cat([tgt,pred], dim=-1)\n",
    "                out = torch.stack(res).transpose(1,0).contiguous()\n",
    "                \n",
    "        #training\n",
    "        else:\n",
    "            feats = self.img_enc(src)\n",
    "            dec_outs = self.transformer(feats, tgt, tgt_mask)\n",
    "            outs = self.lm(dec_outs, tgt)\n",
    "            out = self.transformer.generate(outs)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "d_model = 512\n",
    "em_sz = 256\n",
    "img_encoder = ResnetBase(em_sz, d_model)\n",
    "transformer = make_full_model(len(itos), d_model)\n",
    "lm = LM(len(itos), d_model, 1400, 3)\n",
    "net = Img2Seq(img_encoder, transformer, lm, seq_len)\n",
    "\n",
    "# AdamW16 = partial(optim.Adam, betas=(0.9,0.99), eps=1e-4)  ->  #.to_fp16(max_scale=256)\n",
    "# partial: way to always call a function with a given set of arguments or keywords\n",
    "\n",
    "learn = Learner(data, net, loss_func=LabelSmoothing(smoothing=0.1), metrics=[CER(itos)], callback_fns=TeacherForce)\n",
    "learn.clip_grad(0.25)\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Add LM to model state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd = torch.load(PATH/'models/v1_gelu_512.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lm_sd = torch.load('data/wikitext/models/wiki103_lm_enc.pth', map_location=device)\n",
    "\n",
    "from collections import OrderedDict\n",
    "new_lm_sd = OrderedDict()\n",
    "for k, v in lm_sd.items():\n",
    "    name = 'lm.lm.'+k\n",
    "    new_lm_sd[name] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd['model'].update(new_lm_sd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.model.load_state_dict(sd['model'], strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('v1_gelu_512_wiki103_base_lm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Load and split learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.load('v1_gelu_512_wiki103_base_lm')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.split([learn.model.img_enc, learn.model.transformer, learn.model.lm.lm, learn.model.lm.mixer])\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.freeze_to(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.model.lm.lm.encoder.weight.requires_grad, learn.model.lm.mixer.weight.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(1, max_lr=2e-3)\n",
    "\n",
    "### Training - GeLU, singlehead attn\n",
    "# 3x2\n",
    "# 12.332767\t10.326035\t0.070659   sz:128, bs:100, lr:1e-3   'v1_exp_gelu'   \n",
    "# 2.631910\t2.435502\t0.015014   sz:256, bs:60, lr:1e-3    'v1_gelu_256'\n",
    "# Mix\n",
    "# 29.340302\t22.766819\t0.023261   sz:400, bs:20, lr:1e-3    'v1_gelu_400'\n",
    "# 14.096607\t11.613899\t0.014052   sz:512, bs:15, lr:5e-5    'v1_gelu_512'\n",
    "\n",
    "### LM\n",
    "# 3x2, 128, last_layer\n",
    "# 8.440049\t11.431677\t0.078723\t04:12   CustomSampler        'lm_exp1'\n",
    "# 12.273294\t10.339036\t0.070205\t03:38   w/out CustomSampler  'lm_exp2'\n",
    "# 12.611145\t8.729902\t0.057937\t03:49   wiki103 LM           'lm_exp5'\n",
    "\n",
    "# 11.200677\t9.639755\t0.064239\t04:47   w/out CustomSampler; unfreeze() - lrs:(2e-5, 2e-4, 2e-3)  'lm_exp3'\n",
    "# 11.895980\t9.951677\t0.068518\t03:44   w/out LM or CustomSampler; lrs:(2e-5, 2e-4)               'lm_exp4'\n",
    "# 12.091815\t8.443047\t0.056295\t05:03   wiki103 LM; unfreeze() - lrs:(2e-5, 2e-4, 2e-3)  'lm_exp6'\n",
    "\n",
    "# Mix, 512\n",
    "# last layer\n",
    "# 53.998314\t32.622627\t0.024337   'v1_gelu_512_wiki2_base_lm_last_layer'\n",
    "# 13.940077\t10.796149\t0.010067\t51:19  'v1_gelu_512_wiki103_base_lm_last_layer'\n",
    "\n",
    "# full - unfreeze()\n",
    "# 13.917192\t10.200513\t0.009347\t1:15:24   3cycles 'v1_gelu_512_wiki103_base_lm_full'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('v1_gelu_512_wiki103_base_lm_last_layer')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.unfreeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lrs = slice(2e-5, 2e-4, 2e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(3, max_lr=lrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('v1_gelu_512_wiki103_base_lm_full')\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.load('v1_edited_gelu_400') #, strict=False)\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.data = data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.lr_find()\n",
    "learn.recorder.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# def convert_h5_to_pth(fname, learner):\n",
    "#     sd = torch.load(PATH/f\"models/{fname}.h5\")#, map_location='cpu')\n",
    "#     learner.model.load_state_dict(sd)\n",
    "#     learner.save(f\"v1_{fname}\")\n",
    "\n",
    "# convert_h5_to_pth('tfmr_mix_words_512', learn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(3, max_lr=5e-5)\n",
    "# 'v1_exp_baseline'\n",
    "# 0\t92.709564\t85.651016\t0.706222\t05:29\n",
    "# 1\t61.329845\t50.974323\t0.388731\t04:58\n",
    "# 2\t47.806641\t41.577847\t0.311466\t05:10\n",
    "# 3\t42.097706\t36.231873\t0.265740\t05:03\n",
    "# 4\t40.572376\t34.968216\t0.255079\t05:06\n",
    "\n",
    "# 3x2, sz:128, bs:100, fit_one_cycle(5, 1e-4)\n",
    "# 43.886196\t36.105576\t0.265459   baseline\n",
    "# 40.572376\t34.968216\t0.255079   baseline    'v1_exp_baseline'          ***\n",
    "# 55.031853\t49.186695\t0.373046   remove encoder; multihead(8) attn\n",
    "# 40.953770\t36.203377\t0.269332   multihead(8) attn\n",
    "# 40.215485\t35.141953\t0.259177   multihead(4) attn  'v1_exp_multiattn'   ***\n",
    "\n",
    "# lr: 1e-3\n",
    "# 12.597528\t10.658405\t0.074232   baseline   'v1_exp_mask'   ***\n",
    "# 13.165689\t11.060272\t0.075740   nn.ReLU(True) (new itos.pkl)   'v1_exp_relu'\n",
    "# 22.972273\t20.039871\t0.134025   AdamW16 optimizer\n",
    "# 14.912671\t12.835587\t0.088028   d_model=400, hidden=1152\n",
    "# 12.332767\t10.326035\t0.070659   GeLU activation  'v1_exp_gelu'   ***\n",
    "# 13.439913\t11.265744\t0.077952   Swish activation\n",
    "# 10.887406\t9.458515\t0.064375   GeLU, multihead(4) attn   'v1_exp_gelu_multihead' ***\n",
    "# 26.880289\t23.088377\t0.155487   GeLU, AMSGrad(eps:1e-4)\n",
    "# 12.324612\t10.524032\t0.071833   increase precision of LN (1e-5) and attn mask (1e-6)\n",
    "\n",
    "# edited_data, lr: 1e-3, GeLU\n",
    "# 9.280172\t9.898072\t0.073765   normalize(imagenet_stats), no tfms\n",
    "# 9.247200\t9.739366\t0.072657   normalize(), no tfms\n",
    "\n",
    "# sentencepiece \n",
    "# 14.591818\t13.716891\t0.269826    1k\n",
    "# 15.130762\t13.952784\t0.275594    1k, em_sz:128\n",
    "# 22.799696\t23.611570\t0.543942    8k\n",
    "\n",
    "# different architectures (no pretraining)\n",
    "# 22.601816\t17.765850\t0.122246   xresnet34 [256]  'v1_exp_xres34'\n",
    "# 49.068211\t43.858349\t0.337089   xresnet34 [512]\n",
    "# 19.258896\t15.660585\t0.107812   resnet34           ***\n",
    "# 22.787096\t18.036461\t0.124185   resnet50\n",
    "# 23.279938\t18.810421\t0.132146   xception [728]   'v1_exp_xception'\n",
    "# 26.723740\t22.318148\t0.157468   mobilenet v2 (pretrained, gelu, singlehead)  'v1_gelu_mobilenet'\n",
    "# 16.280537\t13.032956\t0.086539   densenet201 (pre, gelu, single, [0][:-3])    'v1_gelu_dense201'\n",
    "\n",
    "# mixed precision\n",
    "# 13.537477\t11.407567\t0.078925   mixed precision  fp_16(max_scale=256), AdamW16     'v1_exp_fp16'\n",
    "# 18.175873\t15.319780\t0.105136   mixed precision bs: 200 (5749 GPU MB)\n",
    "\n",
    "\n",
    "### Training - GeLU, multihead(4) attn\n",
    "# 3x2\n",
    "# 10.887406\t9.458515\t0.064375   sz:128, bs:100, lr:1e-3   'v1_exp_gelu_multihead'\n",
    "# 2.238591\t2.254460\t0.012747   sz:256, bs:60, lr:1e-3    'v1_gelu_multi_256'\n",
    "# Mix\n",
    "# 23.151602\t19.383263\t0.018239   sz:400, bs:20, lr:1e-3    'v1_gelu_multi_400'\n",
    "# 9.212606\t9.694052\t0.010178   sz:512, bs:15, lr:6e-5    'v1_gelu_multi_512'\n",
    "\n",
    "### Training - GeLU, singlehead attn\n",
    "# 3x2\n",
    "# 12.332767\t10.326035\t0.070659   sz:128, bs:100, lr:1e-3   'v1_exp_gelu'   \n",
    "# 2.631910\t2.435502\t0.015014   sz:256, bs:60, lr:1e-3    \n",
    "# Mix\n",
    "# 29.340302\t22.766819\t0.023261   sz:400, bs:20, lr:1e-3    'v1_gelu_400'\n",
    "# 14.096607\t11.613899\t0.014052   sz:512, bs:15, lr:5e-5    'v1_gelu_512'\n",
    "\n",
    "# with LM\n",
    "# 53.998314\t32.622627\t0.024337   last layer only\n",
    "\n",
    "### Training - Edited Data - GeLU, singlehead attn\n",
    "# sm\n",
    "# 10.841042\t9.550449\t0.071475   sz:128, bs:100, lr:1e-3  'v1_gelu_128'\n",
    "# 2.631370\t2.669811\t0.020308   sz:256, bs:100, lr:1e-3  'v1_gelu_256'\n",
    "# mix\n",
    "# 18.652050\t14.330537\t0.021167   sz:400, bs:20, lr:1e-3   'v1_edited_gelu_400'\n",
    "# 11.513125\t8.577786\t0.016521   sz:512, bs:15, lr:5e-5, 3cycle   'v1_edited_gelu_512'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('v1_edited_gelu_512')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.export()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# learn = load_learner(PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Word + Char Arch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        \n",
    "        self.word_decoder = decoder\n",
    "        self.char_decoder = decoder\n",
    "        \n",
    "        self.src_embed = src_embed\n",
    "        \n",
    "        self.tgt_embed = tgt_embed\n",
    "        \n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, src, tgt, tgt_mask=None):\n",
    "        return self.decode(self.encode(src), tgt, tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_embed(src))\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(self.tgt_embed(tgt), src, tgt_mask)\n",
    "    \n",
    "    def generate(self, outs):\n",
    "        return self.generator(outs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Embeddings(nn.Module):\n",
    "    def __init__(self, d_model, vocab):\n",
    "        super(Embeddings, self).__init__()\n",
    "        self.lut = nn.Embedding(vocab, d_model)\n",
    "        self.d_model = d_model\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.lut(x) * 18"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ImageAdaptor(nn.Module):\n",
    "    def __init__(self, em_sz, d_model, drop=0.2):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.linear = nn.Linear(em_sz, d_model)\n",
    "        self.dropout = nn.Dropout(drop)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        x = self.linear(x) * 4\n",
    "        x = self.dropout(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(c_vocab, w_vocab, em_sz=256, d_model=512, N=4, drop=0.2):\n",
    "    c = deepcopy\n",
    "    attn = SingleHeadedAttention(d_model)\n",
    "#     attn = MultiHeadedAttention(d_model, 4)\n",
    "    ff = PositionwiseFeedForward(d_model, drop)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drop), N),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), drop), N),\n",
    "        ImageAdaptor(em_sz, d_model),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, c_vocab), PositionalEncoding(d_model, 1, drop, 2000)\n",
    "        ),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, w_vocab), PositionalEncoding(d_model, 1, drop, 2000)\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab)\n",
    "    )\n",
    "        \n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "                    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Img2Seq(nn.Module):\n",
    "    def __init__(self, img_encoder, transformer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None):\n",
    "        feats = self.img_enc(src)                            # ([bs, h*w, d_model])\n",
    "        dec_outs = self.transformer(feats, tgt, tgt_mask)    # ([bs, sl, d_model])\n",
    "        out = self.transformer.generate(dec_outs)            # ([bs, sl, vocab])\n",
    "        return out\n",
    "\n",
    "    def greedy_decode(self, src, seq_len):\n",
    "        with torch.no_grad():\n",
    "            feats = self.transformer.encode(self.img_enc(src))\n",
    "            bs = src.size(0)\n",
    "            tgt = torch.ones((bs,1), dtype=torch.long, device=device)\n",
    "\n",
    "            res = []                \n",
    "            for i in tqdm(range(seq_len)):\n",
    "                dec_outs = self.transformer.decode(feats, Variable(tgt))\n",
    "                prob = self.transformer.generate(dec_outs[:,-1])\n",
    "                res.append(prob)\n",
    "                pred = torch.argmax(prob, dim=-1, keepdim=True)\n",
    "                if (pred==0).all(): break\n",
    "                tgt = torch.cat([tgt,pred], dim=-1)\n",
    "            out = torch.stack(res).transpose(1,0).contiguous()\n",
    "            return out      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "d_model = 512\n",
    "em_sz = 256\n",
    "img_encoder = ResnetBase(em_sz)\n",
    "transformer = make_full_model(len(itos), em_sz, d_model)\n",
    "net = Img2Seq(img_encoder, transformer)\n",
    "\n",
    "opt_fn = partial(optim.Adam, betas=(0.7, 0.99)) #, lr=0, betas=(0.9, 0.98), eps=1e-9)\n",
    "# partial: way to always call a function with a given set of arguments or keywords\n",
    "\n",
    "learn = RNN_Learner(data, BasicModel(to_gpu(net)), opt_fn=opt_fn)\n",
    "learn.clip = 0.25\n",
    "learn.crit = LabelSmoothing(smoothing=0.1)   #XE_loss\n",
    "learn.metrics = [cer, acc]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### load state dict for changes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#LM\n",
    "LM_PATH = Path('data/wikitext/wikitext-2-raw')\n",
    "sd = torch.load(LM_PATH/'models'/'LM.h5', map_location=lambda storage, loc: storage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd = torch.load(PATH/'models'/'tfmr_paragraph.h5', map_location=lambda storage, loc: storage)\n",
    "sd.pop('img_enc.linear.bias')   # need to remove mismatched linear weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.model.load_state_dict(sd, strict=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('tfmr_mix_words_512')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## LR find"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.lr_find(stepper=TfmrStepper)\n",
    "learn.sched.plot(n_skip=0, n_skip_end=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Experimentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#gpu\n",
    "lr=1e-4\n",
    "learn.fit(lr, 5, cycle_len=1, stepper=TfmrStepper, use_clr=(20,5))\n",
    "# 3x1: XE/bs, kaiming_normal, no pos_enc, emb/out weight tying, no encoder, 2 layers, 256/256, single attn\n",
    "# 52.199187  50.1736    0.730092    lr:1e-4\n",
    "# 17.318725  13.803769  0.205331    2nd run, lr: 1e-3\n",
    "# 11.96539   9.991971   0.151522    3rd run, lr: 1e-4     **tfmr\n",
    "\n",
    "# Loss fns\n",
    "# 43.177135  41.245482  1.185836    LabelSmoothing (KL/bs)\n",
    "# 4215.8745  4012.3099  0.957795    LabelSmoothing (KL)\n",
    "# 45.200728  43.662235  0.753881    LabelSmoothing (KL/bs w/out padding logic)  perplexity < BLEU/acc\n",
    "# 5340.5770  5129.8346  0.738853    XE                    --- scaling by BS not a factor\n",
    "# 52.199187  50.1736    0.730092    XE/bs  **\n",
    "\n",
    "# Initialization/Activation\n",
    "# 52.199187  50.1736    0.730092    kaiming_normal  (ff activation: leaky_relu)\n",
    "# 52.987157  50.976083  0.758013    kaiming_uniform\n",
    "# 49.553779  47.270616  0.702353    xavier_normal\n",
    "# 49.380884  47.345175  0.673508    xavier_uniform  (ff activation: leaky_relu)   **\n",
    "# 49.849071  47.369221  0.708807    xavier_uniform  (ff activation: relu)\n",
    "\n",
    "# Positional Encoding\n",
    "# 35.426869  30.449724  0.423573    target only; no embed/out weight tying   **\n",
    "# 35.390268  31.216677  0.45097     \"\" ; w/ scaling factor [* math.sqrt(d_model)]\n",
    "# 45.596748  42.213515  0.572662    target only; w/ embed/out weight tying\n",
    "\n",
    "# Encoder\n",
    "# 34.46771   29.207974  0.408425    **\n",
    "\n",
    "# Attention\n",
    "# 32.087941  26.431674  0.362717    SingleHead - linear layers for (q,k,v)\n",
    "# 37.739475  32.962685  0.454753    MultiHead (8)\n",
    "# 36.09057   31.346036  0.432342    MultiHead (4)\n",
    "# 31.069687  26.00539   0.360032    MultiHead (1)     **\n",
    "\n",
    "# N layers (6)\n",
    "# 51.119606  48.809905  0.645302\n",
    "\n",
    "# d_model = 512\n",
    "# 21.032752  17.234559  0.230763    ~11:35  **\n",
    "\n",
    "# em_sz\n",
    "# 27.589121  22.698359  0.307399    128  ~14:09\n",
    "# 28.320952  24.183105  0.338207    512  ~11:31\n",
    "\n",
    "# include layer_norm in encoder/decoder\n",
    "# 24.035412  19.824303  0.270881   modified: self.norm(x + self.dropout(sublayer(x)))\n",
    "# 20.347157  17.23408   0.228328   original:  x + self.dropout(sublayer(self.norm(x)))   **tfmr_experiment\n",
    "# 10.732795  9.875141   0.12633    2nd run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('tfmr_experiment')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Freezing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.freeze_to(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lr=1e-4\n",
    "learn.fit(lr, 3, cycle_len=1, stepper=TfmrStepper, use_clr=(20,5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.unfreeze()\n",
    "lr = np.array([lr/10, lr])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Initial size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lr=1e-5\n",
    "learn.fit(lr, 5, cycle_len=1, stepper=TfmrStepper, use_clr=(20,8), best_save_name=f'best_{FOLDER}_{sz}')\n",
    "# 3x1, sz: 128, bs: 100\n",
    "# 20.347157  17.23408   0.228328\n",
    "# 10.732795  9.875141   0.12633     2nd run\n",
    "\n",
    "# 6.753466   6.750452   0.085309    Full Transformer (4N/4h), 10cycles(20,10)   ~33m   'tfmr_full_3x1'\n",
    "# 7.0736     6.798651   0.085855    \"\", single attn, 10cycles(20,10)            ~38m\n",
    "\n",
    "# 7.161385   8.286392   0.102471   0.92869    exp_3x1\n",
    "\n",
    "# 3x2; sz: 128, bs: 100\n",
    "# 65.659014  59.8375    0.492443   Transformer (4N/8h), 5cycles(20,8)    'tfmr_8head_128'\n",
    "#     PosEnc2d instead of Encoder (4N/4h) 5cycles(20,5)  'posenc2d'\n",
    "\n",
    "# 3x2; sz: 256, bs: 60\n",
    "# 16.186042  14.892255  0.101064   Transformer (4N/8h), 5cycles(20,8)    'tfmr_8head_256'\n",
    "\n",
    "\n",
    "# 3x2; sz: 256, bs: 60\n",
    "# 14.65408   13.10881   0.090017    preload tfmr_3x1_256\n",
    "# 4.464029   5.072127   0.037011    Full Transformer (4N/0h), 10cycles(20,10)   ~1h 46m\n",
    "\n",
    "# 6.853619   7.136484   0.050812   0.962716    exp_3x2\n",
    "\n",
    "# Mixed Synth, sz: 128, bs: 100\n",
    "# 15.684915  14.209854  0.133399    enc/dec/mixer, N:4, pre-loaded LM  ~1h\n",
    "\n",
    "# lg; sz: 256, bs: 60\n",
    "# 161.081839 137.861294 0.231951    enc/dec/mixer, N:4, pre-loaded LM  ~1h 45m\n",
    "\n",
    "# lg; sz: 400, bs: 30\n",
    "# 98.491235  89.396089  0.123849    preload tfmr_3x2      'tfmr_lg'\n",
    "# 101.903571 88.239827  0.122516    -same as above-       'tfmr_lg2'\n",
    "# bs: 45\n",
    "# 161.90945  169.713365 0.207615    dropout:0.5    'tfmr_lg_tmp'\n",
    "# 29.706235  25.84234   0.039261    Full Transformer (4N/0h), 10cycles(20,10)   ~3h 47m\n",
    "\n",
    "# concat lines, sz: 512, bs: 20\n",
    "# LMmixer\n",
    "# 68.669186  56.487096  0.123718    345, preload 'tfmr_lg_LM_mixer'\n",
    "# 144.601336 116.728641 0.153257    678, preload 345\n",
    "#      9-12, preload 678, lr: 1e-5\n",
    "\n",
    "# Full\n",
    "# 26.620874  24.083611  0.05163     345, preload 'tfmr_full_lg'  ~2hr  'tfmr_cat345_full'\n",
    "# 46.148448  37.998259  0.04748     678, preload 345   ~2h 25m  'tfmr_cat678_full'\n",
    "# 56.235635  42.104473  0.03309     9-12, preload 678   ~3h 25m   'tfmr_cat9-12_full'\n",
    "\n",
    "# paragraph; sz:512, bs:30\n",
    "# 289.464371 202.553455 0.175994     preload tfmr_lg2, 4 cycles\n",
    "# 178.432805 145.702918 0.128241     2nd run, 5 cycles\n",
    "# 144.489454 130.875442 0.116076     3rd run              'tfmr_paragraph'\n",
    "# bs: 20\n",
    "# 66.923881  70.373825  0.069164     Full Transformer (4N/0h), 10cycles(20,10)     'tfmr_full_paragraph'\n",
    "# 81.810724  78.974658  0.078329     \"\", reversed src_attn/self_attn\n",
    "\n",
    "# 9-12; sz: 512, bs: 30\n",
    "# 139.165167 102.821483 0.089816     preload 'tfmr_full_paragraph'      '9-12_512'\n",
    "\n",
    "# paragraph; sz: 800, bs: 8\n",
    "# 24.8       45.8       0.0407       preload tfmr_cat9-12_full_800, 3cycles    'tfmr_pg_full_800'\n",
    "\n",
    "# concat lines + pg, sz: 800, bs: 8\n",
    "# 9.978522   9.815259   0.006313     preload tfmr_cat9-12_full_800, 1cycle     'tfmr_catpg_800'\n",
    "\n",
    "\n",
    "# mix_words (new dataset); sz: 400, bs: 30\n",
    "# 60.475961  48.045776  0.059005     Tfmr (4N/0h); preload 'tfmr_full_lg' (512)     'tfmr_mix_words_400'\n",
    "\n",
    "# mix_words (new dataset); sz: 400, bs: 10 \n",
    "# 83.258275  66.635521  0.08263      Transformer (4N/8h), 5cycles(20,8)\n",
    "# 42.608748  38.862213  0.047246     2nd run                                 'tfmr_8head_400_mix'\n",
    "\n",
    "# downloaded_images; sz: 512, bs: 10\n",
    "# 119.754323 31.507531  0.05412      'best_downloaded_images_512'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('tfmr_8head_400_mix')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Increase size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# sz,bs = 512,20\n",
    "# sz,bs = 1024,5\n",
    "# sz,bs = 256,60\n",
    "sz,bs = 400,45"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.set_data(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lr=1e-4\n",
    "learn.fit(lr, 5, cycle_len=1, stepper=TfmrStepper, use_clr=(20,5), best_save_name=f'best_{FOLDER}_{sz}')\n",
    "# 3x1; sz: 256, bs: 60\n",
    "# 8.169301   7.199782   0.087933    fresh start\n",
    "# 4.305115   4.585243   0.054132    2nd run (increase dropout 0.2)     'tfmr_3x1_256'\n",
    "# 3.6729     3.756866   0.047393    resize from 128\n",
    "\n",
    "# 1.489701   2.297418   0.027936   0.981084     exp_3x1_256\n",
    "\n",
    "# 2.450041   2.695453   0.032867    Full Transformer (4N/4h)   ~43m 48s       'tfmr_full_3x1'\n",
    "# 2.577439   2.748884   0.033926    Full Transformer (4N/0h)   ~45m 54s       'tfmr_full_3x1_single_attn'\n",
    "\n",
    "# 3x2; sz: 400, bs: 45\n",
    "# 5.701717   5.83973    0.039572    resize from 256 (increase dropout 0.3)     'tfmr_3x2'\n",
    "# 2.24673    2.478204   0.017335    Full Transformer (4N/0h)   ~2h 31m     'tfmr_full_3x2'\n",
    "\n",
    "# 6.739319   6.472689   0.042743    Transformer (4N/8h), 3cycles(20,5)    'tfmr_8head_400'\n",
    "\n",
    "# Mixed Synth, sz: 256, bs: 60\n",
    "# 3.584531   3.645265   0.031744    enc/dec/mixer, N:4, pre-loaded LM  ~2h 30m\n",
    "\n",
    "# lg; sz: 512, bs: 30\n",
    "# 38.786294  45.805161  0.061061    resize from 400  -quit after 3 epochs-\n",
    "# 8.317326   8.491839   0.010576    Full Transformer (4N/0h)   ~2h 31m     'tfmr_full_lg'\n",
    "# 48.016887  37.537244  0.054766    enc/dec/mixer, N:4, pre-loaded LM  ~6h 53m   'tfmr_lg_LM_mixer'\n",
    "\n",
    "# preload tfmr_lg2; increase dropout to 0.3 in attention and ff; wd; lower lr\n",
    "# 98.273312  60.214675  0.08372    lr: 1.5e-5, wd: 1e-5, 3 cycles(20,4)                    'tfmr_lg_wd'\n",
    "# 81.961758  53.715966  0.075622   lr: 2e-5, no wd, 4 cycles(20,5)  -quit after 2 epochs-  'tfmr_lg_wd_tmp'\n",
    "# 45.696747  37.355512  0.053213   \"\", {drops}, preloaded tfmr_lg_wd_tmp                   'tfmr_lg2'\n",
    "\n",
    "# concat lines, sz: 800, bs: 8\n",
    "# Full\n",
    "# 3.880624   4.531283   0.005706    345, preload 'tfmr_cat9-12_full'  ~  'tfmr_cat345_full_800'\n",
    "# 5.569003   6.855545   0.004001    678, preload above, 5cycles  ~2h 45m  'tfmr_cat678_full_800'\n",
    "# 9.628103   10.939916  0.00433     9-12, preload above, 3cycles  ~2h 18m  'tfmr_cat9-12_full_800'\n",
    "\n",
    "# paragraph; sz:800, bs:10\n",
    "# 100.395217 86.058532  0.074781     resize from 512\n",
    "# 78.234521  81.089716  0.067993     2nd run           'tfmr_paragraph_800'\n",
    "\n",
    "# paragraph; sz: 1024, bs:5\n",
    "# 55.707705  56.274219  0.05457     Full Transformer (4N/0h), 10cycles(20,10)   ~23m   'tfmr_full_paragraph2'\n",
    "\n",
    "# sz: 1000, bs:5\n",
    "# 22.210477  47.557047  0.041825    pretrained on tfmr_catpg_1000    'tfmr_pg_1000'\n",
    "\n",
    "# concat lines + pg, sz: 1000, bs: 5\n",
    "# 8.495862   10.158758  0.0068      3cycles     'tfmr_catpg_1000'\n",
    "\n",
    "\n",
    "# mix_words (new dataset); sz: 512, bs: 20\n",
    "# 25.864771  20.991407  0.028696    3/5 iterations -> tfmr_mix_words_400     'tfmr_mix_words_512'\n",
    "# 20.654874  17.808776  0.024131    5cycles(20,5)     'tfmr_mix_words_512_2'\n",
    "\n",
    "# mix_words (new dataset); sz: 512, bs: 8\n",
    "# 15.195948  14.417563  0.019355    Transformer (4N/8h), 6cycles(20,8)    'tfmr_8head_512_mix'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('tfmr_mix_words_512_2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('tfmr_cat9-12_full_800')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.freeze_groups([0])\n",
    "learn.model.img_enc.trainable, learn.model.transformer.trainable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lr=1e-4\n",
    "learn.fit(lr, 5, cycle_len=1, stepper=TfmrStepper, use_clr=(20,5))\n",
    "\n",
    "# sz: 128\n",
    "# input: 3x1\n",
    "# 7.298699   7.093821   0.091642   0.936467    BASELINE - Full tfmr   ~29m\n",
    "# 12.254474  10.917412  0.141336   0.900215    adaptor: conv/bn + drop, no scaling\n",
    "\n",
    "\n",
    "# 28.902502  24.930458  0.398251   0.732572    tfmr (img_enc frozen)   **3x1_base\n",
    "# 42.132684  35.906375  0.338968   0.770622    \"\", sz: 256, 3cycle(20,5)   ??resizing not an improvement??\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# input: mix (sample: 50000), img_enc frozen\n",
    "# 62.628027  59.523581  0.679014    Full Transformer -- 4N/0h, img_enc:256 w/ linear, w/ scaling factor  ~13m\n",
    "\n",
    "# 81.75129   79.767102  0.806713    \"\", img_enc:256 w/ conv/bn, pos_enc, tgt_emb drop:0.5, w/out scaling factor  ~13m\n",
    "# 76.193838  73.24591   0.757938    \"\", tgt_emb drop:0.1  ~13m\n",
    "# [105.264, 0.655, 13.205]\n",
    "# 75.780639  72.219551  0.761694    \"\", no pos_enc  ~14m\n",
    "# [102.794, 0.679, 10.355]\n",
    "# 81.270569  78.084324  0.788466    \"\", linear embedding w/ weight tying, smoothed: 0.7  ~13m\n",
    "# 72.355454  69.852579  0.733982    2nd run\n",
    "# [69.245, 0.932, 187.734]\n",
    "# 95.489009  91.202586  0.791639    \"\", XE loss\n",
    "\n",
    "\n",
    "# 1e-4, 5cycle(20,10), img_enc frozen\n",
    "# input: mix (sample: 50000)\n",
    "# 53.373692  48.385205  0.544684   0.696936    BASELINE - Full tfmr\n",
    "# greedy:    134.89598  0.643921   0.458704\n",
    "# 44.957038  39.8659    0.437462   0.75087     2nd run\n",
    "\n",
    "# 35.53787   30.969923  0.310315   0.811553    BASELINE - unfrozen img_enc    **base-unfrozen\n",
    "# greedy:    117.58634  0.392103   0.551786\n",
    "\n",
    "# 73.066214  70.389176  0.747244   0.572001    4N/0h, img_enc:256 w/ conv/bn, linear: prob dist (smoothed: 0.7)\n",
    "# 66.118537  63.631764  0.69066    0.605459    2nd run\n",
    "\n",
    "# input: 3x1\n",
    "# 28.902502  24.930458  0.398251   0.732572    BASELINE - Full tfmr (img_enc frozen)   **3x1_base\n",
    "# 42.132684  35.906375  0.338968   0.770622    \"\", sz: 256, 3cycle(20,5)   ??resizing not an improvement??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('3x1_base')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sz,bs = 256,60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.set_data = data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lr=1e-4\n",
    "learn.fit(lr, 10, cycle_len=1, stepper=TfmrStepper, use_clr=(20,10))\n",
    "\n",
    "# sz:256 \n",
    "# 60.575186  52.815338  0.335582   0.803861     Baseline - full tfmr\n",
    "# 32.735003  27.785917  0.147773   0.907831     \"\" , unfrozen    **base\n",
    "# val:       0.00983    0.591396   0.75641\n",
    "# greedy:    0.04553    0.495198   0.509295\n",
    "\n",
    "\n",
    "# 101.034633 95.16612   0.625264   0.642998    4N/0h, img_enc:256 w/ conv/bn, linear: prob dist (smoothed: 0.7)\n",
    "# 79.018152  70.1724    0.456446   0.735568    2nd run  ~28m\n",
    "# 64.660601  55.688069  0.352824   0.792006    3rd run    **prob_embed_experiment\n",
    "# [463.614, 2.677, 1025.052]   FAIL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.unfreeze()\n",
    "learn.model.img_enc.trainable, learn.model.transformer.trainable\n",
    "\n",
    "lrs = np.array([lr/10, lr])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lr=1e-4\n",
    "learn.fit(lrs, 5, cycle_len=1, stepper=TfmrStepper, use_clr=(20,10))\n",
    "\n",
    "# 1e-4, 5cycle(20,10)\n",
    "# input: 3x1, sz: 256\n",
    "# 25.84745   21.318996  0.179581   0.874763    BASELINE - Full tfmr, img_enc unfrozen   **3x1_base\n",
    "# greedy:    115.10662  0.513799   0.356322\n",
    "# val:        36.26467  0.793716   0.598851"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Decode Stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "v_dl = iter(data.val_dl)\n",
    "denorm = data.val_ds.denorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = next(v_dl)\n",
    "imgs = denorm(x)\n",
    "\n",
    "shifted_y = rshift(y).long()\n",
    "tgt_mask = subsequent_mask(shifted_y.size(-1)) #make_tgt_mask(shifted_y)\n",
    "\n",
    "learn.model.eval()\n",
    "\n",
    "v_preds = learn.model(x, shifted_y, tgt_mask)\n",
    "v_res = torch.argmax(v_preds, dim=-1)\n",
    "v_attn = source_attn()\n",
    "\n",
    "g_preds = learn.model.greedy_decode(x, seq_len)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "g_attn = source_attn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "v = [learn.crit(v_preds, y).item(), cer(v_preds, y), acc(v_preds, y).item()]\n",
    "print(f'valid:     {str(v[0])[:7]}   {str(v[1])[:7]}   {str(v[2])[:7]}')\n",
    "\n",
    "g = [learn.crit(g_preds, y).item(), cer(g_preds, y), acc(g_preds, y).item()]\n",
    "print(f'greedy:    {str(g[0])[:7]}   {str(g[1])[:7]}   {str(g[2])[:7]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# tfmr full\n",
    "#             loss       cer        acc\n",
    "# valid:     6.55702   0.08746   0.934839   3x1, 128/100\n",
    "# greedy:    42.5730   0.10094   0.759062\n",
    "\n",
    "# previous best models\n",
    "# valid:     2.42856   0.02981   0.981771   3x1, 256/60, 'tfmr_full_3x1_single_attn'\n",
    "# greedy:    11.9438   0.02745   0.933854\n",
    "\n",
    "# valid:     1.47604   0.00797   0.995238   3x2, 400/45,  'tfmr_full_3x2'\n",
    "# greedy:    21.4682   0.01097   0.937622\n",
    "\n",
    "# valid:     7.30494   0.01206   0.990868   lg,  512,30,  'tfmr_full_lg'\n",
    "# greedy:    866.544   0.14109   0.372146\n",
    "\n",
    "# \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# res = torch.argmax(preds, dim=-1)\n",
    "denorm = data.val_ds.denorm\n",
    "\n",
    "fig, axes = plt.subplots(2,2, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    t = char_label_text(res[i])\n",
    "    ax=show_img(denorm(x[i])[0], ax=ax, title=t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Previous Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# 3x1: XE/bs, kaiming_normal, no pos_enc, emb/out weight tying, no encoder, 2 layers, 256/256, single attn\n",
    "# 52.199187  50.1736    0.730092    lr:1e-4\n",
    "# 17.318725  13.803769  0.205331    2nd run, lr: 1e-3\n",
    "# 11.96539   9.991971   0.151522    3rd run, lr: 1e-4     **tfmr\n",
    "\n",
    "# Loss fns\n",
    "# 43.177135  41.245482  1.185836    LabelSmoothing (KL/bs)\n",
    "# 4215.8745  4012.3099  0.957795    LabelSmoothing (KL)\n",
    "# 45.200728  43.662235  0.753881    LabelSmoothing (KL/bs w/out padding logic)  perplexity < BLEU/acc\n",
    "# 5340.5770  5129.8346  0.738853    XE                    --- scaling by BS not a factor\n",
    "# 52.199187  50.1736    0.730092    XE/bs  **\n",
    "\n",
    "# Initialization/Activation\n",
    "# 52.199187  50.1736    0.730092    kaiming_normal  (ff activation: leaky_relu)\n",
    "# 52.987157  50.976083  0.758013    kaiming_uniform\n",
    "# 49.553779  47.270616  0.702353    xavier_normal\n",
    "# 49.380884  47.345175  0.673508    xavier_uniform  (ff activation: leaky_relu)   **\n",
    "# 49.849071  47.369221  0.708807    xavier_uniform  (ff activation: relu)\n",
    "\n",
    "# Positional Encoding\n",
    "# 35.426869  30.449724  0.423573    target only; no embed/out weight tying   **\n",
    "# 35.390268  31.216677  0.45097     \"\" ; w/ scaling factor [* math.sqrt(d_model)]\n",
    "# 45.596748  42.213515  0.572662    target only; w/ embed/out weight tying\n",
    "\n",
    "# Encoder\n",
    "# 34.46771   29.207974  0.408425    **\n",
    "\n",
    "# Attention\n",
    "# 32.087941  26.431674  0.362717    SingleHead - linear layers for (q,k,v)\n",
    "# 37.739475  32.962685  0.454753    MultiHead (8)\n",
    "# 36.09057   31.346036  0.432342    MultiHead (4)\n",
    "# 31.069687  26.00539   0.360032    MultiHead (1)     **\n",
    "\n",
    "# N layers (6)\n",
    "# 51.119606  48.809905  0.645302\n",
    "\n",
    "# d_model = 512\n",
    "# 21.032752  17.234559  0.230763    ~11:35  **\n",
    "\n",
    "# em_sz\n",
    "# 27.589121  22.698359  0.307399    128  ~14:09\n",
    "# 28.320952  24.183105  0.338207    512  ~11:31\n",
    "\n",
    "# include layer_norm in encoder/decoder\n",
    "# 24.035412  19.824303  0.270881   modified: self.norm(x + self.dropout(sublayer(x)))\n",
    "# 20.347157  17.23408   0.228328   original:  x + self.dropout(sublayer(self.norm(x)))   **tfmr_experiment\n",
    "# 10.732795  9.875141   0.12633    2nd run\n",
    "\n",
    "\n",
    "\n",
    "# 3x2 attention/pos_enc experiments\n",
    "# 44.391676  40.633241  0.308526    baseline (src before self attn)    ~1h 7m\n",
    "# 18.111974  18.16437   0.117756    10 cycles                          ~1h 52m      tmp_base\n",
    "# 57.85619   51.015644  0.418202    pos_enc2d/no encoder; 256 + conv 1   ~44m       \n",
    "# 27.356499  24.327117  0.17227     2nd run - attention not working well            tmp_2\n",
    "# 16.581052  15.659028  0.104724    3rd run                                         tmp_3\n",
    "# 55.944376  52.067722  0.417235    baseline w/ pos_enc2d; 256 + conv1    ~57m      tmp2\n",
    "\n",
    "\n",
    "# 3x1, 1e-4, 10 cycles, N: 2\n",
    "# 23.452119  20.750421  0.30148     conv, enc: pos2d, out: Linear    ~20m\n",
    "# 10.747468  11.945359  0.148952    2nd run\n",
    "\n",
    "# 22.816471  20.652263  0.300892    w/ STN (32/7/5/32/32) before img_encoder\n",
    "# 24.154148  21.474449  0.314114    w/ STN (32/7/5/64/64) before img_encoder\n",
    "# 24.315829  21.928237  0.319867    adding bn after conv (no STN)\n",
    "# 25.852014  23.393568  0.345979    \"\", w/ STN(em_sz, d_model, 2) after img_encoder\n",
    "# 22.586863  20.337497  0.294407    \"\", w/ STN (32/7/5/32/32) before img_encoder     ~21m\n",
    "# 11.062696  11.752348  0.146163    2nd run\n",
    "\n",
    "# 17.356167  15.935936  0.223704    linear, enc: None, out: Linear    ~19m\n",
    "# 15.614188  14.634603  0.205307    conv/bn, enc: None, out: Linear    ~19m   **\n",
    "# 17.618474  15.867109  0.220721    \"\"    ~20m\n",
    "# 8.634967   9.732577   0.13152     2nd run    'tmp'  [src-attn tracks well]\n",
    "\n",
    "# 17.213856  16.036616  0.228302    conv/bn, enc: None, out: MLP    ~20m\n",
    "# 17.269863  16.052437  0.224499    conv/bn, enc: Encoder, out: Linear    ~23m\n",
    "# 17.999301  16.085523  0.223856    conv/bn, enc: pos2d, out: Linear    ~19m\n",
    "# 19.352792  17.537506  0.249201    conv/bn, STN(32/7/5), enc: pos2d, out: Linear ~21m\n",
    "# 21.778735  19.844897  0.288944    conv/bn, STN(32/7/5), enc: pos2d, out: MLP    ~22m\n",
    "\n",
    "# 15.627235  14.005776  0.18804     dec: w/ self-attn + src-attn\n",
    "# 7.056934   7.818964   0.10034     2nd run    'tmp2'    [loss decrease but src-attn tracks worse]\n",
    "\n",
    "# 17.841002  16.682544  0.221875    enc: src-attn, dec: self-attn, mixer: +tanh   ~22m\n",
    "# 6.         8.         0.12        2nd run\n",
    "\n",
    "# 15.780708  14.452696  0.205315    enc: src-attn, dec: self-attn, mixer: cat/lin   ~20m    'tmp3'\n",
    "\n",
    "# N: 4\n",
    "# 14.322309  13.357375  0.190247    enc: src-attn, dec: self-attn, mixer: cat/lin   ~25m\n",
    "# 7.48091    7.132944   0.092341    \"\", dec: preloaded LM\n",
    "\n",
    "# paragraph, 1e-4, 10 cycles, N: 4\n",
    "# 691.015346 621.679181 0.645048    enc: src-attn, dec: preloaded LM, mixer: cat/lin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from scipy.ndimage import gaussian_filter\n",
    "k=16\n",
    "\n",
    "def torch_scale_attns(attns):\n",
    "    bs,sl,hw = attns.shape\n",
    "    num = int(math.sqrt(hw))   # sz // k\n",
    "    mod = attns.view(bs,sl,num,num)\n",
    "    scaled = F.interpolate(mod, size=sz)\n",
    "    return scaled  #([bs, sl, h, w])\n",
    "\n",
    "def g_filter(att):\n",
    "    return gaussian_filter(att, sigma=k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def self_attn(layer=-1): return learn.model.transformer.decoder.layers[layer].self_attn.attn.data.cpu()\n",
    "def source_attn(layer=-1): return learn.model.transformer.decoder.layers[layer].src_attn.attn.data.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def cer(preds, targs):\n",
    "    bs,sl = targs.size()\n",
    "    \n",
    "    res = torch.argmax(preds, dim=2)\n",
    "    error = 0\n",
    "    for i in range(bs):\n",
    "        p = char_label_text(res[i])   #.replace(' ', '')\n",
    "        t = char_label_text(targs[i]) #.replace(' ', '')\n",
    "        error += Lev.distance(t, p)/len(t)\n",
    "    return error/bs\n",
    "\n",
    "def char_label_text(pred):\n",
    "#     return sp.DecodeIds(pred.tolist())\n",
    "    ints = to_np(pred).astype(int)\n",
    "    nonzero = ints[np.nonzero(ints)]\n",
    "    return ''.join([itos[i] for i in nonzero])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def show_img(im, figsize=None, ax=None, alpha=None, title=None):\n",
    "    if not ax: fig,ax = plt.subplots(figsize=figsize)\n",
    "    if im.shape[0] == 3: im = image2np(im.data)\n",
    "    ax.imshow(im, alpha=alpha)\n",
    "    if title: ax.set_title(title)\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = next(iter(learn.data.train_dl))\n",
    "# x,y = learn.data.one_batch(ds_type=DatasetType.Train)\n",
    "imgs = learn.data.denorm(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Uploaded Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def thresh_edit(fname, thresh=100, bg=245):\n",
    "    im = Image.open(fname).convert('L')  #grayscale\n",
    "    np_im = np.array(im)\n",
    "    im.close()\n",
    "    thresh_mask = np_im > thresh\n",
    "    np_im[thresh_mask] = bg\n",
    "    return Image.fromarray(np_im, 'L')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "im = thresh_edit(PATH/'test3.png')\n",
    "show_img(im, figsize=(15,15))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "e = 'edit_'+str(fname.name)\n",
    "edited_fname = PATH/e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "edited_im.save(edited_fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fname = PATH/'test/edit_test4.png'\n",
    "im = open_image(fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "seq,res,preds = learn.predict(im)\n",
    "print(seq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# r = torch.tensor([g_res], dtype=torch.long, device=device)\n",
    "truth = \"This is a test letter. I hope this\\nworks but I'm not sure it will.\\nMy handwriting is not very good.\"\n",
    "Lev.distance(truth, str(seq))/len(truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# x,y = next(v_dl)\n",
    "# imgs = denorm(x)\n",
    "\n",
    "learn.model.eval()\n",
    "\n",
    "shifted_y = rshift(y).long()\n",
    "tgt_mask = subsequent_mask(shifted_y.size(-1))\n",
    "v_preds = learn.model(x, shifted_y, tgt_mask)\n",
    "v_res = torch.argmax(v_preds, dim=-1)\n",
    "# v_attn = source_attn()\n",
    "\n",
    "g_preds = learn.model(x)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "# g_attn = source_attn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "v = [learn.loss_func(v_preds, y).item(), cer(v_preds, y)]\n",
    "print(f'valid:     {str(v[0])[:7]}   {str(v[1])[:7]}')\n",
    "\n",
    "g = [learn.loss_func(g_preds, y).item(), cer(g_preds, y)]\n",
    "print(f'greedy:    {str(g[0])[:7]}   {str(g[1])[:7]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(v_res[i])\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(g_res[i])\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# tfmr full\n",
    "#             loss       cer       acc\n",
    "# valid:     2.42856   0.02981   0.98177   3x1, 256/60, 'tfmr_full_3x1_single_attn'\n",
    "# greedy:    11.9438   0.02745   0.93385\n",
    "\n",
    "# valid:     1.98831   0.01858   0.98505   3x1, 256/60, 'exp_3x1_256'\n",
    "# greedy:    16.3832   0.02167   0.90944\n",
    "\n",
    "# valid:     1.47604   0.00797   0.99523   3x2, 400/45,  'tfmr_full_3x2'\n",
    "# greedy:    21.4682   0.01097   0.93762\n",
    "\n",
    "# valid:     7.30494   0.01206   0.99086   lg,  512/30,  'tfmr_full_lg'\n",
    "# greedy:    434.610   0.02398   0.69553\n",
    "\n",
    "# valid:     25.4131   0.04500   0.96773   lg, 512/30,   'tfmr_lg_LM_mixer'\n",
    "# greedy:    734.834   0.15256   0.48912\n",
    "\n",
    "# valid:     10.1481   0.00501   0.99602   cat9-12,  800/8,  'tfmr_cat9-12_full_800'\n",
    "# greedy:    1330.63   0.04525   0.53181\n",
    "\n",
    "# valid:     62.0793   0.05716   0.96214   pg,  512/20  'tfmr_full_paragraph'  (cpu)\n",
    "# greedy:    1739.62   0.08347   0.43309\n",
    "# beam:                0.07958\n",
    "# valid:     71.9120   0.06819   0.95457   \"\", 2nd batch (gpu)\n",
    "# greedy:    2027.24   0.11823   0.39238\n",
    "# beam:                0.10697\n",
    "# valid:     36.3009   0.03397   0.97633   \"\", 3rd batch (gpu)\n",
    "# greedy:    1759.59   0.07070   0.40949\n",
    "    \n",
    "# valid:     50.1414   0.04137   0.97004   pg,  1000/5,  'tfmr_pg_1000'\n",
    "# greedy:    2579.24   0.34178   0.18623\n",
    "\n",
    "# valid:     56.2722   0.04167   0.96923   pg,  1000/5,  'tfmr_catpg_1000'\n",
    "# greedy:    2432.54   0.37192   0.26174\n",
    "\n",
    "\n",
    "# valid:     3.43745   0.05444   0.98735   mix(new)  'tfmr_mix_words_400'\n",
    "# greedy:    43.6545   0.06126   0.91407\n",
    "\n",
    "# valid:     2.96891   0.03338   0.98932   mix(new)  'tfmr_mix_words_512'\n",
    "# greedy:    28.9982   0.03949   0.94500\n",
    "# valid:     2.02411   0.02128   0.99302   5 more cycles\n",
    "# greedy:    19.4735   0.02450   0.96104\n",
    "\n",
    "# valid:     41.9800   0.03879   0.97535   pg, 'tfmr_mix_words_512'\n",
    "# greedy:    1602.68   0.06259   0.49110\n",
    "\n",
    "\n",
    "# valid:     53.9049   0.04822   0.96622   pg, 'tfmr_8head_512_mix'\n",
    "# greedy:    1384.72   0.06304   0.51696\n",
    "\n",
    "\n",
    "# valid:     5.24461   0.00634   mix   'v1_gelu_512'\n",
    "# greedy:    217.438   0.02028\n",
    "# valid:     24.3609   0.02405   pg\n",
    "# greedy:    1935.30   0.05970\n",
    "\n",
    "# valid:     12.7074   0.00954   mix   'v1_gelu_512_wiki103_base_lm_last_layer'\n",
    "# greedy:    553.201   0.02814\n",
    "\n",
    "# valid:     11.4251   0.00474   mix   'v1_gelu_512_wiki103_base_lm_full'\n",
    "# greedy:    358.030   0.01187\n",
    "\n",
    "# valid:     12.1104   0.01533   mix   'v1_gelu_multi_512'\n",
    "# greedy:    506.505   0.04229\n",
    "\n",
    "\n",
    "# edited data\n",
    "# valid:     8.21437   0.06576   sm   'v1_gelu_128'\n",
    "# greedy:    185.981   0.12890\n",
    "\n",
    "# valid:     2.25060   0.01731   sm   'v1_gelu_256'\n",
    "# greedy:    149.525   0.06003\n",
    "\n",
    "# valid:     2.37294   0.01383   mix  'v1_edited_gelu_400'\n",
    "# greedy:    30.3637   0.01732\n",
    "# greedy:    1810.14   0.11409   test\n",
    "\n",
    "# valid:     9.09986   0.01656   mix  'v1_edited_gelu_512'\n",
    "# greedy:    529.308   0.03104\n",
    "# greedy:    1675.20   0.10819   test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3x1 256"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(v_res[i])\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(g_res[i])\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "fig, axes = plt.subplots(3,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(v_res[i])\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(3,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(g_res[i])\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Source Attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx = 1\n",
    "img = imgs[idx]\n",
    "\n",
    "v_chars = v_res[idx]\n",
    "v_attns = to_np(torch_scale_attns(v_attn)[idx])\n",
    "\n",
    "g_chars = g_res[idx]\n",
    "g_attns = to_np(torch_scale_attns(g_attn)[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "fig, axes = plt.subplots(5,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    a = g_filter(v_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[v_chars[i].item()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(5,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    a = g_filter(g_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[g_chars[i].item()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3x2 400"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(v_res[i])\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(g_res[i])\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Source Attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx = 1\n",
    "img = imgs[idx]\n",
    "\n",
    "v_chars = v_res[idx]\n",
    "v_attns = to_np(torch_scale_attns(v_attn)[idx])\n",
    "\n",
    "g_chars = g_res[idx]\n",
    "g_attns = to_np(torch_scale_attns(g_attn)[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "fig, axes = plt.subplots(5,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    a = g_filter(v_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[v_chars[i].item()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(5,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    a = g_filter(g_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[g_chars[i].item()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## lg 512"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(v_res[i], chunk=55)\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(g_res[i], chunk=55)\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Source Attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx = 1\n",
    "img = imgs[idx]\n",
    "\n",
    "v_chars = v_res[idx]\n",
    "v_attns = to_np(torch_scale_attns(v_attn)[idx])\n",
    "\n",
    "g_chars = g_res[idx]\n",
    "g_attns = to_np(torch_scale_attns(g_attn)[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "fig, axes = plt.subplots(5,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    a = g_filter(v_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[v_chars[i].item()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(5,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    a = g_filter(g_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[g_chars[i].item()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## cat 9-12"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(v_res[i], chunk=55)\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(g_res[i], chunk=55)\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Source Attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx = 1\n",
    "img = imgs[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "v_chars = v_res[idx]\n",
    "v_attns = to_np(torch_scale_attns(v_attn[:idx, :20]))[-1]\n",
    "\n",
    "fig, axes = plt.subplots(5,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    \n",
    "    a = g_filter(v_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[v_chars[i].item()])\n",
    "    \n",
    "del v_attns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "g_chars = g_res[idx]\n",
    "g_attns = to_np(torch_scale_attns(g_attn[:idx, :20]))[-1]\n",
    "\n",
    "fig, axes = plt.subplots(5,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    a = g_filter(g_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[g_chars[i].item()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## pg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(v_res[i], chunk=55)\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(g_res[i], chunk=55)\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Source Attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx = 1\n",
    "img = imgs[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "v_chars = v_res[idx]\n",
    "v_attns = to_np(torch_scale_attns(v_attn[:idx, :20]))[-1]\n",
    "\n",
    "fig, axes = plt.subplots(5,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    \n",
    "    a = g_filter(v_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[v_chars[i].item()])\n",
    "    \n",
    "del v_attns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "g_chars = g_res[idx]\n",
    "g_attns = to_np(torch_scale_attns(g_attn[:idx, :20]))[-1]\n",
    "\n",
    "fig, axes = plt.subplots(5,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    a = g_filter(g_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[g_chars[i].item()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Beam Search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Multi-Batch Beam Decode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def beam_cer(preds, targs):\n",
    "    bs = preds.size(0)\n",
    "    error = 0\n",
    "    for i in range(bs):\n",
    "        p = char_label_text(preds[i])\n",
    "        t = char_label_text(targs[i])\n",
    "        error += _cer(t,p)\n",
    "    return error/bs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def repeat_interleave(tensor,n):\n",
    "    res = []\n",
    "    for i in range(tensor.size(0)):\n",
    "        for _ in range(n): res.append(tensor[i])\n",
    "    return torch.stack(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# def normalize_score(score, i, alpha=0.6):\n",
    "#     length_penalty = math.pow((5 + i)/6, alpha)\n",
    "#     return score/length_penalty"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Beam Object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import heapq\n",
    "\n",
    "class Beam(object):\n",
    "    def __init__(self, beam_width=10):\n",
    "        self.heap = list()\n",
    "        self.beam_width = beam_width\n",
    "        self.best_score = None\n",
    "\n",
    "    def add(self, score, complete, seq):        \n",
    "        heapq.heappush(self.heap, (score, complete, seq))\n",
    "        if len(self.heap) > self.beam_width:\n",
    "            heapq.heappop(self.heap)\n",
    "            \n",
    "    def get_seq(self):\n",
    "        return [b[-1] for b in self.heap]\n",
    "        \n",
    "    def __iter__(self):\n",
    "        return iter(self.heap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class BeamSearch(nn.Module):\n",
    "    def __init__(self, net, beam_width, seq_len, end_tok=0):\n",
    "        super(BeamSearch, self).__init__()\n",
    "        self.img_enc = net.img_enc\n",
    "        self.transformer = net.transformer\n",
    "        self.bw = beam_width\n",
    "        self.seq_len = seq_len\n",
    "        self.end_tok = end_tok\n",
    "        self.feats = None\n",
    "        self.beams = []\n",
    "            \n",
    "    def forward(self, src):\n",
    "        with torch.no_grad():\n",
    "            bs = src.size(0)\n",
    "            \n",
    "            # initialize beam per bs\n",
    "            for _ in range(bs):\n",
    "                beam = Beam(self.bw)\n",
    "                beam.add(0.0, False, [1])\n",
    "                self.beams.append(beam)\n",
    "                \n",
    "            # encode src\n",
    "            self.feats = self.transformer.encode(self.img_enc(src))\n",
    "            \n",
    "            for i in tqdm(range(seq_len)):\n",
    "                # gather sequences from beams; combine into tensor (bs*bw)\n",
    "                prev_seq = torch.from_numpy(np.stack([b.get_seq() for b in self.beams])).view(-1,i+1)\n",
    "                \n",
    "                # generate new possibilities\n",
    "                log_probs, chars = self.prob_func(prev_seq.to(device))\n",
    "                log_probs, chars = log_probs.view(bs,-1,self.bw), chars.view(bs,-1,self.bw)\n",
    "                \n",
    "                for j in range(bs):\n",
    "                    curr_beam = Beam(self.bw)\n",
    "                    for k,(score, complete, seq) in enumerate(self.beams[j]):\n",
    "                        for l,c in zip(log_probs[j,k],chars[j,k]):\n",
    "                            log_prob,char = l.item(), c.item()\n",
    "                            curr_beam.add((score+log_prob), (char==self.end_tok), seq+[char])\n",
    "                    self.beams[j] = curr_beam\n",
    "  \n",
    "                # return if all max beams are complete\n",
    "                if (self.top_complete()==True).all(): break\n",
    "                    \n",
    "                # expand feats to match beam size (only on 2nd run)\n",
    "                if i==0: self.feats = repeat_interleave(self.feats, self.bw)\n",
    "\n",
    "            return self.top_seq()\n",
    "\n",
    "    def top_complete(self): return np.stack([max(b)[1] for b in self.beams])\n",
    "    def top_seq(self): return torch.from_numpy(np.stack([max(b)[-1] for b in self.beams]))[:,1:]\n",
    "\n",
    "    def prob_func(self, tgt):\n",
    "        mask = subsequent_mask(tgt.size(-1))\n",
    "        dec_outs = self.transformer.decode(self.feats, tgt, mask)\n",
    "        logits = self.transformer.generate(dec_outs[:,-1])\n",
    "        log_probs = logits - torch.logsumexp(logits, -1, keepdim=True) # more stable than F.softmax(logits,-1).log()\n",
    "        return torch.topk(log_probs, self.bw, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "search = BeamSearch(learn.model, 3, seq_len)\n",
    "b_res = search(x)\n",
    "\n",
    "# bw: 3, sl: 350, ~56s, 0.02491    # no normalized_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "beam_cer(b_res,y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class BeamSearch(nn.Module):\n",
    "    def __init__(self, net, beam_width, seq_len):\n",
    "        super(BeamSearch, self).__init__()\n",
    "        self.img_enc = net.img_enc\n",
    "        self.transformer = net.transformer\n",
    "        self.bw = beam_width\n",
    "        self.seq_len = seq_len\n",
    "        \n",
    "        self.feats = None\n",
    "        self.beam = None\n",
    "        self.scores = None\n",
    "        \n",
    "        net.eval()\n",
    "    \n",
    "    def forward(self, src):\n",
    "        with torch.no_grad():\n",
    "            bs = src.size(0)\n",
    "            \n",
    "            # encode src\n",
    "            self.feats = self.transformer.encode(self.img_enc(src))\n",
    "            \n",
    "            # initialize globals (beam=1 for first iteration; 3 thereafter)\n",
    "            self.beam = torch.ones((bs,1), device=device, dtype=torch.long)\n",
    "            self.scores = torch.zeros((bs,1), device=device, dtype=torch.float)\n",
    "            \n",
    "            for i in tqdm(range(seq_len)):\n",
    "                # generate new topk chars per beam(bs*bw)\n",
    "                log_probs, chars = self.prob_func(self.beam)  #(bs*beam, 3)\n",
    "\n",
    "                # compute local scores\n",
    "                scores = self.scores + log_probs\n",
    "                \n",
    "                # compute new beams per batch\n",
    "                new_scores, idxs = torch.topk(scores.view(bs,-1), self.bw, dim=-1) #(bs, 3)\n",
    "                self.scores = new_scores.view(-1,1)\n",
    "                \n",
    "                # set up new beam:\n",
    "                nxt = torch.stack([c[i] for c,i in zip(chars.view(bs,-1),idxs)]).view(-1,1)\n",
    "                pre = torch.stack([b[i//self.bw] for b,i in zip(self.beam.view(bs,-1,i+1),idxs)]).view(-1,i+1)\n",
    "                                    \n",
    "                # update globals\n",
    "                self.beam = torch.cat([pre,nxt], dim=1)\n",
    "\n",
    "                # end when top of beams are complete\n",
    "                if self.top_complete(): break\n",
    "                                                        \n",
    "                # expand feats to match beam size (only on 2nd run)\n",
    "                if i==0: self.feats = repeat_interleave(self.feats, self.bw) #.repeat(self.bw,1,1)\n",
    "                    \n",
    "            return self.top_sequences(), self.top_scores() \n",
    "\n",
    "\n",
    "    def top_complete(self): return (self.top_sequences()[:,-1]==0).all().item()   #byte tensor\n",
    "    def top_sequences(self): return self.beam.squeeze()[0::self.bw][:,1:]\n",
    "    def top_scores(self): return self.scores.squeeze()[0::self.bw]\n",
    "\n",
    "    def prob_func(self, tgt):\n",
    "        mask = subsequent_mask(tgt.size(-1))\n",
    "        dec_outs = self.transformer.decode(self.feats, tgt, mask)\n",
    "        logits = self.transformer.generate(dec_outs[:,-1])\n",
    "        log_probs = logits - torch.logsumexp(logits, -1, keepdim=True) # more stable than F.softmax(logits,-1).log()\n",
    "        return torch.topk(log_probs, self.bw, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "search = BeamSearch(learn.model, 3, seq_len)\n",
    "b_res,score = search(x)\n",
    "\n",
    "# lg, sl: 250\n",
    "# bw: 1, ~12s, 0.02398\n",
    "# bw: 3, ~28s, 0.02378\n",
    "# bw: 5, ~46s, 0.02378\n",
    "\n",
    "# pg: 1000,5\n",
    "# bs: 3, ~40s, 0.33227  (greedy: 0.28---)\n",
    "# ''   , ~30s, 0.22635  (greedy: 0.23726)\n",
    "\n",
    "# pg: 800,8\n",
    "# bs: 3, ~51s, 0.25424  (greedy: 0.28348)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "beam_cer(b_res,y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#beam\n",
    "fig, axes = plt.subplots(1,3, gridspec_kw={'hspace': 0.4}, figsize=(20, 10))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(b_res[i], chunk=55)\n",
    "    ax=show_img(imgs[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Single Beam Decode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# https://geekyisawesome.blogspot.com/2016/10/using-beam-search-to-generate-most.html\n",
    "\n",
    "import heapq\n",
    "\n",
    "class Beam(object):\n",
    "    '''\n",
    "    For comparison of prefixes, the tuple (prefix_probability, complete_sentence) is used.\n",
    "    This is so that if two prefixes have equal probabilities then a complete sentence\n",
    "    is preferred over an incomplete one since (0.5, False) < (0.5, True)\n",
    "    '''\n",
    "\n",
    "    def __init__(self, beam_width=10):\n",
    "        self.heap = list()\n",
    "        self.beam_width = beam_width\n",
    "        self.best_score = None\n",
    "\n",
    "    def add(self, score, complete, seq):        \n",
    "        # keep track of best_score so far\n",
    "        if self.best_score is None or score > self.best_score:\n",
    "            self.best_score = score\n",
    "            \n",
    "        # only add to beam if score is not more than beam_width below the best_score\n",
    "        if score > self.best_score-self.beam_width:\n",
    "            heapq.heappush(self.heap, (score, complete, seq))\n",
    "            \n",
    "        # maintain beam_width\n",
    "        if len(self.heap) > self.beam_width:\n",
    "            heapq.heappop(self.heap)\n",
    "                \n",
    "    def __iter__(self):\n",
    "        return iter(self.heap)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def beamsearch(prob_fn, seq_len, beam_width=5, start_tok=1, end_tok=3):\n",
    "    prev_beam = Beam(beam_width)\n",
    "    prev_beam.add(0.0, False, [start_tok])\n",
    "    \n",
    "    for i in tqdm(range(seq_len)):\n",
    "        curr_beam = Beam(beam_width)\n",
    "        \n",
    "        # iterate over each beam\n",
    "        for (score, complete, seq) in prev_beam:\n",
    "            if complete == True:\n",
    "                None  # only keep the completed best beam!!\n",
    "#                 curr_beam.add(score, True, seq)\n",
    "            else:\n",
    "                # iterate through topk chars, calculating scores and adding to the beam.\n",
    "                log_probs, chars = prob_fn(seq)\n",
    "                for log_prob, char in zip(log_probs, chars): \n",
    "                    log_prob,char = log_prob.item(), char.item()\n",
    "                    score += log_prob   #log probabilities are additive\n",
    "#                     score = score_func(score, len(seq))\n",
    "                    curr_beam.add(score, (char==end_tok), seq+[char])\n",
    "        \n",
    "        (best_score, best_complete, best_seq) = max(curr_beam)\n",
    "        if best_complete == True: return (best_seq[1:], best_score)   # returns first complete beam not best...\n",
    "            \n",
    "        prev_beam = curr_beam\n",
    "        \n",
    "    (best_score, best_complete, best_seq) = max(curr_beam)\n",
    "    return (best_seq[1:], best_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def beam_decode(net, src, beam_width, seq_len):\n",
    "    net.eval()\n",
    "    with torch.no_grad():\n",
    "        feats = net.transformer.encode(net.img_enc(src))        \n",
    "        return beamsearch(partial(prob_func, net=net, feats=feats, beam_width=beam_width), seq_len, beam_width)\n",
    "    \n",
    "def prob_func(tgt, net=None, feats=None, beam_width=5):\n",
    "    tgt = torch.tensor([tgt], dtype=torch.long, device=device)\n",
    "    mask = subsequent_mask(tgt.size(-1))\n",
    "    dec_outs = net.transformer.decode(feats, tgt, mask)\n",
    "    logits = net.transformer.generate(dec_outs[:,-1])\n",
    "    \n",
    "    log_probs = logits - torch.logsumexp(logits, 1)  # more numerically stable\n",
    "    # log_probs = F.softmax(logits, -1).log()\n",
    "    \n",
    "    return torch.topk(log_probs.squeeze(0), beam_width, dim=-1)\n",
    "#     return zip(res[0][0].detach(),res[1][0].detach())\n",
    "\n",
    "def score_func(log_probs, i, alpha=0.6):\n",
    "    length_penalty = math.pow((5 + i)/6, alpha)\n",
    "    return log_probs/length_penalty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx = 2\n",
    "x1 = x[idx][None]\n",
    "y1 = y[idx][None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "b_res, score = beam_decode(learn.model, image, 3, seq_len)    #294, 3m18s\n",
    "# 294 - 1m40s\n",
    "# 294 - 1m45s (w/ score_func)\n",
    "# 295 - 22s (beam_width=1 ~ greedy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "r = torch.tensor([b_res], dtype=torch.long, device=device)\n",
    "p = char_label_text(r)\n",
    "\n",
    "_cer(truth, p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# valid\n",
    "p = char_label_text(v_res[idx][None])\n",
    "t = char_label_text(y1[0])\n",
    "_cer(t,p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# greedy\n",
    "p = char_label_text(g_res[idx][None])\n",
    "t = char_label_text(y1[0])\n",
    "_cer(t,p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# beam (sz=3)\n",
    "r = torch.tensor([b_res], dtype=torch.long, device=device)\n",
    "p = char_label_text(r)\n",
    "t = char_label_text(y1[0])\n",
    "_cer(t,p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "stoi = {k:i for i,k in enumerate(itos)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "print(char_label_text(g_res[idx][None]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "st = ''.join([itos[i] for i in b_res])\n",
    "p = '\\n'.join(textwrap.wrap(st, 70))\n",
    "show_img(denorm(x1)[0], figsize=(10,10), title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Source Attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx = 0\n",
    "img = imgs[idx]\n",
    "\n",
    "v_chars = v_res[idx]\n",
    "v_attns = to_np(torch_scale_attns(v_attn)[idx])\n",
    "\n",
    "g_chars = g_res[idx]\n",
    "g_attns = to_np(torch_scale_attns(g_attn)[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#valid\n",
    "fig, axes = plt.subplots(5,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    a = g_filter(vv_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[v_chars[i].item()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(6,4, gridspec_kw={'hspace': 0.3}, figsize=(20, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    a = g_filter(g_attns[i])\n",
    "    ax.imshow(img, alpha=None)\n",
    "    ax.imshow(a, cmap='Blues', interpolation='nearest', alpha=0.3)\n",
    "    ax.set_title(itos[g_chars[i].item()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Attention Visualizations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def transparent_cmap(cmap, N=5):\n",
    "    \"Copy colormap and set alpha values\"\n",
    "    mycmap = matplotlib.cm.get_cmap(cmap, N)\n",
    "    mycmap._init()\n",
    "    mycmap._lut[:,-1] = np.linspace(0, 0.6, N+3)\n",
    "    return mycmap\n",
    "\n",
    "#Use base cmap to create transparent\n",
    "# mycmap = transparent_cmap(plt.cm.Reds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def show_attn(img, attns, chars, ax, color, showChars=True):\n",
    "    for i in range(attns.shape[0]):\n",
    "        c = chars[i].item()\n",
    "        if c not in [0,1,2,3]:\n",
    "            a = g_filter(attns[i])\n",
    "            y,x = scipy.ndimage.center_of_mass(a)\n",
    "            #sns.heatmap(a, cmap=mycmap, cbar=False, ax=ax)\n",
    "            ax.imshow(a, cmap=transparent_cmap(color), interpolation='nearest')\n",
    "            if showChars: ax.text(x-8,y-10,itos[c], fontsize=15)\n",
    "\n",
    "    ax.set_title(char_label_text(chars))\n",
    "    ax.imshow(img, alpha=0.6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def thresh_attn(attn, thresh=0.1):\n",
    "    zeros = torch.zeros_like(attn)\n",
    "    new = torch.where(attn >= thresh, attn, zeros)\n",
    "    \n",
    "    # some attns will not have a value over the thresh in which case\n",
    "    # we need to insert top k value at appropriate index\n",
    "    vals, idxs = torch.topk(attn, 1, dim=-1)\n",
    "    \n",
    "    # reshape\n",
    "    flat_new = new.flatten(0,1)\n",
    "    vals = vals.flatten()\n",
    "    idxs = idxs.flatten()\n",
    "    \n",
    "    for i in range(flat_new.size(0)):\n",
    "        flat_new[i,idxs[i]] = vals[i]\n",
    "\n",
    "    new = flat_new.view_as(new)\n",
    "    return new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Decoder Self-Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sns.set_context(context=\"notebook\")\n",
    "\n",
    "def draw(data, x, y, ax):\n",
    "    return sns.heatmap(data, xticklabels=x, square=True, yticklabels=y, vmin=0.0, vmax=1.0,\n",
    "                       cmap='YlOrRd', linewidths=0.05, cbar=False, ax=ax)\n",
    "\n",
    "for layer in range(4):\n",
    "    print(\"Decoder Self-Attention Layer\", layer+1)\n",
    "\n",
    "    fig, axes = plt.subplots(1,4, figsize=(20, 10))\n",
    "    for i,ax in enumerate(axes.flat):\n",
    "        # greedy decoding (no access to true values)\n",
    "        pred = char_split_text(g_res[i])[20:40]\n",
    "        shifted_y = rshift(g_res.float()).long()\n",
    "        true = char_split_text(shifted_y[i])[20:40]\n",
    "        g = draw(self_attn(layer)[i].data[20:40, 20:40], true, pred, ax=ax)\n",
    "        g.set_yticklabels(g.get_yticklabels(), rotation=0) \n",
    "        g.set_xticklabels(g.get_xticklabels(), rotation=0) \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Decoder Source-Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1,4, gridspec_kw={'hspace': 0.5}, figsize=(20, 10))\n",
    "    \n",
    "for idx in range(len(axes.flat)//4):\n",
    "    img = imgs[idx]\n",
    "    g_chars = g_res[idx]\n",
    "    \n",
    "    # 4 attn layers\n",
    "    for h in range(4):\n",
    "        attn = source_attn(h)\n",
    "        g_attns = to_np(torch_scale_attns(attn)[idx])\n",
    "\n",
    "        show_attn(img, g_attns, g_chars, axes[idx,h], 'YlGn', showChars=False)\n",
    "        axes[idx,h].set_title(f'layer {h+1}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Validation vs Greedy (final layer src-attn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "v_scaled_attns = torch_scale_attns(thresh_attn(v_attn))\n",
    "g_scaled_attns = torch_scale_attns(thresh_attn(g_attn))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2,2, gridspec_kw={'hspace': 0.5}, figsize=(20, 20))\n",
    "for idx in range(len(axes.flat)//2):\n",
    "    img = imgs[idx]\n",
    "\n",
    "    v_chars = v_res[idx]\n",
    "    v_attns = to_np(v_scaled_attns[idx])\n",
    "\n",
    "    g_chars = g_res[idx]\n",
    "    g_attns = to_np(g_scaled_attns[idx])\n",
    "    \n",
    "    # valid\n",
    "    show_attn(img, v_attns, v_chars, axes[idx,0], 'YlOrRd')\n",
    "    # greedy\n",
    "    show_attn(img, g_attns, g_chars, axes[idx,1], 'YlGn')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Individual Examination"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx=0\n",
    "img = imgs[idx]\n",
    "\n",
    "g_chars = g_res[idx]\n",
    "g_scaled_attns = torch_scale_attns(thresh_attn(g_attn)[0:1])  # passing in bs of 1\n",
    "g_attns = to_np(g_scaled_attns[0])  # removing bs\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(20, 20))\n",
    "show_attn(img, g_attns, g_chars, ax, 'YlGn')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sns.set_context(context=\"notebook\")\n",
    "\n",
    "def draw(data, x, y, ax):\n",
    "    mask = np.zeros_like(data)\n",
    "    mask[np.triu_indices_from(mask, k=1)] = True\n",
    "    return sns.heatmap(data, xticklabels=x, square=True, yticklabels=y, vmin=0.0, vmax=1.0,\n",
    "                       mask=mask, cmap='YlOrRd', linewidths=0.05, cbar=False, ax=ax)\n",
    "\n",
    "for layer in range(4):\n",
    "    print(\"Decoder Self-Attention Layer\", layer+1)\n",
    "\n",
    "    fig, ax = plt.subplots(1,1, figsize=(20, 10))\n",
    "    i = 1\n",
    "    # greedy decoding (no access to true values)\n",
    "    pred = char_split_text(g_res[i])[230:280]\n",
    "    shifted_y = rshift(g_res.float()).long()\n",
    "    true = char_split_text(shifted_y[i])[230:280]\n",
    "    g = draw(self_attn(layer)[i].data[230:280, 230:280], true, pred, ax=ax)\n",
    "    g.set_yticklabels(g.get_yticklabels(), rotation=0) \n",
    "    g.set_xticklabels(g.get_xticklabels(), rotation=0) \n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Memory utility methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# prints currently alive Variables     # Tensors and \n",
    "import gc\n",
    "for obj in gc.get_objects():\n",
    "    try:\n",
    "        if (hasattr(obj, 'data') and torch.is_tensor(obj.data)):     # torch.is_tensor(obj) or \n",
    "            print(obj.name, type(obj), obj.size())\n",
    "    except:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "collapsed": true,
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def tensor_size(tensor):\n",
    "    return (tensor.element_size() * tensor.v_res.nelement())\n",
    "\n",
    "tensor_size(v_res)"
   ]
  }
 ],
 "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.5"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}