{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"This notebook borrowed many codes from： \n1. [PyTorch Tutorial](https://pytorch.org/tutorials/beginner/torchtext_translation.html)\n2. [Blog: solving-an-image-captioning-task-using-deep-learning](https://www.analyticsvidhya.com/blog/2018/04/solving-an-image-captioning-task-using-deep-learning/)","metadata":{}},{"cell_type":"markdown","source":"This is a baseline model that using CNN and LSTM to translate image into chemical structure. The total training images are 2424186, that's enormous. So I decided to sample 5% of full training data for training.\n![](https://github.com/yunjey/pytorch-tutorial/raw/master/tutorials/03-advanced/image_captioning/png/model.png)","metadata":{}},{"cell_type":"code","source":"!pip install timm","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport torch\nimport numpy as np\nimport random\n\nfrom torch.utils.data import DataLoader\nimport cv2\n\nimport timm\nfrom pprint import pprint\n# model_names = timm.list_models(pretrained=True)\n# pprint(model_names)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def seed_everything(seed=99):\n    torch.manual_seed(seed)\n    np.random.seed(seed)\n    random.seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\n    \nseed_everything()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cfg = {\n    'version':'version2',\n    'device': 'cuda' if torch.cuda.is_available() else 'cpu',\n    'train_img':'../input/bms-molecular-translation/train',\n    'train_anno':'../input/bms-molecular-translation/train_labels.csv',\n    'sample_ratio': 0.1,\n    'backbone':'efficientnet_b0',\n    'pretrianed':True,\n    'vocab_size':38, # 36 unique char + [SOS, EOS]\n    'max_len':150,\n    'embed_size':16,\n    'hidden_size':64,\n    'image_size':[128, 128],\n    'batch_size':64,\n    'num_workers':6,\n    'n_epochs':10,\n    'lr':1e-3,\n    'min_lr':1e-6,\n    'patience':3,\n    'TTA':5,\n}\n\nif not os.path.isdir(cfg['version']):\n    os.mkdir(cfg['version'])\n    print('create dir')\ncfg","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### build vocab\nadd start token 'SOS 'and end token 'EOS' into vocab","metadata":{}},{"cell_type":"code","source":"class Lang():\n    \"\"\"\n    seq: chemical structure, shape like: 'InChI=1S...'\n    \"\"\"\n    def __init__(self):\n        start_end_token = ['PAD', 'SOS', 'EOS']\n        self.vocab = start_end_token + ['(', ')', '+', ',', '-', '/', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'B', 'C', 'D', 'F', 'H', 'I', 'N', 'O', 'P', 'S', 'T', 'b', 'c', 'h', 'i', 'l', 'm', 'r', 's', 't']\n        self.token_to_idx = {token:idx for idx, token in enumerate(self.vocab)}\n        \n    def seq_to_idx(self, seq):\n        idxs = []\n        seq = seq.replace('InChI=1S', '')  # remove head\n        for token in seq:\n            idxs += [self.token_to_idx[token]]\n        source = [self.token_to_idx['SOS']] + idxs\n        target = idxs + [self.token_to_idx['EOS']]\n        return source, target      # return source and target\n    \n    def idx_to_seq(self, idxs):\n        idxs = self.rm_re_idxs(idxs)  # remove repeated text\n        # add head, remove SOS and EOS\n        seq = 'InChI=1S'+''.join([self.vocab[idx] for idx in idxs])\n        \n        return seq.replace('SOS','').replace('EOS','')\n    \n    def rm_re_idxs(self, idxs):\n        # remove repeated text\n        new_idxs = []\n        for idx in idxs:\n            if idx == self.token_to_idx['EOS'] or idx == self.token_to_idx['PAD']:\n                break\n            new_idxs += [idx]\n        return new_idxs\n    \nlang = Lang()\ncfg['vocab_size'] = len(lang.vocab)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"sentence to idxs, then restore","metadata":{}},{"cell_type":"code","source":"sentence = 'InChI=1S/C13H20OS/c1-9(2)8-15-13-6-5-10(3)7-12(13)11(4)14/h5-7,9,11,14H,8H2,1-4H3'\nsource, target = lang.seq_to_idx(sentence)\nlang.idx_to_seq(source)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"sampling data, then split it into train and test dataset","metadata":{}},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split\n\nfull_df = pd.read_csv(cfg['train_anno'])\n\nsample_df = full_df.sample(frac=cfg['sample_ratio']).reset_index(drop=True)\n\nX = list(range(len(sample_df)))\ntrain_indexs, val_indexs = train_test_split(X, test_size=0.33, random_state=42)\ntrain_df, val_df = sample_df.loc[train_indexs], sample_df.loc[val_indexs]\ntrain_df.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"augmentation pipeline","metadata":{}},{"cell_type":"code","source":"import albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\ntransform = A.Compose([\n    A.Resize(cfg['image_size'][0], cfg['image_size'][1]),\n    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    ToTensorV2()\n])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"dataloader pipeline","metadata":{}},{"cell_type":"code","source":"def convert_image_id_2_path(img_dir:str, image_id: str) -> str:\n    return \"{}/{}/{}/{}/{}.png\".format(img_dir,\n        image_id[0], image_id[1], image_id[2], image_id \n    )\n\ndef data_process(df, img_dir):\n    data = []\n    for idx in range(len(df)):\n        image_id = df.iloc[idx].image_id\n        img_path = convert_image_id_2_path(img_dir, image_id)\n        seq = df.iloc[idx].InChI\n        data += [(img_path, seq)]\n        \n    return data\n    \ndef generate_batch(data_batch, tfs=transform, train=True):\n    \n    img_batch, source_batch, target_batch = [], [], []\n    for (img_path, seq) in data_batch:\n        img = cv2.imread(img_path)\n        if tfs:\n            img = tfs(image=img)['image']\n            \n        img_batch += [img]\n        source, target = lang.seq_to_idx(seq)\n        source_batch += [torch.tensor(source, dtype=torch.long)]\n        target_batch += [torch.tensor(target, dtype=torch.long)]\n    return img_batch, source_batch, target_batch\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_data = data_process(train_df, img_dir=cfg['train_img'])\nval_data = data_process(val_df, img_dir=cfg['train_img'])\ntrain_iter = DataLoader(train_data, batch_size=cfg['batch_size'],\n                        shuffle=True, collate_fn=generate_batch, num_workers=cfg['num_workers'])\nvalid_iter = DataLoader(val_data, batch_size=cfg['batch_size'],\n                        shuffle=False, collate_fn=generate_batch, num_workers=cfg['num_workers'])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"display samples","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\ndef visualize_batch(image, labels):\n    plt.figure(figsize=(16, 12))\n    for ind, (image, label) in enumerate(zip(image, labels)):\n        plt.subplot(3, 3, ind + 1)\n        plt.imshow(image.permute(1, 2, 0))\n        plt.title(f\"{label[:10]}...\", fontsize=10)\n        plt.axis(\"off\")\n    \n    plt.show()\n\nfor imgs, source, target in train_iter:\n    visualize_batch(imgs[:3], target[:3])\n    break","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pack_sequence\nimport torch.nn as nn\n\n# Conv + LSTM\nclass Generator(nn.Module):\n    \"\"\"\n    Conv encoder LSTM decoder\n    x: imgs\n    seqs: padded sequence of idxs that is a batch of InChl\n    lengths: batch of seqs length\n    \"\"\"\n    def __init__(self, cfg):\n        super().__init__()\n        self.cfg = cfg\n        self.backbone = timm.create_model(cfg['backbone'], pretrained=cfg['pretrianed'], num_classes=cfg['hidden_size'])\n        self.emb_layer = nn.Embedding(cfg['vocab_size'], cfg['embed_size'])\n        self.lstm = nn.LSTM(cfg['embed_size'], cfg['hidden_size'])\n        self.out_layer = nn.Linear(cfg['hidden_size'], cfg['vocab_size'])\n        \n    def forward(self, x, source_padded, lens):\n        batch_size = x.size(0)\n        features = self.backbone(x)\n        (h, c) = self.init_state(batch_size)\n        states = (features.unsqueeze(0), c) \n        input_embedding = self.emb_layer(source_padded)    # embedding: len, batch, size\n        packed = pack_padded_sequence(input_embedding, lens)\n        out_packed, _ = self.lstm(packed, states)\n        outputs = self.out_layer(out_packed[0])\n        return outputs   # packed len, batch\n    \n    def greedy_decode(self, x, lens=None):\n        \"\"\"Greedy search\"\"\"\n        max_len = lens if lens is not None else self.cfg['max_len']\n        sampled_ids = []\n        batch_size = x.size(0)\n        logits = torch.zeros(max_len, batch_size, cfg['vocab_size'], device=self.cfg['device'])\n        features = self.backbone(x)\n        (h, c) = self.init_state(batch_size)\n        states = (features.unsqueeze(0), c)\n        # create start \n        SOS = torch.tensor(lang.token_to_idx['SOS'], dtype=torch.long, device=self.cfg['device']).expand((1, batch_size)) # create SOS\n        last_emb = self.emb_layer(SOS) \n        for i in range(max_len):                      # maximum sampling length\n            hiddens, states = self.lstm(last_emb, states)         # (1, batch, hidden_size), \n            outputs = self.out_layer(hiddens.squeeze())           # (batch_size, vocab_size)\n            curr = torch.argmax(outputs, dim=-1, keepdim=True)    # to idxs\n            curr = curr.reshape(1, batch_size)\n            last_emb = self.emb_layer(curr)  # (1, batch, embed_size)\n            \n            sampled_ids.append(curr)\n            logits[i] = outputs\n            if curr[:,0] == lang.token_to_idx['EOS'] or curr[:,0] == lang.token_to_idx['PAD']:\n                break\n        sampled_ids = torch.cat(sampled_ids)                    \n        return sampled_ids, logits  # (seq, batch)\n    \n    def init_state(self, batch_size):\n        return (\n            torch.zeros(1, batch_size, self.cfg['hidden_size']).to(self.cfg['device']),\n            torch.zeros(1, batch_size, self.cfg['hidden_size']).to(self.cfg['device'])\n        )","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import Levenshtein\n\n# metric\ndef get_score(y_pred, y_true):\n    scores = []\n    for true, pred in zip(y_true, y_pred):\n        score = Levenshtein.distance(true, pred)\n        scores.append(score)\n    avg_score = np.mean(scores)\n    return avg_score\n\ndef get_accuracy(y_pred, y_true):\n    return (y_pred == y_true).sum() / len(y_pred)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def pad_packed(packed_data, batch_sizes):\n    max_len, N = len(batch_sizes), batch_sizes[0]\n    padded_packed = np.zeros((max_len, N), dtype=np.int64)\n    lens = np.zeros(N, dtype=np.int64)\n    curr_idx = 0\n    for step, batch_size in enumerate(batch_sizes):\n        curr_batch = packed_data[curr_idx: curr_idx+batch_size]\n        padded_packed[step, :batch_size] = curr_batch\n        lens[:batch_size] += 1\n        curr_idx += batch_size\n        \n    return padded_packed, lens\n\n# preprocess data\ndef load_data(data):\n    imgs_batch, source_batch, target_batch = data\n    imgs_batch = torch.stack(imgs_batch, dim=0).to(cfg['device'])\n    lens = torch.tensor([len(item) for item in source_batch])\n    \n     # sort by length\n    sorted_lens, sorted_indices = torch.sort(lens, descending=True)\n    imgs_batch = imgs_batch[sorted_indices].to(cfg['device'])\n    source_padded = pad_sequence(source_batch)  # padding\n    source_padded = source_padded[:,sorted_indices].to(cfg['device'])\n    target_batch = [target_batch[idx] for idx in sorted_indices]\n    return imgs_batch, source_padded, sorted_lens, target_batch","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import math\n\nclass Eearly_Stopping():\n    def __init__(self, mode: str='min', patience: int=10, apply=True):\n        self.best = math.inf\n        self.mode = mode\n        self.base_patience = patience\n        self.patience = patience\n        self.apply=apply\n        if self.patience <= 0:\n            raise Exception(\"Invalid patience!\", patience)\n            \n    def step(self, model, monitor):\n        if monitor < self.best:\n            self.best = monitor\n            torch.save(model, f'model.pth') # save model\n            self.patience = self.base_patience\n        else:\n            self.patience -= 1 if self.apply else -1\n        \n        return True if self.patience==0 else False","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch.optim as optim\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau \n\n# training setting\ngenerator = Generator(cfg).to(cfg['device'])\noptimizer = optim.Adam(generator.parameters(), lr=cfg['lr'])\ncriterion = nn.CrossEntropyLoss()\nscheduler = ReduceLROnPlateau(optimizer, mode='min', patience=cfg['patience'], factor=0.5, verbose=True)\nes = Eearly_Stopping(mode='min', patience=5, apply=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def evaluate(val_loader, model):\n    pred_text = []\n    gt_text = []\n    bar = tqdm(val_loader, desc='eval')\n    model.eval()\n    with torch.no_grad():\n        for data in bar:\n            imgs_batch, source_padded, lens, target_batch = load_data(data)\n            pred_batch, logits = model.greedy_decode(imgs_batch, lens.max())\n            outs_packed = pack_padded_sequence(logits, lens)\n            target_packed = pack_sequence(target_batch)\n            target_packed = target_packed.to(cfg['device'])\n           \n            loss = criterion(outs_packed.data, target_packed.data)\n            \n            # compute acc\n            pred_packed = torch.argmax(torch.softmax(outs_packed.data, dim=-1), dim=-1)\n            pred_packed = pred_packed.cpu().detach().numpy()\n            target_packed = target_packed.data.cpu().numpy()\n            acc = get_accuracy(pred_packed, target_packed)\n            \n            # compute metric\n            pred_batch = pred_batch.permute(1, 0)  # (seq, batch) -> (batch, seq)\n            preds = [lang.idx_to_seq(item) for item in pred_batch.cpu().numpy()]\n            gts = [lang.idx_to_seq(item) for item in target_batch]\n            score = get_score(preds, gts)\n            \n            bar.set_description(f'Eval loss: {loss.item():.3f} score: {score:.3f} acc: {acc:.3f}')\n            \n            pred_text += preds\n            gt_text += gts\n            \n    score = get_score(pred_text, gt_text)\n    print('val score:', score)    \n    return score\n\ndef train(train_iter, epoch):\n    pred_texts = []\n    gt_texts = []\n    generator.train()\n    bar = tqdm(train_iter, desc='training')\n    for data in bar:\n        imgs_batch, source_padded, lens, target_batch  = load_data(data)\n        outs_packed = generator(imgs_batch, source_padded, lens)\n        target_packed, batch_sizes, _, _ = pack_sequence(target_batch)\n        target_packed = target_packed.to(cfg['device'])\n        loss = criterion(outs_packed, target_packed)\n        \n        ### Backward pass\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n        \n        ## compute acc\n        pred_packed = torch.argmax(torch.softmax(outs_packed, dim=-1), dim=-1)\n        pred_packed = pred_packed.cpu().detach().numpy()\n        target_packed = target_packed.cpu().numpy()\n        acc = get_accuracy(pred_packed, target_packed)\n        \n        ## compute distance\n        pred_padded = pad_packed(pred_packed, batch_sizes)[0].transpose(1, 0) # (seq, N) -> (N, seq)\n        pred_text = [lang.idx_to_seq(item) for item in pred_padded]\n        gt_text = [lang.idx_to_seq(item) for item in target_batch]\n        score = get_score(pred_text, gt_text)\n        \n        ## log\n        bar.set_description(f'Train epoch: {epoch+1} loss: {loss.item():.3f} score: {score:.3f} acc: {acc:.3f}')\n        pred_texts += pred_text\n        gt_texts += gt_text\n        \n    score = get_score(pred_texts, gt_texts)\n    print('train score:', score)\n    return generator","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tqdm import tqdm\n\nfor epoch in range(cfg['n_epochs']):\n    ## train\n    generator = train(train_iter, epoch)\n    ## evaluate\n    score = evaluate(valid_iter, generator)\n    scheduler.step(score)\n    if es.step(generator, score):\n        print('early stop!')\n        break;","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = torch.load('model.pth')\ncheckpoint = {}\ncheckpoint['net'] = model.state_dict()\ncheckpoint['train_cfg'] = cfg\ntorch.save(checkpoint, f\"{cfg['version']}/checkpont.pth\")  # save model","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}