{"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":"## Notebooks referenced (kudos to the creators!)\n\n##### <https://www.kaggle.com/ihelon/molecular-translation-exploratory-data-analysis>\n##### <https://www.kaggle.com/yasufuminakama/inchi-resnet-lstm-with-attention-starter>","metadata":{}},{"cell_type":"markdown","source":"## Import packages","metadata":{}},{"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport os\nimport cv2\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nfrom tqdm import tqdm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.optim import Adam, SGD\nimport torchvision.models as models\nfrom torch.nn.parameter import Parameter\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.nn.utils.rnn import pad_sequence, pack_padded_sequence\nimport Levenshtein\nfrom albumentations import (\n    Compose, OneOf, Normalize, Resize, RandomResizedCrop, RandomCrop, HorizontalFlip, VerticalFlip, \n    RandomBrightness, RandomContrast, RandomBrightnessContrast, Rotate, ShiftScaleRotate, Cutout, \n    IAAAdditiveGaussianNoise, Transpose, Blur\n    )\nfrom albumentations.pytorch import ToTensorV2\nfrom albumentations import ImageOnlyTransform\nimport re\n!pip install timm\nimport timm\nimport warnings \nwarnings.filterwarnings('ignore')\nimport time\nimport math\nimport pickle\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install GPUtil\nimport torch\nfrom GPUtil import showUtilization as gpu_usage\nfrom numba import cuda\n\ndef free_gpu_cache():\n    print(\"Initial GPU Usage\")\n    gpu_usage()                             \n\n    torch.cuda.empty_cache()\n\n    cuda.select_device(0)\n    cuda.close()\n    cuda.select_device(0)\n\n    print(\"GPU Usage after emptying the cache\")\n    gpu_usage()\n\nfree_gpu_cache() ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load files","metadata":{}},{"cell_type":"code","source":"df_train_labels = pd.read_csv('../input/bms-molecular-translation/train_labels.csv')\ndf_sample_submission = pd.read_csv('../input/bms-molecular-translation/sample_submission.csv')\ndf_extra_InChIs = pd.read_csv('../input/bms-molecular-translation/extra_approved_InChIs.csv')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_file_path(image_id, train_flag = True):\n    if train_flag:\n        return \"../input/bms-molecular-translation/train/{}/{}/{}/{}.png\".format(\n            image_id[0], image_id[1], image_id[2], image_id \n        )\n    else:\n        return \"../input/bms-molecular-translation/test/{}/{}/{}/{}.png\".format(\n            image_id[0], image_id[1], image_id[2], image_id \n        )        ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train_labels = df_train_labels.reset_index(drop = True)\ndf_sample_submission = df_sample_submission.reset_index(drop = True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train_labels['filename'] = df_train_labels['image_id'].apply(lambda x: get_file_path(x, train_flag = True))\ndf_sample_submission['filename'] = df_sample_submission['image_id'].apply(lambda x: get_file_path(x, train_flag = False))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(df_train_labels.shape)\ndf_train_labels.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(df_sample_submission.shape)\ndf_sample_submission.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def convert_image_id_2_path(image_id: str) -> str:\n    return \"../input/bms-molecular-translation/train/{}/{}/{}/{}.png\".format(\n        image_id[0], image_id[1], image_id[2], image_id \n    )","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"convert_image_id_2_path('000011a64c74')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def visualize_train_images(image_ids, labels):\n    plt.figure(figsize=(16, 12))\n    \n    for ind, (image_id, label) in enumerate(zip(image_ids, labels)):\n        plt.subplot(3, 3, ind + 1)\n        image = cv2.imread(convert_image_id_2_path(image_id))\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n        plt.imshow(image)\n#         print(f\"{ind}: {label}\")\n        plt.title(f\"{label[:30]}\", fontsize=10)\n        plt.axis(\"off\")\n    \n    plt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample_row = df_train_labels.sample(5)\nvisualize_train_images(\n    sample_row['image_id'], sample_row[\"InChI\"]\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(df_extra_InChIs.shape)\ndf_extra_InChIs.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## EDA & Data processing","metadata":{}},{"cell_type":"code","source":"img1 = Image.open(\"../input/bms-molecular-translation/train/0/0/0/000011a64c74.png\")\nimg2 = Image.open(\"../input/bms-molecular-translation/train/0/0/0/000019cc0cd2.png\")\nprint('Shapes:',np.asarray(img1).shape, np.asarray(img2).shape)\nprint('Palettes:',img1.palette, img2.palette)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"h_shape = []\nw_shape = []\nsample_train_ids = df_train_labels['image_id'].sample(10000)\nfor image_id in tqdm(sample_train_ids):\n    image = cv2.imread(convert_image_id_2_path(image_id))\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n    h_shape.append(image.shape[0])\n    w_shape.append(image.shape[1])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(16, 5))\nplt.subplot(1, 3, 1)\nplt.hist(np.array(h_shape) * np.array(w_shape), bins=50)\nplt.xticks(rotation=45)\nplt.title(\"Area Image Distribution\", fontsize=14)\nplt.subplot(1, 3, 2)\nplt.hist(h_shape, bins=50)\nplt.title(\"Height Image Distribution\", fontsize=14)\nplt.subplot(1, 3, 3)\nplt.hist(w_shape, bins=50)\nplt.title(\"Width Image Distribution\", fontsize=14)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def split_form(form):\n    string = ''\n    for i in re.findall(r\"[A-Z][^A-Z]*\", form):\n        elem = re.match(r\"\\D+\", i).group()\n        num = i.replace(elem, \"\")\n        if num == \"\":\n            string += f\"{elem} \"\n        else:\n            string += f\"{elem} {str(num)} \"\n    return string.rstrip(' ')\n\ndef split_form2(form):\n    string = ''\n    for i in re.findall(r\"[a-z][^a-z]*\", form):\n        elem = i[0]\n        num = i.replace(elem, \"\").replace('/', \"\")\n        num_string = ''\n        for j in re.findall(r\"[0-9]+[^0-9]*\", num):\n            num_list = list(re.findall(r'\\d+', j))\n            assert len(num_list) == 1, f\"len(num_list) != 1\"\n            _num = num_list[0]\n            if j == _num:\n                num_string += f\"{_num} \"\n            else:\n                extra = j.replace(_num, \"\")\n                num_string += f\"{_num} {' '.join(list(extra))} \"\n        string += f\"/{elem} {num_string}\"\n    return string.rstrip(' ')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train_labels['InChI'].iloc[2], df_train_labels['InChI'].iloc[2].split('/')[1], split_form(df_train_labels['InChI'].iloc[2].split('/')[1]), split_form2('/'.join(df_train_labels['InChI'].iloc[2].split('/')[2:]))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Tokenizer(object):\n    \n    def __init__(self):\n        self.stoi = {}\n        self.itos = {}\n        self.count = 0\n\n    def __len__(self):\n        return len(self.stoi)\n    \n    def fit_on_texts(self, texts):\n        vocab = set()\n        for text in texts:\n            vocab.update(text.split(' '))\n        vocab = sorted(vocab)\n        vocab.append('<sos>')\n        vocab.append('<eos>')\n        vocab.append('<pad>')\n        for i, s in enumerate(vocab):\n            if s not in list(self.stoi.keys()):\n                self.stoi[s] = self.count\n                self.count = self.count + 1\n        self.itos = {item[1]: item[0] for item in self.stoi.items()}\n        \n    def text_to_sequence(self, text):\n        sequence = []\n        sequence.append(self.stoi['<sos>'])\n        for s in text.split(' '):\n            sequence.append(self.stoi[s])\n        sequence.append(self.stoi['<eos>'])\n        return sequence\n    \n    def texts_to_sequences(self, texts):\n        sequences = []\n        for text in texts:\n            sequence = self.text_to_sequence(text)\n            sequences.append(sequence)\n        return sequences\n\n    def sequence_to_text(self, sequence):\n        return ''.join(list(map(lambda i: self.itos[i], sequence)))\n    \n    def sequences_to_texts(self, sequences):\n        texts = []\n        for sequence in sequences:\n            text = self.sequence_to_text(sequence)\n            texts.append(text)\n        return texts\n    \n    def predict_caption(self, sequence):\n        caption = ''\n        for i in sequence:\n            if i == self.stoi['<eos>'] or i == self.stoi['<pad>']:\n                break\n            caption += self.itos[i]\n        return caption\n    \n    def predict_captions(self, sequences):\n        captions = []\n        for sequence in sequences:\n            caption = self.predict_caption(sequence)\n            captions.append(caption)\n        return captions","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train_labels['InChI_1'] = df_train_labels['InChI'].apply(lambda x: x.split('/')[1])\ndf_train_labels['InChI_text'] = df_train_labels['InChI_1'].apply(split_form) + ' ' + \\\n                        df_train_labels['InChI'].apply(lambda x: '/'.join(x.split('/')[2:])).apply(split_form2).values\n# ====================================================\n# create tokenizer\n# ====================================================\n# tokenizer = Tokenizer()\n# tokenizer.fit_on_texts(df_train_labels['InChI_text'].values)\n# torch.save(tokenizer, 'tokenizer2.pth')\n\ntokenizer = torch.load('../input/bms-molecular-translation-tokenizer/tokenizer2.pth')\nprint('Saved tokenizer')\n# ====================================================\n# preprocess df_train_labels.csv\n# ====================================================\nlengths = []\ntk0 = tqdm(df_train_labels['InChI_text'].values, total=len(df_train_labels))\nfor text in tk0:\n    seq = tokenizer.text_to_sequence(text)\n    length = len(seq) - 2\n    lengths.append(length)\ndf_train_labels['InChI_length'] = lengths","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train_labels.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.hist(df_train_labels['InChI_length'], bins=50)\nplt.title(\"InChI length Distribution (Train data)\", fontsize=14)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train_labels['InChI_length'].max()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## All characters in InChI\nall_characters_in_InChI = []\nfor i in tqdm(range(0, len(df_train_labels))):\n    characters_in_InChI = list(set(df_train_labels['InChI'].iloc[i]))\n    all_characters_in_InChI = all_characters_in_InChI + list(set(characters_in_InChI).difference(set(all_characters_in_InChI)))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(len(all_characters_in_InChI), 'characters in InChI')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_score(y_true, y_pred):\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","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"get_score('apple','paple')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def init_logger(log_file='train.log'):\n    from logging import getLogger, INFO, FileHandler,  Formatter,  StreamHandler\n    logger = getLogger(__name__)\n    logger.setLevel(INFO)\n    handler1 = StreamHandler()\n    handler1.setFormatter(Formatter(\"%(message)s\"))\n    handler2 = FileHandler(filename=log_file)\n    handler2.setFormatter(Formatter(\"%(message)s\"))\n    logger.addHandler(handler1)\n    logger.addHandler(handler2)\n    return logger\n\nLOGGER = init_logger()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TrainDataset(Dataset):\n    def __init__(self, df, tokenizer, transform=None):\n        super().__init__()\n        self.df = df\n        self.tokenizer = tokenizer\n        self.file_paths = df['filename'].values\n        self.labels = df['InChI_text'].values\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        file_path = self.file_paths[idx]\n        image = cv2.imread(file_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n        label = self.labels[idx]\n        label = self.tokenizer.text_to_sequence(label)\n        label_length = len(label)\n        label_length = torch.LongTensor([label_length])\n        return image, torch.LongTensor(label), label_length","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TestDataset(Dataset):\n    def __init__(self, df, transform=None):\n        super().__init__()\n        self.df = df\n        self.file_paths = df['filename'].values\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        file_path = self.file_paths[idx]\n        image = cv2.imread(file_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n        return image","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def bms_collate(batch):\n    imgs, labels, label_lengths = [], [], []\n    for data_point in batch:\n        imgs.append(data_point[0])\n        labels.append(data_point[1])\n        label_lengths.append(data_point[2])\n    labels = pad_sequence(labels, batch_first=True, padding_value=tokenizer.stoi[\"<pad>\"])\n    return torch.stack(imgs), labels, torch.stack(label_lengths).reshape(-1, 1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#https://pytorch.org/hub/pytorch_vision_resnet/\n\ndef get_transforms(*, data):\n    \n    if data == 'train':\n        return Compose([\n            Resize(500, 500),\n            Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n            ),\n            ToTensorV2(),\n        ])\n    \n    elif data == 'valid':\n        return Compose([\n            Resize(500, 500),\n            Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n            ),\n            ToTensorV2(),\n        ])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train_labels.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = TrainDataset(df_train_labels, tokenizer, transform=get_transforms(data='train'))\n\nfor i in range(1):\n    image, label, label_length = train_dataset[i]\n    text = tokenizer.sequence_to_text(label.numpy())\n    plt.imshow(image.transpose(0, 1).transpose(1, 2))\n    plt.title(f'label: {label}  text: {text}  label_length: {label_length}')\n    plt.show() ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"visualize_train_images(df_train_labels['image_id'].iloc[0:1], df_train_labels[\"InChI\"].iloc[0:1])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"list_InChI_versions = df_train_labels['InChI'].apply(lambda x: x.split('/')[0])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"set(list_InChI_versions)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Build Model","metadata":{}},{"cell_type":"code","source":"class Encoder(nn.Module):\n    def __init__(self, model_name='resnet18', pretrained=False):\n        super().__init__()\n        self.cnn = timm.create_model(model_name, pretrained=pretrained)\n        self.n_features = self.cnn.fc.in_features\n        self.cnn.global_pool = nn.Identity()\n        self.cnn.fc = nn.Identity()\n\n    def forward(self, x):\n        bs = x.size(0)\n        features = self.cnn(x)\n        features = features.permute(0, 2, 3, 1)\n        return features","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Attention(nn.Module):\n    \"\"\"\n    Attention network for calculate attention value\n    \"\"\"\n    def __init__(self, encoder_dim, decoder_dim, attention_dim):\n        \"\"\"\n        :param encoder_dim: input size of encoder network\n        :param decoder_dim: input size of decoder network\n        :param attention_dim: input size of attention network\n        \"\"\"\n        super(Attention, self).__init__()\n        self.encoder_att = nn.Linear(encoder_dim, attention_dim)  # linear layer to transform encoded image\n        self.decoder_att = nn.Linear(decoder_dim, attention_dim)  # linear layer to transform decoder's output\n        self.full_att = nn.Linear(attention_dim, 1)  # linear layer to calculate values to be softmax-ed\n        self.relu = nn.ReLU()\n        self.softmax = nn.Softmax(dim=1)  # softmax layer to calculate weights\n\n    def forward(self, encoder_out, decoder_hidden):\n        att1 = self.encoder_att(encoder_out)  # (batch_size, num_pixels, attention_dim)\n        att2 = self.decoder_att(decoder_hidden)  # (batch_size, attention_dim)\n        att = self.full_att(self.relu(att1 + att2.unsqueeze(1))).squeeze(2)  # (batch_size, num_pixels)\n        alpha = self.softmax(att)  # (batch_size, num_pixels)\n        attention_weighted_encoding = (encoder_out * alpha.unsqueeze(2)).sum(dim=1)  # (batch_size, encoder_dim)\n        return attention_weighted_encoding, alpha","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DecoderWithAttention(nn.Module):\n    \"\"\"\n    Decoder network with attention network used for training\n    \"\"\"\n\n    def __init__(self, attention_dim, embed_dim, decoder_dim, vocab_size, device, encoder_dim=512, dropout=0.5):\n        \"\"\"\n        :param attention_dim: input size of attention network\n        :param embed_dim: input size of embedding network\n        :param decoder_dim: input size of decoder network\n        :param vocab_size: total number of characters used in training\n        :param encoder_dim: input size of encoder network\n        :param dropout: dropout rate\n        \"\"\"\n        super(DecoderWithAttention, self).__init__()\n        self.encoder_dim = encoder_dim\n        self.attention_dim = attention_dim\n        self.embed_dim = embed_dim\n        self.decoder_dim = decoder_dim\n        self.vocab_size = vocab_size\n        self.dropout = dropout\n        self.device = device\n        self.attention = Attention(encoder_dim, decoder_dim, attention_dim)  # attention network\n        self.embedding = nn.Embedding(vocab_size, embed_dim)  # embedding layer\n        self.dropout = nn.Dropout(p=self.dropout)\n        self.decode_step = nn.LSTMCell(embed_dim + encoder_dim, decoder_dim, bias=True)  # decoding LSTMCell\n        self.init_h = nn.Linear(encoder_dim, decoder_dim)  # linear layer to find initial hidden state of LSTMCell\n        self.init_c = nn.Linear(encoder_dim, decoder_dim)  # linear layer to find initial cell state of LSTMCell\n        self.f_beta = nn.Linear(decoder_dim, encoder_dim)  # linear layer to create a sigmoid-activated gate\n        self.sigmoid = nn.Sigmoid()\n        self.fc = nn.Linear(decoder_dim, vocab_size)  # linear layer to find scores over vocabulary\n        self.init_weights()  # initialize some layers with the uniform distribution\n\n    def init_weights(self):\n        self.embedding.weight.data.uniform_(-0.1, 0.1)\n        self.fc.bias.data.fill_(0)\n        self.fc.weight.data.uniform_(-0.1, 0.1)\n\n    def load_pretrained_embeddings(self, embeddings):\n        self.embedding.weight = nn.Parameter(embeddings)\n\n    def fine_tune_embeddings(self, fine_tune=True):\n        for p in self.embedding.parameters():\n            p.requires_grad = fine_tune\n\n    def init_hidden_state(self, encoder_out):\n        mean_encoder_out = encoder_out.mean(dim=1)\n        h = self.init_h(mean_encoder_out)  # (batch_size, decoder_dim)\n        c = self.init_c(mean_encoder_out)\n        return h, c\n\n    def forward(self, encoder_out, encoded_captions, caption_lengths):\n        \"\"\"\n        :param encoder_out: output of encoder network\n        :param encoded_captions: transformed sequence from character to integer\n        :param caption_lengths: length of transformed sequence\n        \"\"\"\n        batch_size = encoder_out.size(0)\n        encoder_dim = encoder_out.size(-1)\n        vocab_size = self.vocab_size\n        encoder_out = encoder_out.view(batch_size, -1, encoder_dim)  # (batch_size, num_pixels, encoder_dim)\n        num_pixels = encoder_out.size(1)\n        caption_lengths, sort_ind = caption_lengths.squeeze(1).sort(dim=0, descending=True)\n        encoder_out = encoder_out[sort_ind]\n        encoded_captions = encoded_captions[sort_ind]\n        # embedding transformed sequence for vector\n        embeddings = self.embedding(encoded_captions)  # (batch_size, max_caption_length, embed_dim)\n        # initialize hidden state and cell state of LSTM cell\n        h, c = self.init_hidden_state(encoder_out)  # (batch_size, decoder_dim)\n        # set decode length by caption length - 1 because of omitting start token\n        decode_lengths = (caption_lengths - 1).tolist()\n        predictions = torch.zeros(batch_size, max(decode_lengths), vocab_size).to(self.device)\n        alphas = torch.zeros(batch_size, max(decode_lengths), num_pixels).to(self.device)\n        # predict sequence\n        for t in range(max(decode_lengths)):\n            batch_size_t = sum([l > t for l in decode_lengths])\n            attention_weighted_encoding, alpha = self.attention(encoder_out[:batch_size_t], h[:batch_size_t])\n            gate = self.sigmoid(self.f_beta(h[:batch_size_t]))  # gating scalar, (batch_size_t, encoder_dim)\n            attention_weighted_encoding = gate * attention_weighted_encoding\n            h, c = self.decode_step(\n                torch.cat([embeddings[:batch_size_t, t, :], attention_weighted_encoding], dim=1),\n                (h[:batch_size_t], c[:batch_size_t]))  # (batch_size_t, decoder_dim)\n            preds = self.fc(self.dropout(h))  # (batch_size_t, vocab_size)\n            predictions[:batch_size_t, t, :] = preds\n            alphas[:batch_size_t, t, :] = alpha\n        return predictions, encoded_captions, decode_lengths, alphas, sort_ind\n    \n    def predict(self, encoder_out, decode_lengths, tokenizer):\n        batch_size = encoder_out.size(0)\n        encoder_dim = encoder_out.size(-1)\n        vocab_size = self.vocab_size\n        encoder_out = encoder_out.view(batch_size, -1, encoder_dim)  # (batch_size, num_pixels, encoder_dim)\n        num_pixels = encoder_out.size(1)\n        # embed start tocken for LSTM input\n        start_tockens = torch.ones(batch_size, dtype=torch.long).to(self.device) * tokenizer.stoi[\"<sos>\"]\n        embeddings = self.embedding(start_tockens)\n        # initialize hidden state and cell state of LSTM cell\n        h, c = self.init_hidden_state(encoder_out)  # (batch_size, decoder_dim)\n        predictions = torch.zeros(batch_size, decode_lengths, vocab_size).to(self.device)\n        # predict sequence\n        for t in range(decode_lengths):\n            attention_weighted_encoding, alpha = self.attention(encoder_out, h)\n            gate = self.sigmoid(self.f_beta(h))  # gating scalar, (batch_size_t, encoder_dim)\n            attention_weighted_encoding = gate * attention_weighted_encoding\n            h, c = self.decode_step(\n                torch.cat([embeddings, attention_weighted_encoding], dim=1),\n                (h, c))  # (batch_size_t, decoder_dim)\n            preds = self.fc(self.dropout(h))  # (batch_size_t, vocab_size)\n            predictions[:, t, :] = preds\n            if np.argmax(preds.detach().cpu().numpy()) == tokenizer.stoi[\"<eos>\"]:\n                break\n            embeddings = self.embedding(torch.argmax(preds, -1))\n        return predictions","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\ndef asMinutes(s):\n    m = math.floor(s / 60)\n    s -= m * 60\n    return '%dm %ds' % (m, s)\n\n\ndef timeSince(since, percent):\n    now = time.time()\n    s = now - since\n    es = s / (percent)\n    rs = es - s\n    return '%s (remain %s)' % (asMinutes(s), asMinutes(rs))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_fn(train_loader, encoder, decoder, criterion, \n             encoder_optimizer, decoder_optimizer, epoch,\n             device, encoder_scheduler=None, decoder_scheduler=None):\n    batch_time = AverageMeter()\n    data_time = AverageMeter()\n    losses = AverageMeter()\n    # switch to train mode\n    encoder.train()\n    decoder.train()\n    start = end = time.time()\n    global_step = 0\n    for step, (images, labels, label_lengths) in enumerate(train_loader):\n        # measure data loading time\n        data_time.update(time.time() - end)\n        images = images.to(device)\n        labels = labels.to(device)\n        label_lengths = label_lengths.to(device)\n        batch_size = images.size(0)\n        features = encoder(images)\n        predictions, caps_sorted, decode_lengths, alphas, sort_ind = decoder(features, labels, label_lengths)\n        targets = caps_sorted[:, 1:]\n        predictions = pack_padded_sequence(predictions, decode_lengths, batch_first=True).data\n        targets = pack_padded_sequence(targets, decode_lengths, batch_first=True).data\n        loss = criterion(predictions, targets)\n        # record loss\n        losses.update(loss.item(), batch_size)\n        loss.backward()\n        encoder_grad_norm = torch.nn.utils.clip_grad_norm_(encoder.parameters(), 5)\n        decoder_grad_norm = torch.nn.utils.clip_grad_norm_(decoder.parameters(), 5)\n        encoder_optimizer.step()\n        decoder_optimizer.step()\n        encoder_optimizer.zero_grad()\n        decoder_optimizer.zero_grad()\n        global_step += 1\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n        if step % 1000 == 0 or step == (len(train_loader)-1):\n            print('Epoch: [{0}][{1}/{2}] '\n                  'Data {data_time.val:.3f} ({data_time.avg:.3f}) '\n                  'Elapsed {remain:s} '\n                  'Loss: {loss.val:.4f}({loss.avg:.4f}) '\n                  'Encoder Grad: {encoder_grad_norm:.4f}  '\n                  'Decoder Grad: {decoder_grad_norm:.4f}  '\n                  .format(\n                   epoch+1, step, len(train_loader), batch_time=batch_time,\n                   data_time=data_time, loss=losses,\n                   remain=timeSince(start, float(step+1)/len(train_loader)),\n                   encoder_grad_norm=encoder_grad_norm,\n                   decoder_grad_norm=decoder_grad_norm\n                   ))\n    return losses.avg","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def valid_fn(valid_loader, encoder, decoder, tokenizer, criterion, device):\n    batch_time = AverageMeter()\n    data_time = AverageMeter()\n    # switch to evaluation mode\n    encoder.eval()\n    decoder.eval()\n    text_preds = []\n    start = end = time.time()\n    for step, (images) in enumerate(valid_loader):\n        # measure data loading time\n        data_time.update(time.time() - end)\n        images = images.to(device)\n        batch_size = images.size(0)\n        with torch.no_grad():\n            features = encoder(images)\n            predictions = decoder.predict(features, 275, tokenizer)\n        predicted_sequence = torch.argmax(predictions.detach().cpu(), -1).numpy()\n        _text_preds = tokenizer.predict_captions(predicted_sequence)\n        text_preds.append(_text_preds)\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n        if step % 1000 == 0 or step == (len(valid_loader)-1):\n            print('EVAL: [{0}/{1}] '\n                  'Data {data_time.val:.3f} ({data_time.avg:.3f}) '\n                  'Elapsed {remain:s} '\n                  .format(\n                   step, len(valid_loader), batch_time=batch_time,\n                   data_time=data_time,\n                   remain=timeSince(start, float(step+1)/len(valid_loader)),\n                   ))\n    text_preds = np.concatenate(text_preds)\n    return text_preds","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = TrainDataset(df_train_labels.sample(100000), tokenizer, transform=get_transforms(data='train'))\nvalid_dataset = TestDataset(df_sample_submission, transform=get_transforms(data='valid'))\n\ntrain_loader = DataLoader(train_dataset, \n                          batch_size=64, \n                          shuffle=True, \n                          num_workers=4, \n                          pin_memory=True,\n                          drop_last=True, \n                          collate_fn=bms_collate)\nvalid_loader = DataLoader(valid_dataset, \n                          batch_size=64, \n                          shuffle=False, \n                          num_workers=4,\n                          pin_memory=True, \n                          drop_last=False)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"encoder = Encoder(model_name='resnet18', pretrained=True)\nencoder.to(device)\nencoder_optimizer = Adam(encoder.parameters(), lr=1e-4, weight_decay=1e-6, amsgrad=False)\n\ndecoder = DecoderWithAttention(attention_dim=256,\n                               embed_dim=256,\n                               decoder_dim=512,\n                               vocab_size=len(tokenizer),\n                               dropout=0.5,\n                               device=device)\ndecoder.to(device)\ndecoder_optimizer = Adam(decoder.parameters(), lr=1e-4, weight_decay=1e-6, amsgrad=False)\ncriterion = nn.CrossEntropyLoss(ignore_index=tokenizer.stoi[\"<pad>\"])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for epoch in range(1):\n    start_time = time.time()\n\n    # train\n    avg_loss = train_fn(train_loader, encoder, decoder, criterion, \n                        encoder_optimizer, decoder_optimizer, epoch, device)\n\n    # eval\n    text_preds = valid_fn(valid_loader, encoder, decoder, tokenizer, criterion, device)\n    text_preds = [f\"InChI=1S/{text}\" for text in text_preds]\n#     LOGGER.info(f\"labels: {valid_labels[:5]}\")\n#     LOGGER.info(f\"preds: {text_preds[:5]}\")\n\n    # scoring\n#     score = get_score(valid_labels, text_preds)\n    elapsed = time.time() - start_time\n\n#     LOGGER.info(f'Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  time: {elapsed:.0f}s')\n#     LOGGER.info(f'Epoch {epoch+1} - Score: {score:.4f}')\n    \n    with open('InChI_output.pkl', 'wb') as f:\n        pickle.dump(text_preds, f)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}