1--Lines.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",
    "from fastai.callbacks.tracker 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')"
   ]
  },
  {
   "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": "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",
    "    ax.imshow(image2np(im.data), alpha=alpha)\n",
    "    if title: ax.set_title(title)\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss, Metrics, Callbacks"
   ]
  },
  {
   "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 = self.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\n",
    "    \n",
    "    def loss_prep(self, pred, target):\n",
    "        \"equalize input/target sl; combine bs/sl dimensions\"\n",
    "        bs,tsl = target.shape\n",
    "        _ ,sl,vocab = pred.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: pred = F.pad(pred, (0,0,0,tsl-sl))\n",
    "        # not ideal => adds 96 logits all 0s...\n",
    "\n",
    "        targ = target.contiguous().view(-1).long()\n",
    "        pred = pred.contiguous().view(-1, vocab)\n",
    "        return pred, targ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def cer(preds, targs):\n",
    "    bs = targs.size(0)\n",
    "    res = torch.argmax(preds, dim=-1)\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) or 1)\n",
    "    return error, bs\n",
    "\n",
    "def char_label_text(pred, sep=''):\n",
    "    ints = to_np(pred).astype(int)\n",
    "    nonzero = ints[np.nonzero(ints)] #[:-1]  #remove eos token\n",
    "    return sep.join([itos[i] for i in nonzero])"
   ]
  },
  {
   "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):\n",
    "        super().__init__()\n",
    "        self.name = 'cer'\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 = 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)"
   ]
  },
  {
   "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((1,size,size), device=device).byte())\n",
    "    #return torch.tril(torch.ones((1,1,size,size), device=device).byte())  # complex batches\n",
    "    \n",
    "def parallelogram_mask(size, diagonal):\n",
    "    mask = torch.ones((1,size,size), device=device).byte()\n",
    "    upper = torch.tril(mask).bool()\n",
    "    lower = torch.triu(mask, diagonal=-diagonal).bool()\n",
    "    return (upper & lower).byte()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class TeacherForce(LearnerCallback):\n",
    "    def __init__(self, learn:Learner):\n",
    "        super().__init__(learn)\n",
    "        \n",
    "    def on_batch_begin(self, last_input, last_target, **kwargs):\n",
    "        s = rshift(last_target).long()\n",
    "        mask = subsequent_mask(s.size(-1))\n",
    "        #mask = parallelogram_mask(s.size(-1), 10)\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": [
    "## Lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "line_labels = pd.read_csv(f'{PATH}/ascii/lines.txt', names=['filename','result','value'], escapechar='\\\\',\n",
    "                          delim_whitespace=True, skiprows=23, header=None, usecols=[0,1,8])\n",
    "line_labels['text'] = line_labels.apply(lambda row: row.value.replace('|', ' '), axis=1)\n",
    "line_labels.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def cleanup(x):\n",
    "    x = re.sub(r\" \\'(\\w+)\", lambda x: x.group(0)[1:], x)     # match all space and apostrophe letters\n",
    "    x = re.sub(r\" [.,?!:;()]\", lambda x: x.group(0)[1:], x)  # match all space and punctuation\n",
    "    return x\n",
    "\n",
    "line_labels['text'] = line_labels.apply(lambda row: cleanup(row.text), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "df = line_labels[['filename', 'text']]\n",
    "# df['filename'] = df['filename'].apply(lambda x: f\"{x}.png\")\n",
    "df['filename'] = df['filename']+'.png'\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "FOLDER = 'lines'\n",
    "sz,bs = (64,512), 50#100\n",
    "# sz,bs = (96,1024), 100\n",
    "seq_len = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## paragraphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fname = 'edited_pg.csv' #'small_synth_words.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,10\n",
    "seq_len = 700"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## sm 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": [
    "num_lines = df.label.apply(lambda x: len(x.split('\\n')))\n",
    "df['num_lines'] = num_lines\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# df = df[:20000]\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",
    "num_lines,seq_len = 4,100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## combo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fname = 'combo_145k.csv'\n",
    "FOLDER = 'combo_cat'\n",
    "\n",
    "CSV = PATH/fname\n",
    "df = pd.read_csv(CSV)\n",
    "\n",
    "sz,bs = 512,10\n",
    "seq_len = 750\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# 6 and fewer\n",
    "df = df[df.num_lines <= 6]\n",
    "sz,bs = 512,30\n",
    "seq_len = 600\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# 6 and greater\n",
    "df = df[df.num_lines >= 6]\n",
    "sz,bs = 512,10\n",
    "seq_len = 750\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## test combo (no fonts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "FOLDER = 'test_combo'\n",
    "CSV = PATH/'test_combo.csv'\n",
    "df = pd.read_csv(CSV)\n",
    "\n",
    "sz,bs = 512,10\n",
    "seq_len = 750\n",
    "len(df)"
   ]
  },
  {
   "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=0, max_warp=0.1)\n",
    "\n",
    "def force_gray(image): return image.convert('L').convert('RGB')\n",
    "\n",
    "def label_collater(samples:BatchSamples, pad_idx:int=0):\n",
    "    \"Function that collect samples and pads ends 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+1).long() + pad_idx  # add 1 to max_len to account for bos token\n",
    "    for i,lbl in enumerate(lbls):\n",
    "        labels[i,:len(lbl)] = torch.from_numpy(lbl)  #padding end    \n",
    "    return imgs, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "itos = pickle.load(open(PATH/'itos.pkl', 'rb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class CharTokenizer(BaseTokenizer):\n",
    "    def tokenizer(self, t:str) -> List[str]: return list(t)+['xxeos']\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",
    "        return sep.join([self.itos[i] for i in nums]) if sep is not None else [self.itos[i] for i in nums]\n",
    "\n",
    "class SequenceList(TextList):    \n",
    "    def __init__(self, items:Iterator, vocab:Vocab, **kwargs):\n",
    "        toknizr = Tokenizer(tok_func=CharTokenizer, pre_rules=[], post_rules=[], special_cases=[])\n",
    "        procs = [TokenizeProcessor(tokenizer=toknizr, include_bos=False),\n",
    "                 NumericalizeProcessor(vocab=vocab)]\n",
    "        super().__init__(items, vocab, sep='', pad_idx=0, processor=procs)\n",
    "    \n",
    "    def analyze_pred(self, pred:Tensor):\n",
    "        return torch.argmax(pred, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\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",
    "        .label_from_df(label_cls=SequenceList, vocab=CharVocab(itos))\n",
    "        .transform([], size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        #.transform(tfms, size=sz, resize_method=ResizeMethod.PAD, padding_mode='border')\n",
    "        # maintains aspect ratio but too small for good results => mostly whitespace\n",
    "        .databunch(bs=bs, device=device, collate_fn=label_collater)\n",
    "        #.normalize() # this sets x values to an odd range (~.3,-6)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=2, ds_type=DatasetType.Train, figsize=(18,10))"
   ]
  },
  {
   "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 = nn.LayerNorm\n",
    "LayerNorm = partial(nn.LayerNorm, eps=1e-4)  # accomodates mixed precision training\n",
    "# LayerNorm = partial(nn.BatchNorm2d, eps=1e-4)"
   ]
  },
  {
   "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))  # acts as a weak LM\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 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",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.w_2(self.dropout(F.gelu(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": {},
   "source": [
    "# img2seq2seq 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, src_adapt, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_adapt = src_adapt\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), self.embed(tgt), tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_adapt(src))\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(tgt, src, tgt_mask)\n",
    "    \n",
    "    def embed(self, tgt):\n",
    "        return self.tgt_embed(tgt)\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 fastai.vision.models.cadene_models import se_resnet50"
   ]
  },
  {
   "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):\n",
    "        super(ResnetBase, self).__init__()\n",
    "        \n",
    "        slices = {512:-4, 1024:-3, 2048:-2}\n",
    "        #512, 8, 64\n",
    "        #1024, 4, 32\n",
    "        #2048, 2, 16\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = se_resnet50(True)\n",
    "        modules = list(net.children())[:s]\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)   \n",
    "        x = self.pool(self.conv(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, em_sz, N=4, drops=0, attn_type='multi', heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    pos = PositionalEncoding(d_model, drops, 2000)\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",
    "            nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8)) #, pos\n",
    "        ),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), pos\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\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):\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, seq_len=700):\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(seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "#                     mask = parallelogram_mask(tgt.size(-1), 10)\n",
    "\n",
    "                    t = self.transformer.embed(tgt)\n",
    "    \n",
    "#                     dec_outs = self.transformer.decode(feats, t[:,-11:])\n",
    "                    dec_outs = self.transformer.decode(feats, t, mask)\n",
    "\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",
    "            seq_outs = self.t2(dec_outs, tgt, tgt_mask)\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": [
    "def init_params(model):\n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    transformer = make_full_model(len(itos), d_model, em_sz, N=N, drops=drops, attn_type=attn_type, heads=heads)\n",
    "    transformer.apply(init_params)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=smoothing),\n",
    "                   metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# SEResnet50 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, src_adapt, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_adapt = src_adapt\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), self.embed(tgt), tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_adapt(src))\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(tgt, src, tgt_mask)\n",
    "    \n",
    "    def embed(self, tgt):\n",
    "        return self.tgt_embed(tgt)\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 fastai.vision.models.cadene_models import se_resnet50"
   ]
  },
  {
   "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):\n",
    "        super(ResnetBase, self).__init__()\n",
    "        \n",
    "        slices = {512:-4, 1024:-3, 2048:-2}\n",
    "        #512, 8, 64\n",
    "        #1024, 4, 32\n",
    "        #2048, 2, 16\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = se_resnet50(True)\n",
    "        modules = list(net.children())[:s]    #512,2,16;  256,4,32;  128,8,64\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "\n",
    "\n",
    "#         self.block1=Block(728,1024,2,(2,1),start_with_relu=True,grow_first=True)\n",
    "#         self.block2=Block(1024,1536,2,(2,1),start_with_relu=True,grow_first=True)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        \n",
    "        x = self.pool(self.conv(x))\n",
    "#         x = self.block2(self.block1(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "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):\n",
    "        super().__init__()\n",
    "        \n",
    "        slices = {512:-4, 1024:-3, 2048:-2}\n",
    "#         slices = {128: -4, 256: -3, 512: -2}\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = se_resnet50(True)\n",
    "        modules = list(net.children())[:s]   #512,2,16;  256,4,32;  128,8,64\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "        self.c1 = conv_layer(em_sz,em_sz,stride=(2,1)) #8,32\n",
    "        self.b1 = res_block(em_sz)\n",
    "        self.c2 = conv_layer(em_sz,em_sz,stride=(2,1)) #4,32\n",
    "        self.b2 = res_block(em_sz)\n",
    "        self.c3 = conv_layer(em_sz,em_sz,stride=(2,1)) #2,32\n",
    "        self.b3 = res_block(em_sz)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = self.pool(self.conv(x))\n",
    "        \n",
    "        x1 = self.b1(self.c1(x))\n",
    "        x2 = self.b2(self.c2(x1))\n",
    "        x3 = self.b3(self.c3(x2))\n",
    "        \n",
    "        out = torch.cat([x, x1, x2, x3], dim=2)\n",
    "        return out.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, em_sz, N=4, drops=0, attn_type='multi', heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    pos = PositionalEncoding(d_model, drops, 2000)\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",
    "            nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8)) #, pos\n",
    "        ),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), pos\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\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):\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, seq_len=700):\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(seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "#                     mask = parallelogram_mask(tgt.size(-1), 10)\n",
    "\n",
    "                    t = self.transformer.embed(tgt)\n",
    "    \n",
    "#                     dec_outs = self.transformer.decode(feats, t[:,-11:])\n",
    "                    dec_outs = self.transformer.decode(feats, t, mask)\n",
    "\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": [
    "def init_params(model):\n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    img_encoder = ResnetBase(em_sz)\n",
    "    transformer = make_full_model(len(itos), d_model, em_sz, N=N, drops=drops, attn_type=attn_type, heads=heads)\n",
    "    transformer.apply(init_params)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=smoothing),\n",
    "                   metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Xception 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, src_adapt, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_adapt = src_adapt\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), self.embed(tgt), tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_adapt(src))\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(tgt, src, tgt_mask)\n",
    "    \n",
    "    def embed(self, tgt):\n",
    "        return self.tgt_embed(tgt)\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 fastai.vision.models.cadene_models import xception_cadene"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class SeparableConv2d(nn.Module):\n",
    "    def __init__(self,in_channels,out_channels,kernel_size=1,stride=1,padding=0,dilation=1,bias=False):\n",
    "        super(SeparableConv2d,self).__init__()\n",
    "\n",
    "        self.conv1 = nn.Conv2d(in_channels,in_channels,kernel_size,stride,\n",
    "                               padding,dilation,groups=in_channels,bias=bias)\n",
    "        self.pointwise = nn.Conv2d(in_channels,out_channels,1,1,0,1,1,bias=bias)\n",
    "\n",
    "    def forward(self,x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.pointwise(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    def __init__(self,in_filters,out_filters,reps,strides=1,start_with_relu=True,grow_first=True):\n",
    "        super(Block, self).__init__()\n",
    "\n",
    "        if out_filters != in_filters or strides!=1:\n",
    "            self.skip = nn.Conv2d(in_filters,out_filters,1,stride=strides, bias=False)\n",
    "            self.skipbn = nn.BatchNorm2d(out_filters)\n",
    "        else:\n",
    "            self.skip=None\n",
    "\n",
    "        rep=[]\n",
    "\n",
    "        filters=in_filters\n",
    "        if grow_first:\n",
    "            rep.append(nn.ReLU(inplace=True))\n",
    "            rep.append(SeparableConv2d(in_filters,out_filters,3,stride=1,padding=1,bias=False))\n",
    "            rep.append(nn.BatchNorm2d(out_filters))\n",
    "            filters = out_filters\n",
    "\n",
    "        for i in range(reps-1):\n",
    "            rep.append(nn.ReLU(inplace=True))\n",
    "            rep.append(SeparableConv2d(filters,filters,3,stride=1,padding=1,bias=False))\n",
    "            rep.append(nn.BatchNorm2d(filters))\n",
    "\n",
    "        if not grow_first:\n",
    "            rep.append(nn.ReLU(inplace=True))\n",
    "            rep.append(SeparableConv2d(in_filters,out_filters,3,stride=1,padding=1,bias=False))\n",
    "            rep.append(nn.BatchNorm2d(out_filters))\n",
    "\n",
    "        if not start_with_relu:\n",
    "            rep = rep[1:]\n",
    "        else:\n",
    "            rep[0] = nn.ReLU(inplace=False)\n",
    "\n",
    "        if strides != 1:\n",
    "            rep.append(nn.MaxPool2d(3,strides,1))\n",
    "        self.rep = nn.Sequential(*rep)\n",
    "\n",
    "    def forward(self,inp):\n",
    "        x = self.rep(inp)\n",
    "\n",
    "        if self.skip is not None:\n",
    "            skip = self.skip(inp)\n",
    "            skip = self.skipbn(skip)\n",
    "        else:\n",
    "            skip = inp\n",
    "\n",
    "        x+=skip\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class XceptionBase(nn.Module):\n",
    "    def __init__(self, em_sz):\n",
    "        super(XceptionBase, self).__init__()\n",
    "        \n",
    "        slices = {728:-6, 1024:-5, 2048:-1}\n",
    "        #728, 4, 32\n",
    "        #1024, 2, 16\n",
    "        #2048, 2, 16\n",
    "        s = slices[em_sz]\n",
    "\n",
    "        net = xception_cadene(True)\n",
    "        self.base = net[:s]\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "\n",
    "#         self.block1=Block(728,1024,2,(2,1),start_with_relu=True,grow_first=True)\n",
    "#         self.block2=Block(1024,1536,2,(2,1),start_with_relu=True,grow_first=True)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        \n",
    "        x = self.pool(self.conv(x))\n",
    "#         x = self.block2(self.block1(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, em_sz, N=4, drops=0, attn_type='multi', heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    pos = PositionalEncoding(d_model, drops, 2000)\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",
    "            nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8)) #, pos\n",
    "        ),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), pos\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\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):\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, seq_len=700):\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(seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "#                     mask = parallelogram_mask(tgt.size(-1), 10)\n",
    "\n",
    "                    t = self.transformer.embed(tgt)\n",
    "    \n",
    "#                     dec_outs = self.transformer.decode(feats, t[:,-11:])\n",
    "                    dec_outs = self.transformer.decode(feats, t, mask)\n",
    "\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": [
    "def init_params(model):\n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    img_encoder = XceptionBase(em_sz)\n",
    "    transformer = make_full_model(len(itos), d_model, em_sz, N=N, drops=drops, attn_type=attn_type, heads=heads)\n",
    "    transformer.apply(init_params)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=smoothing),\n",
    "                   metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Multi-Resolution 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, src_adapt, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_adapt = src_adapt\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), self.embed(tgt), tgt_mask)\n",
    "    \n",
    "    def encode(self, src):\n",
    "        return self.encoder(self.src_adapt(src))\n",
    "    \n",
    "    def decode(self, src, tgt, tgt_mask=None):\n",
    "        return self.decoder(tgt, src, tgt_mask)\n",
    "    \n",
    "    def embed(self, tgt):\n",
    "        return self.tgt_embed(tgt)\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]   #512,2,16;  256,4,32;  128,8,64\n",
    "        self.base = nn.Sequential(*modules)\n",
    "        \n",
    "        self.conv = conv_layer(em_sz,em_sz)\n",
    "        self.pool = nn.AdaptiveMaxPool2d((16,None))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = self.pool(self.conv(x))\n",
    "        return x.flatten(2,3).permute(0,2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_full_model(vocab, d_model, em_sz, N=4, drops=0, attn_type='multi', heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    pos = PositionalEncoding(d_model, drops, 2000)\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",
    "            nn.Linear(em_sz, d_model), Lambda(lambda x: x.mul_(8)) #, pos\n",
    "        ),\n",
    "        nn.Sequential(\n",
    "            Embeddings(d_model, vocab), pos\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\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):\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, seq_len=700):\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(seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "#                     mask = parallelogram_mask(tgt.size(-1), 10)\n",
    "\n",
    "                    t = self.transformer.embed(tgt)\n",
    "    \n",
    "#                     dec_outs = self.transformer.decode(feats, t[:,-11:])\n",
    "                    dec_outs = self.transformer.decode(feats, t, mask)\n",
    "\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": [
    "def init_params(model):\n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.1, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    img_encoder = ResnetBase(em_sz, d_model)\n",
    "    transformer = make_full_model(len(itos), d_model, em_sz, N=N, drops=drops, attn_type=attn_type, heads=heads)\n",
    "    transformer.apply(init_params)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=smoothing),\n",
    "                   metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experimentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 256, N=4, drops=0, attn_type='multi', heads=8)\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(5, max_lr=1e-3)\n",
    "#### Paragraphs ####\n",
    "# sz:512, bs:10, resnet34, 5cycle(1e-3), N:4, multi:8\n",
    "# 713.113770\t723.217712\t0.640478\t01:21   em_sz:256, 32x32\n",
    "# 722.943359\t732.392761\t0.650870\t01:00   em_sz:512, 16x16\n",
    "# -----      CUDA memory error     ------       em_sz:128, 64x64             \n",
    "\n",
    "# 722.326721\t732.504883\t0.651015\t01:01   em_sz:128, conv/pool(8),  8x64\n",
    "# 695.523865\t704.048645\t0.623089\t01:18   em_sz:128, conv/pool(16), 16x64\n",
    "# 685.498596\t694.236023\t0.614935\t01:02   em_sz:256, conv/pool(16), 16x32  **\n",
    "# 694.890808\t703.071411\t0.623319\t00:58   em_sz:512, conv/pool(16), 16x16\n",
    "\n",
    "# 681.818176\t690.215942\t0.612152\t01:26   se_resnet50, em_sz:1024, conv/pool(16), 16x32 **\n",
    "# 684.835876\t694.623779\t0.614549\t01:21   resnet50, em_sz:1024, conv/pool(16), 16x32\n",
    "# 699.177185\t708.581848\t0.627579\t01:47   xception, em_sz:728, conv/pool(16), 16x32\n",
    "\n",
    "# 697.656616\t698.180603\t0.618370\t01:17   em_sz:256, conv/pool(16), 16x32, N:6\n",
    "# 698.445801\t712.330444\t0.630357\t01:10   em_sz:256, conv/pool(16), 16x32, multi:16\n",
    "\n",
    "# 693.106262\t698.929504\t0.619492\t01:25   em_sz:256, conv/pool(16), 4res w/ blocks:(16,8,4,2x32)\n",
    "\n",
    "#### LINES ####\n",
    "# sz:(48,512), bs:100, resnet34, 5cycle(1e-3), em_sz:512, N:4, multi:8\n",
    "\n",
    "# 16.167061\t19.248831\t0.130215   maxpool, linear+mul_8   'lines1'\n",
    "# 16.855366\t19.663122\t0.134168\n",
    "# 27.059019\t27.176353\t0.196770   maxpool, ff\n",
    "# 19.401184\t21.343414\t0.147686   maxpool, linear+pos\n",
    "# 20.828600\t22.394470\t0.154748   maxpool, linear\n",
    "# 19.685211\t21.401609\t0.148922   no maxpool, no linear+mul_8   'lines2'\n",
    "# 15.595810\t18.914587\t0.129315   no maxpool, linear+mul_8   'lines3'  ---  linear+mul_8 crucial!\n",
    "\n",
    "# 12.375642\t14.939847\t0.103393   em_sz:256, cat convs(3,2,1 x 32)  'lines4'\n",
    "# 17.737593\t20.629824\t0.142271   em_sz:512, cat convs(2,1 x 16)\n",
    "# 8.757012\t11.217060\t0.084667\t07:18   em_sz:128, bs:25, 10 resolutions   'lines5'\n",
    "# 9.839996\t12.166018\t0.090686\t03:09   em_sz:128, bs:25, 4 resolutions(6,3,2,1)    'lines6'\n",
    "# 18.749489\t19.622530\t0.141156\t02:20   \"\", d_model:128\n",
    "\n",
    "# sz:(64,512), bs:50\n",
    "# 10.431609\t12.825503\t0.092731\t01:30   em_sz:128, 4 resolutions(8,4,2,1) of 1 line\n",
    "# 11.442437\t13.249302\t0.096183\t03:35   em_sz:128, 4 resolutions(8,4,2,1)\n",
    "# greedy:    2.99910   .21647\n",
    "\n",
    "# bs: 50\n",
    "# 6.106063\t9.794361\t0.070278\t01:10   em_sz: 256, 4x32\n",
    "# 12.173990\t14.315442\t0.096518\t01:03   \"\", bs:100\n",
    "# 7.294389\t10.227886\t0.072702\t01:29   resnet50, em_sz: 1024, 4x32\n",
    "\n",
    "# sz:(64,512), bs:100\n",
    "# 10.622000\t13.152803\t0.091343\t01:21   em_sz:1024, 1 resolution: (1x64)   'lines7'\n",
    "# 15.201996\t17.066908\t0.121637\t00:51   pool(32,None) before base, em_sz:512, 1x64\n",
    "# 12.205278\t14.841956\t0.104063\t01:20   pool(48,None) before base, em_sz:1024, 1x64\n",
    "# 12.217975\t13.971797\t0.097959\t00:57   pool(4,None) after base, em_sz:512, 1x64\n",
    "# 10.639203\t13.198123\t0.092142\t01:22   pool(8,None) after base, em_sz:1024, 1x64\n",
    "# 10.091258\t12.612871\t0.088563\t01:36   \"\", N:6    'lines8'\n",
    "# 8.227972\t12.432637\t0.087113\t01:35   2nd run, 5cycle(4e-6)\n",
    "# 10.719064\t13.117971\t0.090754\t01:50   \"\", N:8\n",
    "# 11.734365\t13.847672\t0.096678\t01:09   pool(8,None) after, em_sz:512, N:4, 3 resolutions\n",
    "# 10.949694\t13.447043\t0.095033\t01:11   \"\", em_sz:1024\n",
    "# 11.054914\t13.662763\t0.096324\t01:13   \"\", padding rather than stride 4->2\n",
    "# 11.858457\t13.944701\t0.097301\t00:59   \"\", last conv k=2, 1x63\n",
    "# 10.656734\t13.374981\t0.093117\t01:23   \"\", last conv k=4, 1x63\n",
    "# 9.708715\t12.423415\t0.085666\t02:45   pool(8,64) after base, em_sz:1024, w/res_blocks, 1x64   'lines9'\n",
    "# greedy:    1.28441   .12355\n",
    "\n",
    "# bs: 50\n",
    "# xception\n",
    "# 13.364486\t17.758265\t0.127069\t01:37   xception(2048) 2x16\n",
    "# 5.375062\t10.732620\t0.077556\t01:49   xception(728) 4x32\n",
    "# 7.213572\t12.425294\t0.091167\t01:37   xception(728) 1x32 (2 conv layers)\n",
    "# 6.362609\t11.896906\t0.086484\t01:39   xception(728) 1x32 (2 blocks)\n",
    "# 10.085477\t12.171700\t0.087609\t03:12   custom xception(256) 8x64\n",
    "\n",
    "# se_resnet50\n",
    "# 5.696907\t9.483581\t0.067969\t01:31   em_sz:1024, 4x32\n",
    "# 9.689358\t12.082891\t0.088165\t02:47   em_sz:512, 8x64\n",
    "# 7.772936\t12.125853\t0.083637\t01:22   em_sz:2048, 2x16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.load('lines9'); None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Image PreProcessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class ImageProcessor(nn.Module):\n",
    "    def __init__(self, d_model, max_lines=14):\n",
    "        super().__init__()\n",
    "\n",
    "        net = models.resnet34(True)\n",
    "        modules = list(net.children())[:-2]\n",
    "        \n",
    "        self.base = nn.Sequential(*modules)\n",
    "        self.head = create_head(1024, max_lines)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        with torch.no_grad():\n",
    "            base = self.base(x)\n",
    "            n_lines = self.head(base)\n",
    "            \n",
    "            b,c,h,w = base.shape\n",
    "            #res = torch.zeros((b,c,longest))\n",
    "        \n",
    "            preds = torch.argmax(n_lines, dim=-1)\n",
    "            lines = torch.unique(preds)\n",
    "            \n",
    "            longest = lines[-1]*w\n",
    "            res = torch.zeros((b,c,longest), device=device)\n",
    "            for l in lines:\n",
    "                idxs = (preds==l).nonzero()[:,0]\n",
    "                outs = F.adaptive_max_pool2d(base[idxs], (l,None)).flatten(2,3)\n",
    "                # could add positional encoding here\n",
    "                pad = longest - outs.size(-1)\n",
    "                outs = F.pad(outs, (0,pad), \"constant\", 0)\n",
    "                res[idxs] = outs\n",
    "                \n",
    "        return res.permute(0,2,1)"
   ]
  },
  {
   "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, attn_type='multi', heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "    pos = PositionalEncoding(d_model, drops, 2000)\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), pos\n",
    "        ),\n",
    "        nn.Linear(d_model, vocab),\n",
    "    )\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):\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, seq_len=700):\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(seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "#                     mask = parallelogram_mask(tgt.size(-1), 10)\n",
    "\n",
    "                    t = self.transformer.embed(tgt)\n",
    "    \n",
    "#                     dec_outs = self.transformer.decode(feats, t[:,-11:])\n",
    "                    dec_outs = self.transformer.decode(feats, t, mask)\n",
    "\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": [
    "def init_params(model):\n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, N=4, drops=0.2, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    img_encoder = ImageProcessor(d_model)\n",
    "    transformer = make_full_model(len(itos), d_model, N=N, drops=drops, attn_type=attn_type, heads=heads)\n",
    "    transformer.apply(init_params)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=smoothing),\n",
    "                   metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## Combine/load state_dict "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, N=6, drops=0, attn_type='multi', heads=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "new_sd = OrderedDict()\n",
    "\n",
    "line_sd = torch.load(PATH/'models/line_98acc.pth', map_location=device)\n",
    "\n",
    "for k, v in line_sd['model'].items():\n",
    "    if k.startswith('0'):\n",
    "        name = re.sub(\"^(0)\",'img_enc.base',k)\n",
    "        new_sd[name] = v\n",
    "    if k.startswith('1'):\n",
    "        name = re.sub(\"^(1)\",'img_enc.head',k)\n",
    "        new_sd[name] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.model.load_state_dict(new_sd, strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# sd = torch.load(PATH/'models/combo_512_mix_res9.pth', map_location=device)\n",
    "\n",
    "# for k, v in sd['model'].items():\n",
    "#     if not k.startswith('img_enc'):\n",
    "#         new_sd[k] = v"
   ]
  },
  {
   "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",
    "        self.linear = nn.Linear(em_sz, d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        return self.linear(x) * 8"
   ]
  },
  {
   "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, attn_type='multi', heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \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",
    "    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):\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, seq_len=700):\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(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": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.2, attn_type='multi', heads=8, smoothing=0.1):\n",
    "    img_encoder = ResnetBase(em_sz, d_model)\n",
    "    transformer = make_full_model(len(itos), d_model, N=N, drops=drops, attn_type=attn_type, heads=heads)\n",
    "    transformer.apply(init_params)\n",
    "    #transformer.generator.weight = transformer.tgt_embed[0].lut.weight\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=smoothing),\n",
    "                   metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# TransformerXL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class PositionalEncoding(Module):\n",
    "    \"Encode the position with a sinusoid.\"\n",
    "    def __init__(self, d_model:int):\n",
    "        self.register_buffer('freq', 1 / (10000 ** (torch.arange(0., d_model, 2.) / d_model)))\n",
    "\n",
    "    def forward(self, pos:Tensor):\n",
    "        inp = torch.ger(pos, self.freq)  #(outer) matrix product of 2 vectors\n",
    "        enc = torch.cat([inp.sin(), inp.cos()], dim=-1)\n",
    "        return enc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class FeedForward(Module):\n",
    "    def __init__(self, d_model:int, drops=0.1):\n",
    "        self.core = nn.Sequential(\n",
    "            nn.Linear(d_model, d_model*4), nn.ReLU(inplace=True), nn.Dropout(drops),\n",
    "            nn.Linear(d_model*4, d_model), nn.Dropout(drops)\n",
    "        )\n",
    "        self.ln = nn.LayerNorm(d_model)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        out = self.core(x)\n",
    "        return self.ln(x + out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MultiHeadRelativeAttention(Module):\n",
    "    def __init__(self, n_heads:int, d_model:int, drops=0.1, bias=False):\n",
    "        d_head = d_model//n_heads\n",
    "        self.n_heads, self.d_head = n_heads, d_head\n",
    "        self.attention = nn.Linear(d_model, 3 * n_heads * d_head, bias=bias)\n",
    "        self.out = nn.Linear(n_heads * d_head, d_model, bias=bias)\n",
    "        self.drop_att,self.drop_res = nn.Dropout(drops),nn.Dropout(drops)\n",
    "        self.ln = nn.LayerNorm(d_model)\n",
    "        self.r_attn = nn.Linear(d_model, n_heads * d_head, bias=bias)\n",
    "\n",
    "    def forward(self, x:Tensor, mask:Tensor=None, **kwargs):\n",
    "        return self.ln(x + self.drop_res(self.out(self._mhra(x, mask=mask, **kwargs))))\n",
    "\n",
    "    def _mhra(self, x:Tensor, r:Tensor=None, u:Tensor=None, v:Tensor=None, mask:Tensor=None, mem:Tensor=None):\n",
    "        #Notations from the paper:\n",
    "        #x: input, r: vector of relative distance between two elements\n",
    "        #u,v: learnable parameters of the model common between layers\n",
    "        #mask: to avoid cheating, mem: previous hidden states\n",
    "                \n",
    "        #x: [bs, sl, d_model]\n",
    "        #u/v: [n_heads, 1, d_head]\n",
    "        #r: [sl, d_model]\n",
    "        #mem: 1st:[0]; 2nd:[bs, sl, d_model]; nth: sl*i-1 up to mem_len\n",
    "        bs,x_len,seq_len = x.size(0),x.size(1),r.size(0)\n",
    "        context = x if mem is None else torch.cat([mem, x], dim=1)\n",
    "        wq,wk,wv = torch.chunk(self.attention(context), 3, dim=-1)\n",
    "        wq = wq[:,-x_len:]\n",
    "        wq,wk,wv = map(lambda x:x.view(bs, x.size(1), self.n_heads, self.d_head), (wq,wk,wv))\n",
    "        # [bs, sl, n_heads, d_head]\n",
    "        wq,wk,wv = wq.permute(0, 2, 1, 3),wk.permute(0, 2, 3, 1),wv.permute(0, 2, 1, 3)   #wk: transposed(-2,-1)\n",
    "        wkr = self.r_attn(r)\n",
    "        wkr = wkr.view(seq_len, self.n_heads, self.d_head)\n",
    "        wkr = wkr.permute(1,2,0)  #transposed ala wk w/out bs\n",
    "        #### compute attention score (AC is (a) + (c) and BD is (b) + (d) in the paper)\n",
    "        AC = torch.matmul(wq+u,wk)\n",
    "        BD = _line_shift(torch.matmul(wq+v, wkr))\n",
    "        attn_score = (AC + BD).div_(self.d_head ** 0.5)  #scale\n",
    "        if mask is not None:\n",
    "            attn_score = attn_score.float().masked_fill(mask, -float('inf')).type_as(attn_score)\n",
    "        attn_prob = self.drop_att(F.softmax(attn_score, dim=-1))\n",
    "        attn_vec = torch.matmul(attn_prob, wv)\n",
    "        return attn_vec.permute(0, 2, 1, 3).contiguous().view(bs, x_len, -1)\n",
    "    \n",
    "def _left_shift(x:Tensor):\n",
    "    \"Shift the line i of `x` by p-i elements to the left.\"\n",
    "    bs,nh,n,p = x.size()\n",
    "    x_pad = torch.cat([x.new_zeros(bs,nh,n,1), x], dim=3)\n",
    "    x_shift = x_pad.view(bs,nh,p + 1,n)[:,:,1:].view_as(x)\n",
    "    return x_shift"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DecoderLayerXL(Module):\n",
    "    def __init__(self, n_heads:int, d_model:int, drops=0.1):\n",
    "        self.mhra = MultiHeadRelativeAttention(n_heads, d_model, drops=drops)\n",
    "        self.ff   = FeedForward(d_model, drops=drops)\n",
    "\n",
    "    def forward(self, x:Tensor, mask:Tensor=None, **kwargs):\n",
    "        return self.ff(self.mhra(x, mask=mask, **kwargs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class TransformerXL(Module):\n",
    "    \"TransformerXL model: https://arxiv.org/abs/1901.02860.\"\n",
    "    def __init__(self, vocab_sz:int, d_model:int, n_layers:int, n_heads:int, drops:float=0.1, mem_len:int=150):\n",
    "        d_head = d_model//n_heads\n",
    "        self.emb = nn.Embedding(vocab_sz, d_model)\n",
    "        self.drop_emb = nn.Dropout(drops)\n",
    "        self.pos_enc = PositionalEncoding(d_model)\n",
    "        self.u = nn.Parameter(torch.Tensor(n_heads, 1, d_head))\n",
    "        self.v = nn.Parameter(torch.Tensor(n_heads, 1, d_head))\n",
    "        self.layers = nn.ModuleList([DecoderLayerXL(n_heads, d_model, drops) for k in range(n_layers)])\n",
    "        self.init = False\n",
    "\n",
    "    def forward(self, x):\n",
    "        if not self.init:\n",
    "            self.reset()\n",
    "            self.init = True\n",
    "        bs,x_len = x.size()\n",
    "        inp = self.drop_emb(self.emb(x)) #.mul_(self.d_model ** 0.5)\n",
    "        m_len = self.hidden[0].size(1) if hasattr(self, 'hidden') and len(self.hidden[0].size()) > 1 else 0\n",
    "        seq_len = m_len + x_len\n",
    "        mask = torch.triu(x.new_ones(x_len, seq_len), diagonal=1+m_len).bool()[None,None]\n",
    "        pos = torch.arange(seq_len-1, -1, -1, device=inp.device, dtype=inp.dtype)\n",
    "        pos_enc = self.pos_enc(pos)\n",
    "        hids = []\n",
    "        hids.append(inp)\n",
    "        for i, layer in enumerate(self.layers):\n",
    "            mem = self.hidden[i]\n",
    "            inp = layer(inp, r=pos_enc, u=self.u, v=self.v, mask=mask, mem=mem)\n",
    "            hids.append(inp)\n",
    "        core_out = inp[:,-x_len:]\n",
    "        self._update_mems(hids)\n",
    "        return self.hidden,[core_out]\n",
    "    \n",
    "    def _update_mems(self, hids):\n",
    "        if not getattr(self, 'hidden', False): return None\n",
    "        assert len(hids) == len(self.hidden), 'len(hids) != len(self.hidden)'\n",
    "        with torch.no_grad():\n",
    "            for i in range(len(hids)):\n",
    "                cat = torch.cat([self.hidden[i], hids[i]], dim=1)\n",
    "                self.hidden[i] = cat[:,-self.mem_len:].detach()\n",
    "\n",
    "    def reset(self):\n",
    "        \"Reset the internal memory.\"\n",
    "        self.hidden = [next(self.parameters()).data.new(0) for i in range(self.n_layers+1)]\n",
    "\n",
    "    def select_hidden(self, idxs):\n",
    "        # used in beam search only...\n",
    "        self.hidden = [h[idxs] for h in self.hidden]"
   ]
  },
  {
   "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, generator):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.generator = generator\n",
    "        \n",
    "    def forward(self, src, tgt, tgt_mask=None):\n",
    "        return self.decode(self.encode(src), tgt)\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(tgt, src)\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)\n",
    "        x = 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": [
    "class Decoder(nn.Module):\n",
    "    \"Generic N layer decoder with masking.\"\n",
    "    def __init__(self, lm, layer, N):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.lm_layer = lm\n",
    "        self.layers = clones(layer, N)\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "        \n",
    "    def forward(self, x, src):\n",
    "        x = self.lm_layer(x)\n",
    "        \n",
    "        for layer in self.layers:\n",
    "            x = layer(x, src)\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, src_attn, feed_forward, dropout):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.size = size\n",
    "        self.src_attn = src_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 2)  # wraps layer in residual,dropout,norm\n",
    " \n",
    "    def forward(self, x, src):\n",
    "        x = self.sublayer[0](x, lambda x: self.src_attn(x, src, src))\n",
    "        return self.sublayer[1](x, self.feed_forward)"
   ]
  },
  {
   "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, attn_type='multi', attn_heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    lm = TransformerXL(vocab, d_model, n_layers=10, n_heads=8, drops=0.1)\n",
    "    attn = MultiHeadedAttention(d_model, attn_heads) \n",
    "    ff   = PositionwiseFeedForward(d_model, drops)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(lm, DecoderLayer(d_model, c(attn), c(ff), drops), 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):\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, seq_len=700):\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(seq_len)):\n",
    "                    dec_outs = self.transformer.decode(feats, 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",
    "                \n",
    "        #training\n",
    "        else:\n",
    "            feats = self.img_enc(src)\n",
    "            dec_outs = self.transformer(feats, tgt)\n",
    "            out = self.transformer.generate(dec_outs)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.2, attn_type='multi', attn_heads=8):\n",
    "    img_encoder = ResnetBase(em_sz, d_model)\n",
    "    transformer = make_full_model(len(itos), d_model, N=N, drops=drops, attn_type=attn_type, attn_heads=attn_heads)\n",
    "    net = Img2Seq(img_encoder, transformer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=0.1),\n",
    "                    metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# w/ Transformer LM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class DecoderLayerWithLM(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 = clones(feed_forward, 3)\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 5)  # wraps layer in residual,dropout,norm\n",
    " \n",
    "    def forward(self, x, src, tgt_mask=None):\n",
    "        # lm\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))  # shared btw lm and decoder\n",
    "        lm = self.sublayer[1](x, self.feed_forward[0])\n",
    "        \n",
    "        # decoder layer\n",
    "        dec = self.sublayer[2](x, lambda x: self.src_attn(x, src, src))\n",
    "        dec = self.sublayer[3](dec, self.feed_forward[1])\n",
    "        return self.sublayer[4](dec+lm, self.feed_forward[2])"
   ]
  },
  {
   "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, attn_type='multi', attn_heads=8, weight_tying=False):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, attn_heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \n",
    "    ff = PositionwiseFeedForward(d_model, drops)\n",
    "\n",
    "    model = EncoderDecoder(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), drops), N),\n",
    "        Decoder(DecoderLayerWithLM(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",
    "    if weight_tying:\n",
    "        model.generator.weight = model.tgt_embed[0].lut.weight\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# w/ AWDLSTM LM integrated"
   ]
  },
  {
   "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, pos_enc):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.pos_enc = pos_enc\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.pos_enc(tgt), src, tgt_mask)"
   ]
  },
  {
   "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)\n",
    "        x = x.flatten(2,3).permute(0,2,1)\n",
    "        x = self.linear(x) * 8   #(cube root of d_model?)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def lm_mod_model(vocab, d_model, N=4, drops=0.2, attn_type='multi', attn_heads=8):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, attn_heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \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",
    "        PositionalEncoding(d_model, drops, 2000),\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",
    "                \n",
    "    def forward(self, tgts):\n",
    "        res = self.lm(tgts, from_embeddings=True)\n",
    "        pdb.set_trace()\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Mixer(nn.Module):\n",
    "    def __init__(self, d_model, vocab, drops=0.2):\n",
    "        super(Mixer, self).__init__()\n",
    "        self.mixer = nn.Sequential(PositionwiseFeedForward(d_model, drops), LayerNorm(d_model))\n",
    "        self.generator = nn.Linear(d_model, vocab)\n",
    "        \n",
    "    def forward(self, dec_outs, lm_outs):\n",
    "        mix = self.mixer((lm_outs+outs)/2)\n",
    "        return self.generator(mix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class Embedding(nn.Module):\n",
    "    def __init__(self, d_model, vocab, drops, pad_tok=0):\n",
    "        super(Embedding, self).__init__()\n",
    "        self.emb = nn.Embedding(vocab, d_model, padding_idx=pad_tok)\n",
    "        self.emb_drop = EmbeddingDropout(self.emb, drops)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.emb_drop(x)  #self.lut(x) * math.sqrt(self.d_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, embedding, lm, mixer):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.embedding = embedding\n",
    "        self.transformer = transformer\n",
    "        self.lm = lm\n",
    "        self.mixer = mixer\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None, seq_len=700):\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(seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "                    tgt_emb = self.embedding(tgt)\n",
    "                    dec_outs = self.transformer.decode(feats, tgt_emb, mask)\n",
    "                    lm_outs = self.lm(tgt_emb)\n",
    "                    prob = self.mixer(dec_outs[:,-1], lm_outs[:,-1])\n",
    "                    #outs = self.lm(dec_outs[:,-1], self.transformer.embed, 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",
    "            tgt_emb = self.embedding(tgt)\n",
    "            dec_outs = self.transformer(feats, tgt_emb, tgt_mask)    # ([bs, sl, d_model])\n",
    "            lm_outs = self.lm(tgt_emb)\n",
    "            out = self.mixer(dec_outs, lm_outs)\n",
    "            #out = self.transformer.generate(outs)            # ([bs, sl, vocab])\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.2, attn_type='multi', attn_heads=8):\n",
    "    vocab = len(data.vocab.itos)\n",
    "    img_encoder = ResnetBase(em_sz, d_model)\n",
    "    transformer = lm_mod_model(vocab, d_model, N=N, drops=drops, attn_type=attn_type, attn_heads=attn_heads)\n",
    "    embedding = Embedding(d_model, vocab, drops, pad_tok=0)\n",
    "    lm = LM(vocab, d_model, 1400, 3)\n",
    "    mixer = Mixer(d_model, vocab, drops)\n",
    "    net = Img2Seq(img_encoder, transformer, embedding, lm, mixer)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=0.1),\n",
    "                    metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=6, drops=0.1, attn_type='multi')\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Add LM to model state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd = torch.load(PATH/'models/combo_512_9.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# update existing model according to LM modifications\n",
    "sd['model']['embedding.emb.weight'] = sd['model']['transformer.tgt_embed.0.lut.weight']\n",
    "sd['model']['embedding.emb_drop.emb.weight'] = sd['model']['transformer.tgt_embed.0.lut.weight']\n",
    "sd['model'][\"transformer.pos_enc.pe\"] = sd['model']['transformer.tgt_embed.1.pe']\n",
    "sd['model']['mixer.generator.weight'] = sd['model']['transformer.generator.weight']\n",
    "sd['model']['mixer.generator.bias'] = sd['model']['transformer.generator.bias']\n",
    "\n",
    "del sd['model']['transformer.tgt_embed.1.pe']\n",
    "del sd['model']['transformer.tgt_embed.0.lut.weight']\n",
    "del sd['model']['transformer.generator.weight']\n",
    "del sd['model']['transformer.generator.bias']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lm_sd = torch.load('data/wikitext/models/wiki103_lm_enc.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "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": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd['model'].update(new_lm_sd)\n",
    "\n",
    "learn.model.load_state_dict(sd['model'], strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# tie weights of transformer embedding and generator w/ lm encodings\n",
    "learn.model.embedding.emb.weight = learn.model.lm.lm.encoder.weight\n",
    "learn.model.embedding.emb_drop.emb.weight = learn.model.lm.lm.encoder_dp.emb.weight\n",
    "learn.model.mixer.generator.weight = learn.model.lm.lm.encoder.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('combo_512_9_wiki103_base_lm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "## load and split learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('combo_512_9_wiki103_base_lm')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.split([learn.model.img_enc, learn.model.embedding, learn.model.transformer, learn.model.lm, learn.model.mixer])\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.layer_groups[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.freeze_to(-3)\n",
    "learn.model.img_enc.linear.weight.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lrs = slice(2e-5, 2e-4, 2e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# w/ AWDLSTM added"
   ]
  },
  {
   "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)\n",
    "        x = 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, attn_type='multi', attn_heads=8, weight_tying=False):\n",
    "    c = deepcopy\n",
    "    \n",
    "    if attn_type=='multi':\n",
    "        attn = MultiHeadedAttention(d_model, attn_heads)\n",
    "    else:\n",
    "        attn = SingleHeadedAttention(d_model)\n",
    "        \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",
    "    if weight_tying:\n",
    "        model.generator.weight = model.tgt_embed[0].lut.weight\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, lm):\n",
    "        super(Img2Seq, self).__init__()\n",
    "        self.img_enc = img_encoder\n",
    "        self.transformer = transformer\n",
    "        self.lm = lm\n",
    "        \n",
    "    def forward(self, src, tgt=None, tgt_mask=None, seq_len=700):\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(seq_len)):\n",
    "                    mask = subsequent_mask(tgt.size(-1))\n",
    "                    dec_outs = self.transformer.decode(feats, tgt, mask)\n",
    "                    dec_prob = self.transformer.generate(dec_outs[:,-1])\n",
    "                    lm_prob,_,_ = self.lm(tgt)\n",
    "                    prob = (dec_prob + lm_prob[:,-1])/2\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": [
    "lm_config = dict(emb_sz=512, n_hid=1400, n_layers=3, pad_token=0, qrnn=False, bidir=False, output_p=0.2,\n",
    "              hidden_p=0.2, input_p=0.5, embed_p=0.1, weight_p=0.4, tie_weights=True, out_bias=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def make_learner(data, d_model, em_sz, N=4, drops=0.2, attn_type='multi', attn_heads=8):\n",
    "    vocab = len(data.vocab.itos)\n",
    "    img_encoder = ResnetBase(em_sz, d_model)\n",
    "    transformer = make_full_model(vocab, d_model, N=N, drops=drops, attn_type=attn_type, attn_heads=attn_heads)\n",
    "    lm = get_language_model(AWD_LSTM, vocab, lm_config, drop_mult=0.5)\n",
    "    net = Img2Seq(img_encoder, transformer, lm)\n",
    "    return Learner(data, net, loss_func=LabelSmoothing(smoothing=0.1),\n",
    "                    metrics=[CER()], callback_fns=[TeacherForce])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=6, drops=0.1, attn_type='multi')\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Add LM to model state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.model.lm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd = torch.load(PATH/'models/combo_512_9.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lm_sd = torch.load(PATH/'models/wiki2_lm.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "new_lm_sd = OrderedDict()\n",
    "for k, v in lm_sd['model'].items():\n",
    "    name = 'lm.'+k\n",
    "    new_lm_sd[name] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "lm_sd['model'].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd['model'].update(new_lm_sd)\n",
    "\n",
    "learn.model.load_state_dict(sd['model'], strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('combo_512_9_wiki2_lm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Num Lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "                .split_by_rand_pct(valid_pct=0.15, seed=42)\n",
    "                .label_from_df(cols=2)\n",
    "                .transform(tfms, size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "                .databunch(bs=bs, val_bs=bs*2, device=device)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2 = cnn_learner(data, models.resnet34, metrics=accuracy)\n",
    "# Note: already split and frozen:  m[0][6], m[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# load saved sd into l2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "new_sd = OrderedDict()\n",
    "\n",
    "sd = torch.load(PATH/'models/combo_512_mix_res9.pth', map_location=device)\n",
    "for k, v in sd['model'].items():\n",
    "    if k.startswith('img_enc.base'):\n",
    "        name = k.replace('img_enc.base', '0')\n",
    "        new_sd[name] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2_sd = torch.load(PATH/'models/line_97acc.pth', map_location=device)\n",
    "for k, v in l2_sd['model'].items():\n",
    "    if k.startswith('1.'):\n",
    "        new_sd[k] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2.model.load_state_dict(new_sd, strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2.layer_groups[1][47].weight.requires_grad, l2.layer_groups[2][9].weight.requires_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "l2.lr_find()\n",
    "l2.recorder.plot(suggestion=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2.fit_one_cycle(1,1e-3)\n",
    "# 0.171519\t0.069522\t0.976006\t12:57     'line_97acc'\n",
    "# 0.094188\t0.050964\t0.988564\t12:01  preload img_enc: combo_512_mix_res9 and head: line_97acc   'line_98acc'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "l2.save('line_98acc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# learn = cnn_learner(data, models.resnet34, custom_head=create_head(1024, 1), metrics=accuracy, loss_func=MSELossFlat())\n",
    "# learn.fit_one_cycle(1,1e-3)\n",
    "# # 1.229723\t0.113453\t0.9057\t12:43\n",
    "\n",
    "# preds,y,losses = learn.get_preds(with_loss=True)\n",
    "\n",
    "# # actual accuracy\n",
    "# preds.apply_(round)\n",
    "# (preds.long().squeeze(-1) == y).float().mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=4, drops=0, attn_type='multi', heads=8)\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('combo_512_mix_res8')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.lr_find()\n",
    "learn.recorder.plot(suggestion=True)"
   ]
  },
  {
   "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
   },
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(5, max_lr=1e-5, callbacks=[SaveModelCallback(learn, name='sm_256_lines2')])\n",
    "#sm, 5cycle, 1e-3\n",
    "\n",
    "# Image Preprocessing by num_lines\n",
    "# 11.384931\t10.208405\t0.122443   'sm_256_lines'\n",
    "# 10.548290\t9.845284\t0.116961   2nd run; 5cycle(1e-5); 'sm_256_lines2'\n",
    "\n",
    "# \n",
    "\n",
    "# 2.474220\t4.081600\t0.036482  N:4,F.gelu,sz:256,em_sz:512,single,drop:0  'sm_256_1'\n",
    "# greedy:    1.36697   .066\n",
    "# 4.430638\t4.402001\t0.034730  2nd run, lr:1.5e-5, add tfms, drop:0.2   'sm_256_2'\n",
    "# greedy:    0.37206   .04145\n",
    "\n",
    "# 5.158808\t4.964441\t0.042544  \"\", w/ tfms, drop:0.2   'sm_256_3'\n",
    "# greedy:    0.50646   .04103\n",
    "# 4.266788\t4.694098\t0.039822  2nd run, lr:1.5e-5   'sm_256_4'\n",
    "# greedy:    0.38972   .03879\n",
    "\n",
    "# 3.678168\t3.962710\t0.035466  N:8,tfms   'sm_256_5'\n",
    "# greedy:    0.38097   .04405\n",
    "\n",
    "# 2.840161\t3.429177\t0.030243  N:4,tfms,multi(8)  'sm_256_6'\n",
    "# greedy:    0.32775   .04201\n",
    "# greedy:    1.38865   .06943\n",
    "\n",
    "# 3.377438\t3.786088\t0.033931   N:6,tfms,drop:0.1,multi(8)   'sm_256_7'\n",
    "# greedy:    0.36201   .04266  \n",
    "# greedy:    1.40557   .06902\n",
    "\n",
    "# 3.989137\t4.353169\t0.037717   N:4,weight-tying,multi(16)   'sm_256_8'\n",
    "# greedy:    1.44291   .07299\n",
    "\n",
    "# 4.806501\t5.007887\t0.042339   N:4,weight-tying,multi(8),drops:0.2   'sm_256_9'\n",
    "# greedy:    1.46753   .07424\n",
    "\n",
    "# 12.817862\t11.842313\t0.105486   N:4,no init/weight tying,multi(8),drops:0.1   'sm_256_10'\n",
    "\n",
    "# 3.515456\t3.880955\t0.033763   N:4,multi(8),drops:0.1   'sm_256_11'\n",
    "# greedy:    1.40945   .06964\n",
    "\n",
    "# combo_cat6 - preload 'sm_256_7'\n",
    "# 7.570516\t6.670641\t0.015070    'combo_512_7'\n",
    "# greedy:    6.29406   .02186\n",
    "\n",
    "\n",
    "# combo 6lines and under, preload 'sm_256_6', 10cycle(1e-3)\n",
    "# 10.784849\t9.144131\t0.024880     'combo_512_13'\n",
    "# greedy:    9.21715   .02788\n",
    "# upload:    103.456   .44797\n",
    "#     pg:    148.473   .69869\n",
    "\n",
    "# combo 6lines and over, preload 'combo_512_13', 5cycle(2e-5)\n",
    "# 56.732105\t38.604137\t0.026153    'combo_512_14'\n",
    "# greedy:    141.022   .02933\n",
    "# upload:    179.606   .81724\n",
    "\n",
    "# combo_cat_pg - preload 'combo_512_7'\n",
    "# 25.674347\t21.161726\t0.015869    'combo_512_8'\n",
    "# greedy:    114.878   .03654\n",
    "#   test:    115.247   .05014\n",
    "\n",
    "# combo_cat_pg_dl_sorted - preload 'combo_512_8', 5cycle, 1.5e-4\n",
    "# 9.800321\t5.676855\t0.006939    'combo_512_9'\n",
    "# greedy:    34.8757   .01523\n",
    "\n",
    "#   test:    91.3226   .05102\n",
    "#   test:    104.685   .05483\n",
    "#   test:    116.467   .05315\n",
    "\n",
    "# upload:    115.149   .66801\n",
    "\n",
    "# imdb_wiki_combo - preload 'combo_512_9'; split: img_enc, transformer; 3cycle, slice(5e-5, 1e-3)\n",
    "# 20.390352\t16.698141\t0.018726    'combo_512_12'\n",
    "# greedy:    60.8384   .02474\n",
    "#     pg:    112.755   .05204\n",
    "\n",
    "# combo_145k - fit:1cycle(1e-3); preload combo_512_9; drops:0.2\n",
    "# 30.075714\t31.024599\t0.031703\t1:50:33\n",
    "\n",
    "# combo_145k <= 6 - 5cycle(1e-4); preload sm_256_mix_res\n",
    "# 19.502026\t17.365313\t0.044436    'combo_512_mix_res'\n",
    "# greedy:    14.6870   .04744\n",
    "# upload:    75.1357   .47955\n",
    "#     pg:    142.710   .66233\n",
    "\n",
    "# test_combo 65k (sm,cat,pg,dl)\n",
    "# 18.262995\t15.091379\t0.039351   5cycle(1e-5)  'combo_512_mix_res5'\n",
    "# greedy:    129.968   .03932\n",
    "# upload:    81.6835   .21794\n",
    "#     pg:    127.317   .05613\n",
    "\n",
    "# 16.678631\t12.867742\t0.035985   1cycle(1e-5), drops:0.1   'combo_512_mix_res6'\n",
    "# 14.878900\t12.797726\t0.035834   1cycle(1e-6), drops:0.1   'combo_512_mix_res7'\n",
    "# greedy:    34.3988   .05552\n",
    "# upload:    75.5511   .21010\n",
    "#     pg:    118.963   .05287\n",
    "\n",
    "# 16.086544\t12.320695\t0.034263   1cycle(1e-5), drops:0.1   'combo_512_mix_res8'\n",
    "# 11.640953\t11.339241\t0.032201   1cycle(1e-5), drops:0     'combo_512_mix_res9'\n",
    "# greedy:    31.5506   .05410\n",
    "#   test:    74.4820   .22067\n",
    "#     pg:    105.182   .05076\n",
    "\n",
    "\n",
    "# 79.324310\t69.801376\t0.325578\t23:04   preload lines/combo_512_mix_res9;  1cycle(1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Experiment - Words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=6, drops=0.1, attn_type='multi')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('combo_512_9')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.lr_find()\n",
    "learn.recorder.plot(suggestion=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.fit_one_cycle(15, max_lr=1e-3, callbacks=[SaveModelCallback(learn, name='word_sm_2')])\n",
    "\n",
    "# word 60k\n",
    "# sm, 5cycle, 1e-3\n",
    "\n",
    "# 22.903589\t22.401066\t0.520125   N:6,em_sz:512,tfms,drop:0.1,multi(8)  'word_sm_1'\n",
    "# 32.389709\t31.828810\t0.746749   N:4,em_sz:512,tfms,drop:0.1,single\n",
    "# 34.492157\t32.719650\t0.776098   \"\",em_sz:256\n",
    "\n",
    "# word 7k (auto-tokenize from data)\n",
    "# sm, 5cycle, 1e-3\n",
    "# 16.987526\t16.639275\t0.421535   N:4,em_sz:512,tfms,drop:0.1,single\n",
    "# 10.224819\t10.781199\t0.297889   2nd run\n",
    "# 6.486265\t7.934221\t0.225055   3rd run    'word_sm_4'\n",
    "\n",
    "# word 10k (itos from larger txt files)\n",
    "# sm, 15cycle, 1e-3\n",
    "# 3.893518\t6.037587\t0.209027   N:4,em_sz:512,tfms,drop:0.1,single   'word_sm_2'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Greedy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=6, drops=0.1, attn_type='multi')\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd = torch.load(PATH/'models/combo_512_9.pth', map_location=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sd['model'][\"transformer.src_adapt.0.weight\"] = sd['model']['img_enc.linear.weight']\n",
    "sd['model'][\"transformer.src_adapt.0.bias\"] = sd['model']['img_enc.linear.bias']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.model.load_state_dict(sd['model'], strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# learn.load('combo_512_9_wiki2_lm')\n",
    "learn.load('combo_512_mix_res')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def full_test(learn, sl, dl=data.valid_dl, batches=10):\n",
    "    learn.model.eval()\n",
    "    iterable = iter(dl)\n",
    "    g_loss,g_cer=0,0\n",
    "    if batches is None:\n",
    "        batches = len(dl.dl.dataset)//bs\n",
    "    for i in progress_bar(range(batches)):\n",
    "        x,y = next(iterable)\n",
    "        g_preds = learn.model(x, seq_len=sl)\n",
    "        g_res = torch.argmax(g_preds, dim=-1)\n",
    "        g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y)[0]/bs]\n",
    "        g_loss+=g[0]\n",
    "        g_cer+=g[1]\n",
    "    return [g_loss/batches, g_cer/batches]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# set_seed()\n",
    "g = full_test(learn, seq_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "print(f'greedy:    {str(g[0])[:7]}   {str(g[1])[1:7]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# losses = np.array([learn.loss_func(g_preds[i:i+1],y[i:i+1]).item() for i in range(bs)])\n",
    "# cers = np.array([cer(g_preds[i:i+1],y[i:i+1])[0] for i in range(bs)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# set_seed()\n",
    "x,y = next(iter(learn.data.valid_dl))\n",
    "\n",
    "g_preds = learn.model(x, seq_len=seq_len)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y)[0]/bs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#greedy\n",
    "fig, axes = plt.subplots(2,3, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    p = char_label_text(g_res[i], sep='')\n",
    "    ax=show_img(x[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "FOLDER = 'uploads'\n",
    "df = pd.read_csv(PATH/'uploads.csv')\n",
    "len(df)\n",
    "\n",
    "sz,bs = 512,14\n",
    "seq_len = 700"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "FOLDER = 'paragraphs'\n",
    "df = pd.read_csv(PATH/'test_pg.csv')\n",
    "len(df)\n",
    "\n",
    "sz,bs = 512,15\n",
    "seq_len = 700"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "## Test dataset only!!!\n",
    "# set_seed()  # reproducibility\n",
    "data = (ImageList.from_df(df, path=PATH, folder=FOLDER, after_open=force_gray)\n",
    "        .split_none()\n",
    "        .label_from_df(label_cls=SequenceList, vocab=CharVocab(itos))\n",
    "        .transform([], size=sz, resize_method=ResizeMethod.SQUISH)\n",
    "        .databunch(bs=bs, device=device, collate_fn=label_collater)\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = next(iter(data.train_dl))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### learner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, 512, N=4, drops=0.1, attn_type='multi')\n",
    "# Total # of trainable params:\n",
    "# N=6: 65,786,272\n",
    "# N=4: 51,073,440"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# learn.load('combo_512_9_wiki2_lm')\n",
    "learn.load('combo_512_mix_res')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# learn.data = data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GPU batch testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = learn.model(x, seq_len=seq_len)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "g = [learn.loss_func(g_preds, y).item()/bs, cer(g_preds, y)[0]/bs]\n",
    "\n",
    "print(f'  test:    {str(g[0])[:7]}   {str(g[1])[1:7]}')\n",
    "\n",
    "# test:    397.121   .29261  ~3:55   combo_512_9_wiki2_lm  (LM isn't aware of '\\n'!!)\n",
    "# test:    101.374   .05139  ~27s    combo_512_9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "#test\n",
    "fig, axes = plt.subplots(2,2, gridspec_kw={'hspace': 0.4}, figsize=(18, 20))\n",
    "for i,ax in enumerate(axes.flat):\n",
    "    i +=4\n",
    "    p = char_label_text(g_res[i])\n",
    "    ax=show_img(x[i], ax=ax, title=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing Lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# test data\n",
    "x,y = next(iter(data.train_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def lines_into_paragraph(res):\n",
    "    out = []\n",
    "    for r in res:\n",
    "        s = np.split(r, np.where(r == 2)[0])[0].numpy()\n",
    "        out.append(np.append(s,4))\n",
    "    out = np.concatenate(out)\n",
    "    out[-1] = 2  # replace final '\\n' with 'eos'\n",
    "    return torch.from_numpy(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def target_into_lines(targ, split_idx=4):\n",
    "    nonzero = targ[targ.nonzero()].flatten()\n",
    "    lines = np.split(nonzero, np.where(nonzero == split_idx)[0])\n",
    "    maxlen = len(max(lines,key=len))\n",
    "    res = torch.zeros((len(lines),maxlen))\n",
    "    for i,arr in enumerate(lines):\n",
    "        res[i,:len(arr)] = arr\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from scipy import signal,ndimage\n",
    "import statistics\n",
    "\n",
    "def img_into_line_tensor(img):\n",
    "    arr = img[0].numpy()\n",
    "    w = arr.shape[-1]\n",
    "    heights = []\n",
    "    lengths = []\n",
    "    stds = arr.std(axis=1)\n",
    "    g_stds = scipy.ndimage.gaussian_filter1d(stds, 5)\n",
    "    peaks,_ = scipy.signal.find_peaks(g_stds, prominence=stds.std()//3, distance=20)\n",
    "    mins = scipy.signal.argrelextrema(g_stds, np.less_equal)[0]  # np.less_equal critical for flat minima, edges\n",
    "    for p in peaks:\n",
    "        rows = range(mins[mins < p][-1], mins[mins > p][0])\n",
    "        lengths.append(len(rows))\n",
    "        heights.append(arr[rows])\n",
    "    max_len = max(lengths)\n",
    "    outs = torch.ones((len(heights),3 , max_len, w))\n",
    "    for i,h in enumerate(heights):\n",
    "        outs[i,:,:len(h)] = torch.from_numpy(h)\n",
    "    return outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idx = 1\n",
    "inps = img_into_line_tensor(x[idx])\n",
    "targs = target_into_lines(y[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = learn.model(inps, seq_len=seq_len)\n",
    "g_res = torch.argmax(g_preds, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "print(f'test cer:   {str(cer(g_preds,targs)[0]/g_preds.size(0))[1:7]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {}
   },
   "source": [
    "## Single Image (cpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = data.train_ds[5]\n",
    "pred = learn.predict(x)\n",
    "show_img(x, title=str(pred[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "pred[2].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "probs = F.softmax(pred[2], dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "scores, idxs = torch.topk(probs, 3, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "scores[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "for s,i in zip(scores,idxs):\n",
    "    new = {}\n",
    "    for score, idx in zip(s,i):\n",
    "        new[itos[idx]] = score.data\n",
    "    print(new)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "chars = idxs.numpy()\n",
    "chars.appl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "idxs.apply_(lambda x: itos[x])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "for i in range(30):\n",
    "    print(char_label_text(idxs[i], sep=' '))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# learn.show_results(ds_type=data.train_ds, rows=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "i = 0\n",
    "x = xs[i][None]\n",
    "y = ys[i][None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "batch = data.one_item(xs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "xs,ys = data.one_batch(detach=False, denorm=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "g_preds = learn.model(x, seq_len=seq_len)\n",
    "g_res = torch.argmax(g_preds, dim=-1)\n",
    "g = [learn.loss_func(g_preds, y).item(), cer(g_preds, y)[0]]\n",
    "\n",
    "print(f'  test:    {str(g[0])[:7]}   {str(g[1])[1:7]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "p = char_label_text(g_res[0])\n",
    "show_img(x[0], figsize=(18,10), title=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "im = PATH/'uploads'/'test1.png'\n",
    "img = open_image(im)\n",
    "prediction = learn.predict(img)[0]\n",
    "show_img(img, title=str(prediction))"
   ]
  }
 ],
 "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
}