image2latex.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from collections import Counter\n",
    "from collections import namedtuple\n",
    "from os.path import join\n",
    "from functools import partial\n",
    "\n",
    "import numpy as np\n",
    "import PIL\n",
    "import torch\n",
    "import torchvision\n",
    "\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import init\n",
    "from torch.distributions.uniform import Uniform\n",
    "from torch.nn.utils import clip_grad_norm_\n",
    "from torch.distributions.bernoulli import Bernoulli\n",
    "import torch.optim as optim\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_dir = \"formula_images_processed\"\n",
    "formula_list = \"im2latex_formulas.norm.lst\"\n",
    "train_list = \"im2latex_train_filter.lst\"\n",
    "validate_list = \"im2latex_validate_filter.lst\"\n",
    "test_list = \"im2latex_test_filter.lst\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Image2LatexDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, images_dir, formula_list, train_list):\n",
    "        self.images_dir = images_dir\n",
    "\n",
    "        with open(formula_list, \"r\", encoding=\"utf-8\", errors=\"ignore\", newline=\"\\n\") as f1:\n",
    "            self.formulas = [formula.replace(\"\\n\", \"\").replace(\"\\t\", \" \") for formula in f1.readlines()]\n",
    "        \n",
    "        with open(train_list, \"r\", encoding=\"utf-8\", errors=\"ignore\", newline=\"\\n\") as f2:\n",
    "            self.train_set = [t.replace(\"\\n\", \"\").split() for t in f2.readlines()]\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        item = self.train_set[idx]\n",
    "        filename = item[0]\n",
    "        formula = self.formulas[int(item[1])]\n",
    "#         render_type = item[2]\n",
    "        image = PIL.Image.open(self.images_dir + \"/\" + filename)\n",
    "        return torchvision.transforms.ToTensor()(image), formula\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.train_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set = Image2LatexDataset(images_dir, formula_list, train_list)\n",
    "validate_set = Image2LatexDataset(images_dir, formula_list, validate_list)\n",
    "test_set = Image2LatexDataset(images_dir, formula_list, test_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "START_TOKEN = 0\n",
    "PAD_TOKEN = 1\n",
    "END_TOKEN = 2\n",
    "UNK_TOKEN = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Vocab(object):\n",
    "    def __init__(self):\n",
    "        self.sign2id = {\"<s>\": START_TOKEN, \"</s>\": END_TOKEN,\n",
    "                        \"<pad>\": PAD_TOKEN, \"<unk>\": UNK_TOKEN}\n",
    "        self.id2sign = dict((idx, token)\n",
    "                            for token, idx in self.sign2id.items())\n",
    "        self.length = 4\n",
    "\n",
    "    def add_sign(self, sign):\n",
    "        if sign not in self.sign2id:\n",
    "            self.sign2id[sign] = self.length\n",
    "            self.id2sign[self.length] = sign\n",
    "            self.length += 1\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "\n",
    "\n",
    "def build_vocab(min_count=10):\n",
    "    \"\"\"\n",
    "    traverse training formulas to make vocab\n",
    "    and store the vocab in the file\n",
    "    \"\"\"\n",
    "    vocab = Vocab()\n",
    "    counter = Counter()\n",
    "\n",
    "    formulas_file = formula_list\n",
    "    with open(formula_list, \"r\", encoding=\"utf-8\", errors=\"ignore\", newline=\"\\n\") as f1:\n",
    "        formulas = [formula.replace(\"\\n\", \"\").replace(\"\\t\", \" \") for formula in f1.readlines()]\n",
    "        \n",
    "    with open(train_list, \"r\", encoding=\"utf-8\", errors=\"ignore\", newline=\"\\n\") as f2:\n",
    "        for line in f2:\n",
    "            img_filename, idx = line.strip('\\n').split()\n",
    "            idx = int(idx)\n",
    "            formula = formulas[idx].split()\n",
    "            counter.update(formula)\n",
    "\n",
    "    for word, count in counter.most_common():\n",
    "        if count >= min_count:\n",
    "            vocab.add_sign(word)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = build_vocab()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, out_channels=512, add_pos_feat=True):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.cnn = nn.Sequential(\n",
    "            nn.Conv2d(3, 64, 3, 1, 1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2, 1),\n",
    "\n",
    "            nn.Conv2d(64, 128, 3, 1, 1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2, 1),\n",
    "\n",
    "            nn.Conv2d(128, 256, 3, 1, 1),\n",
    "            nn.ReLU(),\n",
    "            \n",
    "            nn.Conv2d(256, 256, 3, 1, 1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d((2, 1), (2, 1), 0),\n",
    "\n",
    "            nn.Conv2d(256, out_channels, 3, 1, 0),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "\n",
    "        self.add_pos_feat = add_pos_feat\n",
    "    \n",
    "    def forward(self, images):\n",
    "        encoded_images = self.cnn(images)\n",
    "        encoded_images = encoded_images.permute(0, 2, 3, 1)  # [B, H', W', 512]\n",
    "        B, H, W, C = encoded_imgs.shape\n",
    "        encoded_images = encoded_images.contiguous().view(B, H*W, C)\n",
    "        if self.add_pos_feat:\n",
    "            encoded_images = self.add_positional_features(encoded_images)\n",
    "        return encoded_images\n",
    "    \n",
    "    def add_positional_features(self, tensor: torch.Tensor, min_timescale: float = 1.0, max_timescale: float = 1.0e4):\n",
    "        \"\"\"\n",
    "        Implements the frequency-based positional encoding described\n",
    "        in `Attention is all you Need\n",
    "        Parameters\n",
    "        ----------\n",
    "        tensor : ``torch.Tensor``\n",
    "            a Tensor with shape (batch_size, timesteps, hidden_dim).\n",
    "        min_timescale : ``float``, optional (default = 1.0)\n",
    "            The largest timescale to use.\n",
    "        Returns\n",
    "        -------\n",
    "        The input tensor augmented with the sinusoidal frequencies.\n",
    "        \"\"\"\n",
    "        _, timesteps, hidden_dim = tensor.size()\n",
    "\n",
    "        timestep_range = get_range_vector(timesteps, tensor.device).data.float()\n",
    "        # We're generating both cos and sin frequencies,\n",
    "        # so half for each.\n",
    "        num_timescales = hidden_dim // 2\n",
    "        timescale_range = get_range_vector(\n",
    "            num_timescales, tensor.device).data.float()\n",
    "\n",
    "        log_timescale_increments = math.log(\n",
    "            float(max_timescale) / float(min_timescale)) / float(num_timescales - 1)\n",
    "        inverse_timescales = min_timescale *         torch.exp(timescale_range * -log_timescale_increments)\n",
    "\n",
    "        # Broadcasted multiplication - shape (timesteps, num_timescales)\n",
    "        scaled_time = timestep_range.unsqueeze(1) * inverse_timescales.unsqueeze(0)\n",
    "        # shape (timesteps, 2 * num_timescales)\n",
    "        sinusoids = torch.randn(\n",
    "            scaled_time.size(0), 2*scaled_time.size(1), device=tensor.device)\n",
    "        sinusoids[:, ::2] = torch.sin(scaled_time)\n",
    "        sinusoids[:, 1::2] = torch.sin(scaled_time)\n",
    "        if hidden_dim % 2 != 0:\n",
    "            # if the number of dimensions is odd, the cos and sin\n",
    "            # timescales had size (hidden_dim - 1) / 2, so we need\n",
    "            # to add a row of zeros to make up the difference.\n",
    "            sinusoids = torch.cat(\n",
    "                [sinusoids, sinusoids.new_zeros(timesteps, 1)], 1)\n",
    "        return tensor + sinusoids.unsqueeze(0)\n",
    "\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, encoder_outdim, decoder_rnn_hidden, embed_size, output_size):\n",
    "        super(Decoder, self).__init__()\n",
    "\n",
    "        self.rnn_decoder = nn.LSTMCell(decoder_rnn_hidden+embed_size, decoder_rnn_hidden)\n",
    "        self.embedding = nn.Embedding(output_size, embed_size)\n",
    "\n",
    "        self.init_wh = nn.Linear(encoder_outdim, decoder_rnn_hidden)\n",
    "        self.init_wc = nn.Linear(encoder_outdim, decoder_rnn_hidden)\n",
    "        self.init_wo = nn.Linear(encoder_outdim, decoder_rnn_hidden)\n",
    "\n",
    "    def forward(self, encoder_output):\n",
    "        mean_encoder_output = encoder_output.mean(dim=1)\n",
    "        h = torch.tanh(self.init_wh(mean_enc_out))\n",
    "        c = torch.tanh(self.init_wc(mean_enc_out))\n",
    "        o = torch.tanh(self.init_wo(mean_enc_out))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(sign2id, batch):\n",
    "    size = batch[0][0].size()\n",
    "    batch = [img_formula for img_formula in batch\n",
    "             if img_formula[0].size() == size]\n",
    "    batch.sort(key=lambda img_formula: len(img_formula[1].split()),\n",
    "               reverse=True)\n",
    "\n",
    "    imgs, formulas = zip(*batch)\n",
    "    formulas = [formula.split() for formula in formulas]\n",
    "    tgt4training = formulas2tensor([['<s>']+formula for formula in formulas], sign2id)\n",
    "    tgt4cal_loss = formulas2tensor([formula+['</s>'] for formula in formulas], sign2id)\n",
    "    imgs = torch.stack(imgs, dim=0)\n",
    "    return imgs, tgt4training, tgt4cal_loss\n",
    "\n",
    "\n",
    "def formulas2tensor(formulas, sign2id):\n",
    "    \"\"\"convert formula to tensor\"\"\"\n",
    "\n",
    "    batch_size = len(formulas)\n",
    "    max_len = len(formulas[0])\n",
    "    tensors = torch.ones(batch_size, max_len, dtype=torch.long) * PAD_TOKEN\n",
    "    for i, formula in enumerate(formulas):\n",
    "        for j, sign in enumerate(formula):\n",
    "            tensors[i][j] = sign2id.get(sign, UNK_TOKEN)\n",
    "    return tensors\n",
    "\n",
    "\n",
    "def count_parameters(model):\n",
    "    \"\"\"count model parameters\"\"\"\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "\n",
    "def tile(x, count, dim=0):\n",
    "    \"\"\"\n",
    "    Tiles x on dimension dim count times.\n",
    "    \"\"\"\n",
    "    perm = list(range(len(x.size())))\n",
    "    if dim != 0:\n",
    "        perm[0], perm[dim] = perm[dim], perm[0]\n",
    "        x = x.permute(perm).contiguous()\n",
    "    out_size = list(x.size())\n",
    "    out_size[0] *= count\n",
    "    batch = x.size(0)\n",
    "    x = x.view(batch, -1).transpose(0, 1).repeat(count, 1).transpose(0, 1)          .contiguous()          .view(*out_size)\n",
    "    if dim != 0:\n",
    "        x = x.permute(perm).contiguous()\n",
    "    return x\n",
    "\n",
    "\n",
    "def load_formulas(filename):\n",
    "    formulas = dict()\n",
    "    with open(filename) as f:\n",
    "        for idx, line in enumerate(f):\n",
    "            formulas[idx] = line.strip()\n",
    "    print(\"Loaded {} formulas from {}\".format(len(formulas), filename))\n",
    "    return formulas\n",
    "\n",
    "\n",
    "def cal_loss(logits, targets):\n",
    "    \"\"\"args:\n",
    "        logits: probability distribution return by model\n",
    "                [B, MAX_LEN, voc_size]\n",
    "        targets: target formulas\n",
    "                [B, MAX_LEN]\n",
    "    \"\"\"\n",
    "    padding = torch.ones_like(targets) * PAD_TOKEN\n",
    "    mask = (targets != padding)\n",
    "\n",
    "    targets = targets.masked_select(mask)\n",
    "    logits = logits.masked_select(\n",
    "        mask.unsqueeze(2).expand(-1, -1, logits.size(2))\n",
    "    ).contiguous().view(-1, logits.size(2))\n",
    "    logits = torch.log(logits)\n",
    "\n",
    "    assert logits.size(0) == targets.size(0)\n",
    "\n",
    "    loss = F.nll_loss(logits, targets)\n",
    "    return loss\n",
    "\n",
    "\n",
    "def get_checkpoint(ckpt_dir):\n",
    "    \"\"\"return full path if there is ckpt in ckpt_dir else None\"\"\"\n",
    "    if not os.path.isdir(ckpt_dir):\n",
    "        raise FileNotFoundError(\"No checkpoint found in {}\".format(ckpt_dir))\n",
    "\n",
    "    ckpts = [f for f in os.listdir(ckpt_dir) if f.startswith('ckpt')]\n",
    "    if not ckpts:\n",
    "        raise FileNotFoundError(\"No checkpoint found in {}\".format(ckpt_dir))\n",
    "\n",
    "    last_ckpt, max_epoch = None, 0\n",
    "    for ckpt in ckpts:\n",
    "        epoch = int(ckpt.split('-')[1])\n",
    "        if epoch > max_epoch:\n",
    "            max_epoch = epoch\n",
    "            last_ckpt = ckpt\n",
    "    full_path = os.path.join(ckpt_dir, last_ckpt)\n",
    "    print(\"Get checkpoint from {} for training\".format(full_path))\n",
    "    return full_path\n",
    "\n",
    "\n",
    "def schedule_sample(prev_logit, prev_tgt, epsilon):\n",
    "    prev_out = torch.argmax(prev_logit, dim=1, keepdim=True)\n",
    "    prev_choices = torch.cat([prev_out, prev_tgt], dim=1)  # [B, 2]\n",
    "    batch_size = prev_choices.size(0)\n",
    "    prob = Bernoulli(torch.tensor([epsilon]*batch_size).unsqueeze(1))\n",
    "    # sampling\n",
    "    sample = prob.sample().long().to(prev_tgt.device)\n",
    "    next_inp = torch.gather(prev_choices, 1, sample)\n",
    "    return next_inp\n",
    "\n",
    "\n",
    "def cal_epsilon(k, step, method):\n",
    "    \"\"\"\n",
    "    Reference:\n",
    "        Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks\n",
    "        See details in https://arxiv.org/pdf/1506.03099.pdf\n",
    "    \"\"\"\n",
    "    assert method in ['inv_sigmoid', 'exp', 'teacher_forcing']\n",
    "\n",
    "    if method == 'exp':\n",
    "        return k**step\n",
    "    elif method == 'inv_sigmoid':\n",
    "        return k/(k+math.exp(step/k))\n",
    "    else:\n",
    "        return 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class Model:\n",
    "    def __init__(self, optimizer, model, lr_scheduler,\n",
    "                 train_loader, val_loader, args,\n",
    "                 use_cuda=True, max_epoch=25):\n",
    "\n",
    "        self.optimizer = optimizer\n",
    "        self.model = model\n",
    "        self.lr_scheduler = lr_scheduler\n",
    "        self.train_loader = train_loader\n",
    "        self.val_loader = val_loader\n",
    "        self.args = args\n",
    "\n",
    "        self.step = 0\n",
    "        self.epoch = 0\n",
    "        self.total_step = 0\n",
    "        self.last_epoch = max_epoch\n",
    "        self.best_val_loss = None\n",
    "        self.device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "\n",
    "    def train(self):\n",
    "        message = \"Epoch {}, step:{}/{} {:.2f}%, Loss:{:.4f}, Perplexity:{:.4f}\"\n",
    "\n",
    "        while self.epoch <= self.last_epoch:\n",
    "            self.model.train()\n",
    "            losses = 0.0\n",
    "            for imgs, tgt4training, tgt4cal_loss in self.train_loader:\n",
    "                self.optimizer.zero_grad()\n",
    "\n",
    "                imgs = imgs.to(self.device)\n",
    "                tgt4training = tgt4training.to(self.device)\n",
    "                tgt4cal_loss = tgt4cal_loss.to(self.device)\n",
    "                epsilon = cal_epsilon(\n",
    "                    self.args.decay_k, self.total_step, self.args.sample_method)\n",
    "                pred = self.model(imgs, tgt4training, epsilon)\n",
    "\n",
    "                # calculate loss\n",
    "                loss = cal_loss(pred, tgt4cal_loss)\n",
    "                self.step += 1\n",
    "                self.total_step += 1\n",
    "                loss.backward()\n",
    "                clip_grad_norm_(self.model.parameters(), self.args.clip)\n",
    "                self.optimizer.step()\n",
    "\n",
    "                step_loss = loss.item()\n",
    "                losses += step_loss\n",
    "\n",
    "                # log message\n",
    "                if self.step % self.args.print_freq == 0:\n",
    "                    avg_loss = losses / self.args.print_freq\n",
    "                    print(message.format(\n",
    "                        self.epoch, self.step, len(self.train_loader),\n",
    "                        100 * self.step / len(self.train_loader),\n",
    "                        avg_loss,\n",
    "                        2**avg_loss\n",
    "                    ))\n",
    "                    losses = 0.0\n",
    "\n",
    "            val_loss = self.validate()\n",
    "            self.lr_scheduler.step(val_loss)\n",
    "\n",
    "            self.save_model('ckpt-{}-{:.4f}'.format(self.epoch, val_loss))\n",
    "            self.epoch += 1\n",
    "            self.step = 0\n",
    "\n",
    "    def validate(self):\n",
    "        self.model.eval()\n",
    "        val_total_loss = 0.0\n",
    "        mes = \"Epoch {}, validation average loss:{:.4f}, Perplexity:{:.4f}\"\n",
    "        with torch.no_grad():\n",
    "            for imgs, tgt4training, tgt4cal_loss in self.val_loader:\n",
    "                imgs = imgs.to(self.device)\n",
    "                tgt4training = tgt4training.to(self.device)\n",
    "                tgt4cal_loss = tgt4cal_loss.to(self.device)\n",
    "\n",
    "                epsilon = cal_epsilon(\n",
    "                    self.args.decay_k, self.total_step, self.args.sample_method)\n",
    "                pred = self.model(imgs, tgt4training, epsilon)\n",
    "                loss = cal_loss(pred, tgt4cal_loss)\n",
    "                val_total_loss += loss\n",
    "            avg_loss = val_total_loss / len(self.val_loader)\n",
    "            print(mes.format(\n",
    "                self.epoch, avg_loss, 2**avg_loss\n",
    "            ))\n",
    "        if self.best_val_loss is None or avg_loss < self.best_val_loss:\n",
    "            self.best_val_loss = avg_loss\n",
    "            self.save_model('best_ckpt')\n",
    "        return avg_loss\n",
    "    \n",
    "    def predict(self):\n",
    "        self.model.eval()\n",
    "        predictions = []\n",
    "        with torch.no_grad():\n",
    "            for imgs, tgt4training, tgt4cal_loss in self.val_loader:\n",
    "                imgs = imgs.to(self.device)\n",
    "                tgt4training = tgt4training.to(self.device)\n",
    "                tgt4cal_loss = tgt4cal_loss.to(self.device)\n",
    "\n",
    "                epsilon = cal_epsilon(\n",
    "                    self.args.decay_k, self.total_step, self.args.sample_method)\n",
    "                pred = self.model(imgs, tgt4training, epsilon)\n",
    "                # TODO\n",
    "                predictions.append(pred)\n",
    "            # return        \n",
    "\n",
    "    def save_model(self, model_name):\n",
    "        if not os.path.isdir(self.args.save_dir):\n",
    "            os.makedirs(self.args.save_dir)\n",
    "        save_path = join(self.args.save_dir, model_name+'.pt')\n",
    "        print(\"Saving checkpoint to {}\".format(save_path))\n",
    "\n",
    "        torch.save({\n",
    "            'model_state_dict': self.model.state_dict(),\n",
    "            'optimizer_state_dict': self.optimizer.state_dict(),\n",
    "            'lr_sche': self.lr_scheduler.state_dict(),\n",
    "            'args': self.args\n",
    "        }, save_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO\n",
    "Arguments = namedtuple(\"Arguments\", [\"emb_dim\", \"dec_rnn_h\", \"add_position_features\",\"max_len\",\n",
    "                       \"dropout\", \"cuda\", \"batch_size\", \"epoches\", \n",
    "                       \"lr\", \"min_lr\", \"sample_method\", \"decay_k\",\n",
    "                        \"lr_decay\", \"lr_patience\", \"clip\", \"save_dir\",\n",
    "                       \"print_freq\", \"seed\", \"from_check_point\"]\n",
    "            )\n",
    "\n",
    "args = Arguments(80, 512, True, 150, 0.2, True, 32, 5, 3e-4,\n",
    "                         3e-5, \"teacher_forcing\", 1., 0.5, 3, 2.0,\n",
    "                         os.getcwd(), 100, 2020, True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def train(args):\n",
    "    max_epoch = args.epoches\n",
    "    from_check_point = args.from_check_point\n",
    "    if from_check_point:\n",
    "        checkpoint_path = get_checkpoint(args.save_dir)\n",
    "        checkpoint = torch.load(checkpoint_path)\n",
    "        args = checkpoint['args']\n",
    "    print(\"Training args:\", args)\n",
    "\n",
    "    torch.manual_seed(args.seed)\n",
    "    torch.cuda.manual_seed(args.seed)\n",
    "\n",
    "    use_cuda = True if args.cuda and torch.cuda.is_available() else False\n",
    "    device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "\n",
    "    # data loader\n",
    "    print(\"Construct data loader...\")\n",
    "    train_loader = DataLoader(\n",
    "        train_set,\n",
    "        batch_size=args.batch_size,\n",
    "        collate_fn=partial(collate_fn, vocab.sign2id),\n",
    "        pin_memory=True if use_cuda else False,\n",
    "        num_workers=4)\n",
    "    val_loader = DataLoader(\n",
    "        validate_set,\n",
    "        batch_size=args.batch_size,\n",
    "        collate_fn=partial(collate_fn, vocab.sign2id),\n",
    "        pin_memory=True if use_cuda else False,\n",
    "        num_workers=4)\n",
    "\n",
    "    # construct model\n",
    "    print(\"Construct model\")\n",
    "    vocab_size = len(vocab)\n",
    "    model = Image2LatexModel(\n",
    "        vocab_size, args.emb_dim, args.dec_rnn_h,\n",
    "        add_pos_feat=args.add_position_features,\n",
    "        dropout=args.dropout\n",
    "    )\n",
    "    model = model.to(device)\n",
    "    print(\"Model Settings:\")\n",
    "    print(model)\n",
    "\n",
    "    # construct optimizer\n",
    "    optimizer = optim.Adam(model.parameters(), lr=args.lr)\n",
    "\n",
    "    lr_scheduler = ReduceLROnPlateau(\n",
    "        optimizer,\n",
    "        \"min\",\n",
    "        factor=args.lr_decay,\n",
    "        patience=args.lr_patience,\n",
    "        verbose=True,\n",
    "        min_lr=args.min_lr)\n",
    "\n",
    "    if from_check_point:\n",
    "        model.load_state_dict(checkpoint['model_state_dict'])\n",
    "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "        epoch = checkpoint['epoch']\n",
    "        lr_scheduler.load_state_dict(checkpoint['lr_sche'])\n",
    "        # init model from checkpoint\n",
    "        model = Model(optimizer, model, lr_scheduler,\n",
    "                          train_loader, val_loader, args,\n",
    "                          use_cuda=use_cuda,\n",
    "                          init_epoch=epoch, last_epoch=max_epoch)\n",
    "    else:\n",
    "        model = Model(optimizer, model, lr_scheduler,\n",
    "                          train_loader, val_loader, args,\n",
    "                          use_cuda=use_cuda,\n",
    "                          init_epoch=1, last_epoch=args.epoches)\n",
    "\n",
    "    model.train()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "args_cuda = True\n",
    "args_model_path = \"best_ckpt.pt\"\n",
    "args_batch_size = 32\n",
    "args_result_path = \"result.txt\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def test():\n",
    "    checkpoint = torch.load(join(args_model_path))\n",
    "    model_args = checkpoint['args']\n",
    "    use_cuda = True if args_cuda and torch.cuda.is_available() else False\n",
    "\n",
    "    data_loader = DataLoader(\n",
    "        test_set,\n",
    "        batch_size=args_batch_size,\n",
    "        collate_fn=partial(collate_fn, vocab.sign2id),\n",
    "        pin_memory=True if use_cuda else False,\n",
    "        num_workers=4\n",
    "    )\n",
    "\n",
    "    model = Model(optimizer, model, lr_scheduler,\n",
    "        train_loader, val_loader, args,\n",
    "        use_cuda=use_cuda,\n",
    "        init_epoch=epoch, last_epoch=max_epoch\n",
    "    )\n",
    "\n",
    "    predictions = model.predict()\n",
    "    # TODO\n",
    "    # Get the index of corresponding formulas and match them simply if they are same of not\n",
    "    # Write the predicted formula's index to result file along with image name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}