1-ArchLM.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": [
    "## 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": [
    "# sz,bs = 128,100\n",
    "sz,bs = 256,50\n",
    "\n",
    "num_lines,seq_len = 4,100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "df = df[:20000]\n",
    "len(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": [
    "# set_seed()  # reproducibility\n",
    "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(tfms, 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": [
    "# 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, 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), self.embed(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(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, 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())[:-2]    \n",
    "        self.base = nn.Sequential(*modules)                  #bs,512,16,16\n",
    "        \n",
    "        self.c1 = conv_layer(d_model, d_model, stride=(2,1))     #8\n",
    "        self.c2 = conv_layer(d_model, d_model, stride=(2,1))     #4\n",
    "        self.c3 = conv_layer(d_model, d_model, stride=(2,1))     #2\n",
    "#         self.c4 = conv_layer(d_model, d_model, stride=(2,1))     #4\n",
    "#         self.c5 = conv_layer(d_model, d_model, stride=(2,1))     #2\n",
    "                \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x1 = self.c1(x)\n",
    "        x2 = self.c2(x1)\n",
    "        x3 = self.c3(x2)\n",
    "                \n",
    "#         cat = torch.cat([x,x1,x2,x3], dim=2)\n",
    "        \n",
    "        cat = torch.cat([x.flatten(2,3),x1.flatten(2,3),x2.flatten(2,3),x3.flatten(2,3)], dim=-1)\n",
    "        return cat.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 = ResnetBase(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": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = make_learner(data, 512, N=4, drops=0, attn_type='multi', heads=8)"
   ]
  },
  {
   "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",
    "\n",
    "# 5cycle,1e-3\n",
    "# 3.510155\t3.892951\t0.033866  baseline     'sm_256_base'\n",
    "# greedy:    0.43533   .04138\n",
    "# 3.156924\t3.474473\t0.030223  p_mask(2)  'sm_256_pmask2'\n",
    "# greedy:    0.45049   .04231\n",
    "# greedy:    1.09986   .16270   w/ truncation -> much faster\n",
    "# 3.666398\t3.924176\t0.032230  p_mask(2),drops:0.2\n",
    "# greedy:    0.46281   .04217\n",
    "\n",
    "# 3.974781\t4.305206\t0.036939   mix_res,p_mask(10)   'sm_256_pmask10'\n",
    "# greedy:    0.47993   .04458\n",
    "# greedy:    1.04631   .11497    w/ truncation\n",
    "\n",
    "# 4.040658\t4.392500\t0.037479   mix_res   'sm_256_mix_res'\n",
    "# greedy:    0.50428   .04335\n",
    "\n",
    "# bs:50; multi-resolution; stride(2,1); no src adaptor\n",
    "# 2.408194\t3.065314\t0.028297   'sm_256_multi_res'\n",
    "# greedy:    0.74505   .03624\n",
    "\n",
    "# 2.532134\t3.096449\t0.028337   head:32  'sm_256_multi_res2'\n",
    "# greedy:    0.75551   .03680\n",
    "\n",
    "# combo_145k <= 6; 5cycle(1e-3)\n",
    "# 5.069114\t5.064365\t0.015234\n",
    "# "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('combo_512_multi_res')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.load('sm_256_multi_res'); None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Mix-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]\n",
    "        self.base = nn.Sequential(*modules)                  #16x16\n",
    "        self.c1 = conv_layer(d_model, d_model, stride=2)     #8x8\n",
    "        self.c2 = conv_layer(d_model, d_model, stride=2)     #4x4\n",
    "        self.c3 = conv_layer(d_model, d_model, stride=2)     #2x2\n",
    "                \n",
    "    def forward(self, x):\n",
    "        x = self.base(x)\n",
    "        x1 = self.c1(x)\n",
    "        x2 = self.c2(x1)\n",
    "        x3 = self.c3(x2)\n",
    "        cat = torch.cat([x.flatten(2,3),x1.flatten(2,3),x2.flatten(2,3),x3.flatten(2,3)], dim=-1)\n",
    "        return cat.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.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",
    "            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.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, em_sz, 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,
    "hidden": true
   },
   "source": [
    "## Experimentation"
   ]
  },
  {
   "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', 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": [
    "tot = 0\n",
    "for p in learn.model.parameters():\n",
    "    tot += p.sum()\n",
    "tot.data\n",
    "# -3712.5525\n",
    "# -3709.0364  w/ weight tying\n",
    "# 1315.4849   \"\", N=6\n",
    "# -8694.8057  \"\", N=2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# set_seed()\n",
    "learn.fit(1, lr=1e-3)\n",
    "\n",
    "# sm(20k subset), 1cycle, 1e-3, d_model/em_sz:512, N:4, multi(8), drops:0.1\n",
    "# 66.245880\t61.198948\t0.716475\t06:51  baseline\n",
    "# 66.746071\t61.807365\t0.705620\t06:51  w/ weight tying    ***seems to help CER...\n",
    "# 79.020279\t76.020065\t0.936690\t06:54  split(img_enc, tfmr), lrs(1e-4,1e-2)\n",
    "# 68.644821\t63.186871\t0.724228\t07:17  N=6\n",
    "# 68.506302\t62.820026\t0.762108\t07:13  N=6 (w/out weight tying)\n",
    "# 66.554543\t61.799850\t0.848753\t06:24  N=2\n",
    "# 71.178879\t65.140076\t0.779076\t08:08  d_model: 768\n",
    "# 66.679169\t60.917271\t0.793617\t06:53  attn: single\n",
    "# 66.829041\t61.852840\t0.724487\t06:54  attn: multi(4)\n",
    "# 66.263069\t61.742168\t0.699408\t07:13  attn: multi(16)  *\n",
    "# 76.158379\t70.973106\t0.710026\t06:36  smoothing: 0.0\n",
    "# 60.895245\t54.817890\t0.746580\t06:39  smoothing: 0.2\n",
    "# 69.718803\t64.245934\t0.784929\t06:37  w/out src_adaptor .mul_(8)\n",
    "# 66.709183\t61.714687\t0.707358\t06:38  F.relu\n",
    "\n",
    "# 66.384956\t61.029957\t0.737656\t02:07  baseline\n",
    "# 64.477470\t59.358719\t0.679445\t02:09  parallelogram_mask(10)\n",
    "# 64.270401\t59.660389\t0.679594\t02:06  pos_enc after img_enc\n",
    "# 64.752815\t60.505432\t0.697868\t02:07  parallelogram_mask(20)\n",
    "# 63.160973\t57.148361\t0.674351\t02:08  parallelogram_mask(5)\n",
    "# 59.376633\t53.333813\t0.649422\t02:07  parallelogram_mask(1)\n",
    "# 64.695885\t61.779552\t0.791945\t02:08  parallelogram_mask(0)\n",
    "# 59.957340\t52.987080\t0.611505\t02:10  parallelogram_mask(2)  *\n",
    "# 61.474442\t55.054062\t0.649738\t02:09  parallelogram_mask(3)\n",
    "# 61.727688\t54.237797\t0.633881\t02:09  weight tying\n",
    "# 58.971718\t52.863750\t0.643140\t02:33  N=6, drops:0.2\n",
    "# 58.866985\t52.895126\t0.655549\t02:33  N=6, drops:0.1\n",
    "# 60.226372\t53.237236\t0.609828\t02:12  multi(16), drops:0.2\n",
    "# 58.924366\t52.224678\t0.671493\t02:14  multi(16), drops:0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# set_seed()\n",
    "learn.fit_one_cycle(5, max_lr=1e-3)\n",
    "\n",
    "# 5cycle,1e-3\n",
    "# 3.510155\t3.892951\t0.033866  baseline     'sm_256_base'\n",
    "# greedy:    0.43533   .04138\n",
    "# 3.156924\t3.474473\t0.030223  p_mask(2)  'sm_256_pmask2'\n",
    "# greedy:    0.45049   .04231\n",
    "# greedy:    1.09986   .16270   w/ truncation -> much faster\n",
    "# 3.666398\t3.924176\t0.032230  p_mask(2),drops:0.2\n",
    "# greedy:    0.46281   .04217\n",
    "\n",
    "# 3.974781\t4.305206\t0.036939   mix_res,p_mask(10)   'sm_256_pmask10'\n",
    "# greedy:    0.47993   .04458\n",
    "# greedy:    1.04631   .11497    w/ truncation\n",
    "\n",
    "# 4.040658\t4.392500\t0.037479   mix_res   'sm_256_mix_res'\n",
    "# greedy:    0.50428   .04335"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn.save('sm_256_mix_res')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn.load('sm_256_mix_res'); None"
   ]
  },
  {
   "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": {
    "heading_collapsed": true,
    "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": {},
   "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=3e-5, callbacks=[SaveModelCallback(learn, name='combo_512_multi_res2')])\n",
    "#sm, 5cycle, 1e-3\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",
    "# upload:    74.4820   .22067\n",
    "#     pg:    105.182   .05076\n",
    "\n",
    "# test_combo 65k - sz:1024\n",
    "# --fail--   preload combo_512_mix_res5 drop:0.1, 3cycle(2e-4)  \n",
    "\n",
    "\n",
    "# combo_145k, 5cycle(3e-5)\n",
    "# 21.098228\t15.413793\t0.014271\t1:32:22  'combo_512_multi_res2'\n",
    "# greedy:    59.5139   .02229\n",
    "# upload:    71.6485   .28271\n",
    "# upload:    62.4138   .26386\n",
    "#     pg:    117.644   .06028"
   ]
  },
  {
   "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,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data.show_batch(rows=4, ds_type=DatasetType.Train, figsize=(18,20))"
   ]
  },
  {
   "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\n",
    "\n",
    "learn = make_learner(data, 512, N=4, drops=0, attn_type='multi', heads=8)"
   ]
  },
  {
   "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_multi_res2')\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": [
    "### test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = next(iter(learn.data.train_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]\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": {
    "ExecuteTime": {},
    "heading_collapsed": true
   },
   "source": [
    "## Single Image (cpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "x,y = data.train_ds[13]\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "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": [
    "outs,targs = pred[2][None],tensor(y.data)[None]\n",
    "g = [learn.loss_func(outs, targs).item(), cer(outs, targs)[0]]\n",
    "print(f'image:    {str(g[0])[:7]}   {str(g[1])[1:7]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {},
    "heading_collapsed": true,
    "hidden": true
   },
   "outputs": [],
   "source": [
    "probs = F.softmax(pred[2], dim=-1)\n",
    "scores, idxs = torch.topk(probs, 3, dim=-1)"
   ]
  },
  {
   "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.7.3"
  },
  "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
}