{"cells":[{"metadata":{},"cell_type":"markdown","source":"# CNN to GRU with Attention and Teacher Forcing\n\nThis is a baseline for using aa CNN into LSTM type model for image captioning. The LSTM decoder uses an attention over the CNN activations for decoding. Teacher Forcing is used during training to speed up convergence.\n\nTo aavoid having a long notebook run time, this notebook by default doesn't use the full training dataset, nor does it predict on all items in the test dataset. Where this happens, I've added the full dataset version commented out. \n\nRunning this notebook on the entire dataset (four training epochs) took about 30 hours on a 2080 ti GPU. The resulting submission scored 55.6 on the public leaderboard.\n\nThis notebook draws from [this example](https://github.com/fastai/course-nlp/blob/master/7b-seq2seq-attention-translation.ipynb)\n\nNote that installing RDKit on a GPU accelerated instance usually takes around 15 minutes due to conda install issues."},{"metadata":{"trusted":true},"cell_type":"code","source":"!conda install -y -c rdkit rdkit","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom multiprocessing import Pool\nimport os\nfrom skimage import color\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nimport skimage\nimport math\nimport random\nfrom functools import partial\nimport gc\n\nfrom rdkit import Chem\nfrom rdkit.Chem import Draw\nfrom rdkit import RDLogger\n\nRDLogger.DisableLog('rdApp.*') \n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.transforms import *\nfrom torchvision.models import resnet34\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from fastai import *\nfrom fastai.data.core import DataLoaders\nfrom fastai.learner import *\nfrom fastai.callback.all import *\nfrom fastai.losses import *\nfrom fastai.metrics import *","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def smile_to_mol(smile):\n    return Chem.MolFromSmiles(smile)\n\ndef inchi_to_mol(inchi):\n    return Chem.inchi.MolFromInchi(inchi)\n\ndef mol_to_smile(mol):\n    return Chem.MolToSmiles(mol)\n\ndef mol_to_inchi(mol):\n    return Chem.inchi.MolToInchi(mol)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def pad_square(image):\n    w,h = image.shape[-2:]\n    max_wh = max([w,h])\n    hp = int((max_wh-w)/2)\n    vp = int((max_wh-h)/2)\n    padding = (vp, hp)\n    return Pad(padding, padding_mode='edge')(image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageDataset(Dataset):\n    def __init__(self, filenames, file_prefix, y_vals, itos, stoi, size, return_inchi=True):\n        self.filenames = filenames\n        self.file_prefix = file_prefix\n        self.y_vals = y_vals\n        self.itos = itos\n        self.stoi = stoi\n        self.size = size\n        self.return_inchi = return_inchi\n        self.resize = transforms.Resize((size, size))\n        \n    def __len__(self):\n        return len(self.filenames)\n    \n    def __getitem__(self, idx):\n        \n        filename = self.filenames[idx]\n\n        image = np.array(Image.open(f'{self.file_prefix}/{filename[0]}/{filename[1]}/{filename[2]}/{filename}.png'))\n        image = torch.FloatTensor(image)[None,:,:].repeat((3,1,1))/255.\n        \n        if not image.shape[-1]==image.shape[-2]:\n            image = pad_square(image)\n            \n        image = self.resize(image)\n            \n        output = self.y_vals[idx]\n\n        if 'InChI' in output:\n            is_inchi = True\n        else:\n            is_inchi = False\n\n        if (is_inchi and self.return_inchi) or (not is_inchi and not self.return_inchi):\n            out_string = output\n\n        elif is_inchi and not self.return_inchi:\n            out_string = mol_to_smile(inchi_to_mol(output))\n\n        else:\n            out_string = mol_to_inchi(smile_to_mol(output))\n\n        out_ints = [self.stoi['bos']] + [stoi[i] for i in out_string] + [self.stoi['eos']]\n\n        return image.data, out_ints","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data Setup\n\nOne thing I've played around with is comparing generating InChI strings directly, or generating something like a SMILES string (a shorter sequence) that is then converted to an InChI. It seems that a SMILES approach results in more correct structures (in terms of resolving to a valid compound), but predicting InChIs gives better performance overall due to the fact that any incorrect SMILES strings are basically lost (ie can't be converted to an InChI).\n\nThe `ImageDataset` dataset will return InChI strings by default, but will return SMILES strings if `return_inchi=False` is passed.\n\nI also decided to remove compounds from the dataset that had an InChI string longer than 250 characters. These long tail sequences can sneak up on you and give you a Cuda memory error."},{"metadata":{"trusted":true},"cell_type":"code","source":"# smiles itos\n# itos = ['bos', 'pad', 'eos',\n#  'N', '1', '(', 'P', '8', 'S',\n#  'H', ']', 'B', '#', '=', ')', \n#  's', '-', 'r', '7', '4',\n#  '3', '[', 'c', '6', 'n', '2',\n#  'i', 'o', 'O', '@', 'F', '/',\n#  'I', 'l', '\\\\', '5', '+', 'C',\n#  '%', '.', '0', '9', 'b', 'p']\n\n# inchi itos\nitos = ['bos', 'pad', 'eos',\n '=', '/', '+',\n '6', 'O', 'C', '5', 'D', ',', 'c',\n '7', '2', 'b', 'i', '0', 'B', '1',\n '8', 'h', ')', 'n', '9', '-', 'S',\n '(', '4', 'H', 'm', '3', 'F', 't',\n 'P', 'l', 'N', 'T', 's', 'r', 'I']\nstoi = {itos[i]:i for i in range(len(itos))}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# loading just the first 10000 rows, uncomment full load below for using the entire dataset\ndf = next(pd.read_csv('../input/bms-molecular-translation/train_labels.csv', chunksize=10000))\n\n# df = pd.read_csv('../input/bms-molecular-translation/train_labels.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"lens = df.InChI.map(lambda x: len(x))\ndf = df[lens<250]\ndf.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cut = int(0.97*len(df))\ndf_train = df[:cut]\ndf_valid = df[cut:]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_pefix = '../input/bms-molecular-translation/train'\n\n# Note image size is hard-coded to 256x256\ntrain_data = ImageDataset(df_train.image_id.values, train_pefix, df_train.InChI.values,\n                          itos, stoi, 256, return_inchi=True)\n\nvalid_data = ImageDataset(df_valid.image_id.values, train_pefix, df_valid.InChI.values,\n                          itos, stoi, 256, return_inchi=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.imshow(train_data[0][0].permute(1,2,0))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_function(batch, pad=1):\n    \n    images = torch.stack([i[0] for i in batch])\n    \n    max_len = max([len(i[1]) for i in batch])\n    res_y = torch.zeros(len(batch), max_len).long() + pad\n    \n    for i, s in enumerate(batch):\n        res_y[i,:len(s[1])] = torch.LongTensor(s[1])\n    \n    return images, res_y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dl = DataLoader(train_data, batch_size=120, collate_fn=collate_function, shuffle=True, num_workers=4)\nvalid_dl = DataLoader(valid_data, batch_size=120, collate_fn=collate_function, shuffle=False, num_workers=4)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dls = DataLoaders(train_dl, valid_dl)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"x,y = next(iter(dls.loaders[0]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"x.shape, y.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"x.device","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model Architecture\n\nThe model first maps the input images down to a set of activations using the `image_encoder`, which in this notebook is a Resnet34 model. Then the model decodes the output sequence one token at a time using a GRU. At each decoding step, the model computes attention over the image activations.\n\nThis model also has a teacher forcing parameter `pr_force`. At each decoding step, with the probability of `pr_force`, the model is given the ground truth answer. This helps speed up convergence early in training. We start with `pr_force=1` and decay the value to zero over the course of training."},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageCaption(nn.Module):\n    def __init__(self, image_encoder, d_enc_out, nh, emb_sz_dec, voc_sz_dec, out_sl, nl=2, bos_idx=0, pad_idx=1):\n        super().__init__()\n        \n        self.nl = nl\n        self.nh = nh\n        self.bos_idx = bos_idx\n        self.pad_idx = pad_idx\n        self.out_sl = out_sl\n        self.emb_sz_dec = emb_sz_dec\n        self.voc_sz_dec = voc_sz_dec\n        self.pr_force = 0.\n        \n        self.encoder = image_encoder\n        self.init_hidden = nn.Linear(d_enc_out, nh*nl*2)\n        self.out_enc = nn.Linear(2*nh, self.emb_sz_dec, bias=False)\n        \n        self.emb_dec = nn.Embedding(voc_sz_dec, emb_sz_dec)\n        self.gru_dec = nn.GRU(self.emb_sz_dec + 2*nh, self.emb_sz_dec, num_layers=nl,\n                              dropout=0.1, batch_first=True)\n        self.out_drop = nn.Dropout(0.35)\n        self.out = nn.Linear(self.emb_sz_dec, self.voc_sz_dec)\n        self.out.weight.data = self.emb_dec.weight.data\n        \n        self.enc_projection = nn.Linear(d_enc_out, nh*2)\n        self.enc_att = nn.Linear(nh*2, self.emb_sz_dec, bias=False)\n        self.hid_att = nn.Linear(self.emb_sz_dec, self.emb_sz_dec)\n        self.V =  self.init_param(self.emb_sz_dec)\n        \n    def decoder(self, dec_inp, hid, enc_att, enc_out):\n        hid_att = self.hid_att(hid[-1])\n        u = torch.tanh(enc_att + hid_att[:,None])\n        attn_wgts = F.softmax(u @ self.V, 1)\n        ctx = (attn_wgts[...,None] * enc_out).sum(1)\n        emb = self.emb_dec(dec_inp)\n        outp, hid = self.gru_dec(torch.cat([emb, ctx], 1)[:,None], hid)\n        outp = self.out(self.out_drop(outp[:,0]))\n        return hid, outp\n        \n    def forward(self, x, targ=None):\n        \n        enc_out = self.encoder(x)\n        bs, sl, _ = enc_out.shape\n        \n        out_sl = self.out_sl if targ is None else targ.shape[1]\n        \n        mean_encoder_out = enc_out.mean(dim=1)\n        hid = self.init_hidden(mean_encoder_out)\n        hid = hid.view(2, self.nl, bs, self.nh).permute(1,2,0,3).contiguous()\n        hid = hid.view(self.nl, bs, 2*self.nh)\n        hid = self.out_enc(hid)\n\n        enc_out = self.enc_projection(enc_out)\n        enc_att = self.enc_att(enc_out)\n        \n        dec_inp = x.new_zeros(bs).long() + self.bos_idx\n        res = []\n        \n        for i in range(out_sl):\n            hid, outp = self.decoder(dec_inp, hid, enc_att, enc_out)\n            res.append(outp)\n            dec_inp = outp.max(1)[1]\n            if (dec_inp==self.pad_idx).all(): break\n            if (targ is not None) and (random.random()<self.pr_force):\n                if i>=targ.shape[1]: continue\n                dec_inp = targ[:,i]\n        return torch.stack(res, dim=1)\n            \n    def init_param(self, *sz): return nn.Parameter(torch.randn(sz)/math.sqrt(sz[0]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageEncoder(nn.Module):\n    # wrapper for torchvision model\n    def __init__(self, image_encoder):\n        super().__init__()\n        \n        modules = list(image_encoder.children())[:-2]\n        self.image_encoder = nn.Sequential(*modules)\n        \n    def forward(self, x):\n        x = self.image_encoder(x) # (bs, ch, h, w)\n        x = x.permute(0, 2, 3, 1) # (bs, h, w, ch)\n        x = x.view(x.size(0), -1, x.size(-1)) # (bs, h*w, ch)\n\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_encoder = ImageEncoder(resnet34())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"d_enc_out = 512\nnh = 256\nemb_sz_dec = 128\nvoc_sz_dec = len(itos)\nout_sl = 220\n\nic = ImageCaption(image_encoder, d_enc_out, nh, emb_sz_dec, voc_sz_dec, out_sl)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seq2seq_loss(out, targ, pad_idx=1):\n    bs,targ_len = targ.size()\n    _,out_len,vs = out.size()\n    if targ_len>out_len: out  = F.pad(out,  (0,0,0,targ_len-out_len,0,0), value=pad_idx)\n    if out_len>targ_len: targ = F.pad(targ, (0,out_len-targ_len,0,0), value=pad_idx)\n    return CrossEntropyLossFlat()(out, targ)\n\ndef seq2seq_acc(out, targ, pad_idx=1):\n    bs,targ_len = targ.size()\n    _,out_len,vs = out.size()\n    if targ_len>out_len: out  = F.pad(out,  (0,0,0,targ_len-out_len,0,0), value=pad_idx)\n    if out_len>targ_len: targ = F.pad(targ, (0,out_len-targ_len,0,0), value=pad_idx)\n    out = out.argmax(2)\n    return (out==targ).float().mean()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TeacherForcingCallback(Callback):\n    # teacher forcing callback\n    def __init__(self, start_batch, end_batch):\n        self.start_batch = start_batch\n        self.end_batch = end_batch\n\n    def before_batch(self):\n        \n        n_iter = self.train_iter\n        \n        if n_iter < self.start_batch:\n            self.learn.model.pr_force = 1.\n            \n        elif n_iter > self.end_batch:\n            self.learn.model.pr_force = 0.\n            \n        else:\n            self.learn.model.pr_force = 1 - (n_iter - self.start_batch)/(self.end_batch - self.start_batch)\n        \n        if self.training:\n            x,y = self.x, self.y\n            self.learn.xb = (x,y)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn = Learner(dls, ic, loss_func=seq2seq_loss, cbs=[CudaCallback, \n                                                      TeacherForcingCallback(start_batch=4000, end_batch=25000)], \n                metrics=[seq2seq_acc, CorpusBLEUMetric(len(itos))])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.lr_find()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.fit_one_cycle(4, 3e-3)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prediction"},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df = pd.read_csv('../input/bms-molecular-translation/sample_submission.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_prefix = '../input/bms-molecular-translation/test'\n\ntest_df = test_df[:1000] # predict on first 1000\n\ntest_data = ImageDataset(test_df.image_id.values, test_prefix, test_df.InChI.values,\n                          itos, stoi, 256, return_inchi=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_dl = DataLoader(test_data, batch_size=256, collate_fn=collate_function, shuffle=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.model.eval();","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"preds_list = []\n\nwith torch.no_grad():\n    for i, batch in enumerate(test_dl):\n        if i%500 == 0:\n            print(i)\n        \n        x,y = batch\n        preds = learn.model(x.cuda())\n        preds = F.softmax(preds, -1).argmax(-1)\n        preds_list.append(preds.detach().cpu())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pred_strs = []\n\nfor k, pred in enumerate(preds_list):\n    if k%1000 == 0:\n        print(k)\n    \n    gc.collect()\n    for p in pred:\n        pred_str = [itos[i] for i in p][1:]\n        \n        if 'eos' in pred_str:\n            pred_str = pred_str[:pred_str.index('eos')]\n            \n        if 'pad' in pred_str:\n            pred_str = pred_str[:pred_str.index('pad')]\n            \n        pred_str = ''.join(pred_str)\n        \n        pred_strs.append(pred_str)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_df['preds'] = pred_strs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission = test_df[['image_id', 'preds']]\nsubmission.columns = ['image_id', 'InChI']","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submission.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Suggested Improvements\n\nThis notebook is fairly basic. There's a lot of simple improvements to be made from here.\n\n### Data\n\nNo data augmentation is used. This could be added, along with synthetic data from other datasets (ie use RDKit to generate images). The decision to predict SMILES instead of InChI can also be revisited.\n\n### Model\n\nThe current setup is hard-coded to work with 256x256 images. This can be changed by adding an adaptive pooling layer after the CNN encoder. The decoder is actually fairly small. Most of the parameters are in the image encoder. Model size can be expanded on. It would also be good to investigate larger resnet backbones.\n\n### Metrics\n\nI used BLEU and accuracy to evaluate the model. A better metric would be to actually convert SMILES to InChI and calculate the actual Levenshtein distance.\n\n### Inference\n\nInference shown here is extremely simple, using single prediction argmax samping. This could be improved by adding beam search decoding and test-time augmentation to the input images. If you wanted to get really fancy, you could use TTA+beam search to generate several pedictions, then find the concensus string between all pedictions."},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}