{"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":"# **Pytorch Lightning BMS Molecular Translation Train**\n\n# **References**\n* Version 1:PL Encoder(Transformer:tnt) + Decoder(transformer) code https://www.kaggle.com/c/bms-molecular-translation/discussion/231190\n  \n* Version 2:PL Encoder(Resnet34) + Decoder(LSTM with attention) code https://www.kaggle.com/yasufuminakama/inchi-resnet-lstm-with-attention-starter\n\n* Inference: https://www.kaggle.com/yingpengchen/pl-bms-molecular-translation\n\nIf these notebook are helpful, feel free to upvote.","metadata":{}},{"cell_type":"markdown","source":"# **Import 3rdparty**","metadata":{}},{"cell_type":"code","source":"import os\nimport math\nimport time\nimport random\nimport shutil\nfrom pathlib import Path\nfrom contextlib import contextmanager\nfrom collections import defaultdict, Counter, OrderedDict\nimport argparse\nimport gc\n\nimport scipy as sp\nimport pandas as pd\nimport numpy as np\n\nimport Levenshtein\nfrom sklearn import preprocessing\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import StratifiedKFold, GroupKFold, KFold\n\nfrom tqdm.auto import tqdm\nfrom functools import partial\n\nimport cv2\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.optim import Adam, SGD\nfrom torch.nn.parameter import Parameter\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.nn.utils.rnn import pad_sequence, pack_padded_sequence\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau\nimport torch.distributed as dist\n\nimport pytorch_lightning as pl\nimport warnings\n\nimport sys\n# sys.path.append('../input/pytorchimagemodels210218/pytorch-image-models-master')\nsys.path.append('../input/pytorchimagemodels210417/pytorch-image-models-master')\nimport timm\nfrom timm.models.tnt import tnt_s_patch16_224 \n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Config**","metadata":{}},{"cell_type":"code","source":"class Config:\n    data_dir = '../input/bms-molecular-translation/'\n    train_dir = data_dir + 'train/'\n    test_dir = data_dir + 'test/'\n    train_csv_path = data_dir + 'train_labels.csv'\n    test_csv_path = data_dir + 'sample_submission.csv'\n\n    train_folds_path = '../input/bms-train-folds/train_folds.csv'\n    tokenizer_path = '../input/bmstokenizer/bms_tokenizer.pth'\n\n    \n    device = 'cuda'             ### set gpu or cpu mode\n    debug = True              ### debug flag for checking your modify code\n    \n    gpus = 1                 ### gpu numbers\n    precision = 16             ### training precision 8, 16,32, etc\n    batch_size = 32       ### total batch size\n    epochs = 10               ### total training epochs\n    encoder_lr = 1e-4      ### encode learning rate default 1e-4\n    decoder_lr = 1e-4      ### decode learning rate default 4e-4\n    min_lr = 1e-6              ### min learning rate\n    weight_decay = 1e-6\n    gradient_accumulation_steps = 1\n#     max_grad_norm =1000         ### 5\n    num_workers = 16            ### number workers\n    print_freq = 100            ### print log frequency\n    \n#     decoder_mode = 'lstm'         ## lstm, transformer\n#     encoder_name = 'resnet34'\n#     image_size = 224            ### input size in training\n#     attention_dim=512\n#     embed_dim=512\n#     encoder_dim = 512\n#     decoder_dim = 512\n#     max_length = 275 # 275\n    \n    decoder_mode = 'transformer'\n    encoder_name = 'tnt_s_patch16_224'\n    image_size = 224            ### input size in training\n    image_dim = 384\n    text_dim = 384\n    decoder_dim = 384\n    ff_dim = 512\n    num_layer = 3\n    num_head = 8\n    max_length = 300 # 275\n    \n    dropout = 0.5\n\n    seed = 42\n    n_fold = 5\n    trn_fold = [0] # [0,1,2,3,4]\n    fold=0\n    train=True\n    \n    scheduler = 'CosineAnnealingWarmRestarts' # ['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts']\n\n#     # ReduceLROnPlateau\n#     factor=0.2 # ReduceLROnPlateau\n#     patience=4 # ReduceLROnPlateau\n#     eps=1e-6 # ReduceLROnPlateau\n    \n    ## CosineAnnealingLR\n    # T_max=4 # CosineAnnealingLR\n\n    ## CosineAnnealingWarmRestarts\n    T_0 = 10\n    T_mult = 1","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Load Tokenizer**","metadata":{}},{"cell_type":"code","source":"class Tokenizer(object):\n    \n    def __init__(self):\n        self.stoi = {}\n        self.itos = {}\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            self.stoi[s] = i\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 = 'InChI=1S/'\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\n    \n    def one_predict_to_inchi(self, predict):\n        inchi = 'InChI=1S/'\n        for p in predict:\n            if p == self.stoi['<eos>'] or p == self.stoi['<pad>']:\n                break\n            inchi += self.itos[p]\n        return inchi\n    \n    def predict_to_inchi(self, predict):\n        inchi = [\n            self.one_predict_to_inchi(p)\n            for p in predict\n        ]\n        return inchi","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"CFG = Config\ntokenizer = torch.load(CFG.tokenizer_path)\nprint(f\"tokenizer.stoi: {tokenizer.stoi}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Load Train Folds**","metadata":{}},{"cell_type":"code","source":"folds = pd.read_csv(CFG.train_folds_path)\nfolds.head()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Dataset**","metadata":{}},{"cell_type":"code","source":"# ====================================================\n# Dataset\n# ====================================================\nclass TrainDataset(Dataset):\n    def __init__(self, CFG, df, tokenizer, transform=None):\n        super().__init__()\n        self.CFG = CFG\n        self.df = df\n        self.tokenizer = tokenizer\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        image_id = self.df.iloc[idx]['image_id']\n        image_path = os.path.join(self.CFG.train_dir,image_id[0],image_id[1],image_id[2], f'{image_id}.png')\n        image = cv2.imread(image_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\n    \n\nclass TestDataset(Dataset):\n    def __init__(self, CFG, df, transform=None):\n        super().__init__()\n        self.CFG = CFG\n        self.df = df\n        self.labels = df['InChI'].values\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        image_id = self.df.iloc[idx]['image_id']\n        image_path = os.path.join(self.CFG.train_dir,image_id[0],image_id[1],image_id[2], f'{image_id}.png')\n        image = cv2.imread(image_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] ## text labels\n        return image, label\n\n\ndef get_train_transforms(CFG):\n    return A.Compose([\n            A.Resize(height=CFG.image_size, width=CFG.image_size,p=1.0),\n            # A.Transpose(p=0.5),\n            # A.HorizontalFlip(p=0.5),\n            # A.VerticalFlip(p=0.5),\n            A.RandomRotate90(p=0.5),\n#             A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),\n#             A.RandomBrightnessContrast(brightness_limit=(-0.2,0.2), contrast_limit=(-0.2, 0.2), p=0.5),\n#             A.ShiftScaleRotate(shift_limit=0.0625,scale_limit=0.20,rotate_limit=20, p=0.5),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0),\n            ToTensorV2(),\n        ],p=1.0)\n\ndef get_val_transforms(CFG):\n    return A.Compose([\n            A.Resize(CFG.image_size, CFG.image_size),\n            A.CenterCrop(height=CFG.image_size, width=CFG.image_size,p=1.0),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),\n            ToTensorV2(),\n        ],p=1.0)\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Encoder Module**","metadata":{}},{"cell_type":"code","source":"class Encoder(nn.Module):\n    def __init__(self, model_name='resnet18', pretrained=False):\n        super().__init__()\n        self.model_name = model_name\n        if model_name == 'resnet18' or model_name == 'resnet34' or model_name == 'resnet200d':\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        if model_name == 'tf_efficientnet_b0_ns' or model_name == 'tf_efficientnet_b3_ns' or model_name == 'tf_efficientnet_b4_ns':\n            self.cnn = timm.create_model(model_name, pretrained=pretrained)\n            self.n_features = self.cnn.classifier.in_features\n            self.cnn.global_pool = nn.Identity()\n            self.cnn.classifier = nn.Identity()\n        \n        if model_name == 'vit_base_patch16_224' or model_name == 'vit_base_patch16_384':\n            self.cnn = timm.create_model(model_name, pretrained=pretrained)\n            self.n_features = self.cnn.head.in_features \n#             self.cnn.norm = nn.Identity()\n            self.cnn.head = nn.Identity()\n            \n        if model_name == 'tnt_s_patch16_224':\n            self.cnn = tnt_s_patch16_224(pretrained=pretrained)\n         \n        if model_name == 'tnt_s_patch16_384':  # need modify tnt.py by yourself\n            self.cnn = tnt_s_patch16_384(pretrained=False)\n        \n    def forward(self, x):\n        if self.model_name == 'vit_base_patch16_224':\n            B = x.shape[0]\n#             x = 2 * x - 1 \n            #print(f'input x size is {x.size()}')\n        \n            x = self.cnn.patch_embed(x)\n            cls_tokens = self.cnn.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks\n            x = torch.cat((cls_tokens, x), dim=1)\n            x = x + self.cnn.pos_embed\n            x = self.cnn.pos_drop(x)\n\n            for blk in self.cnn.blocks:\n                x = blk(x)\n\n#             enc_self_attns = self.cnn.norm(x)[:, 0]\n#             encoder_out = self.cnn.norm(x)[:, 1:]\n            x = self.cnn.norm(x) ### {batch_size, 197, 768]\n            \n            return x\n        \n        elif self.model_name == 'tnt_s_patch16_224' or self.model_name == 'tnt_s_patch16_384':\n            batch_size, C, H, W = x.shape\n            x = 2 * x - 1  # ; print('input ',   x.size())\n\n            pixel_embed = self.cnn.pixel_embed(x, self.cnn.pixel_pos)\n\n            patch_embed = self.cnn.norm2_proj(self.cnn.proj(self.cnn.norm1_proj(pixel_embed.reshape(batch_size, self.cnn.num_patches, -1))))\n            patch_embed = torch.cat((self.cnn.cls_token.expand(batch_size, -1, -1), patch_embed), dim=1)\n            patch_embed = patch_embed + self.cnn.patch_pos\n            patch_embed = self.cnn.pos_drop(patch_embed)\n\n            for blk in self.cnn.blocks:\n                pixel_embed, patch_embed = blk(pixel_embed, patch_embed)\n\n            patch_embed = self.cnn.norm(patch_embed) #torch.Size([7, 197, 384])\n\n            x = patch_embed\n            return x\n        \n        else:\n            bs = x.size(0)\n            features = self.cnn(x)\n            features = features.permute(0, 2, 3, 1)\n            \n            return features","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Decoder with LSTM**","metadata":{}},{"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\n\n### LSTM decoder with attention\nclass 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, 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.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        device = encoder_out.device\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        embeddings = self.embedding(encoded_captions).to(encoder_out.dtype) # for 16bit precision added by chen\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(device)\n        alphas = torch.zeros(batch_size, max(decode_lengths), num_pixels).to(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            attention_weighted_encoding = attention_weighted_encoding.to(encoder_out.dtype) ## added by chen for 16bit precision\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        device = encoder_out.device\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(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(device)\n        \n#         eos = tokenizer.stoi['<eos>']\n#         pad = tokenizer.stoi['<pad>']\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\n    \n    def forward_step(self, prev_tokens, hidden, encoder_out, function):\n        assert len(hidden) == 2\n        h, c = hidden\n        h, c = h.squeeze(0), c.squeeze(0)\n\n        embeddings = self.embedding(prev_tokens)\n        if embeddings.dim() == 3:\n            embeddings = embeddings.squeeze(1)\n\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\n        hidden = (h.unsqueeze(0), c.unsqueeze(0))\n        predicted_softmax = function(preds, dim=1)\n        return predicted_softmax, hidden, None   ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Decoder with Transformer**","metadata":{}},{"cell_type":"code","source":"import torch\nfrom torch import nn\nimport torch.nn.functional as F\nimport numpy as np\nimport math\n# Table 1: Post-LN Transformer v.s. Pre-LN Transformer\n# ON LAYER NORMALIZATION IN THE TRANSFORMER ARCHITECTURE - ICLR 2020\n# https://openreview.net/pdf?id=B1x8anVFPr\n\n# https://towardsdatascience.com/how-to-code-the-transformer-in-pytorch-24db27c8f9ec\n# https://scale.com/blog/pytorch-improvements\n# Making Pytorch Transformer Twice as Fast on Sequence Generation.\n# https://towardsdatascience.com/how-to-code-the-transformer-in-pytorch-24db27c8f9ec\n\n# ------------------------------------------------------\n# https://kazemnejad.com/blog/transformer_architecture_positional_encoding/\n# https://stackoverflow.com/questions/46452020/sinusoidal-embedding-attention-is-all-you-need\n\nclass PositionEncode1D(nn.Module):\n    def __init__(self, dim, max_length):\n        super().__init__()\n        assert (dim % 2 == 0)\n        self.max_length = max_length\n\n        d = torch.exp(torch.arange(0., dim, 2)* (-math.log(10000.0) / dim))\n        position = torch.arange(0., max_length).unsqueeze(1)\n        pos = torch.zeros(1, max_length, dim)\n        pos[0, :, 0::2] = torch.sin(position * d)\n        pos[0, :, 1::2] = torch.cos(position * d)\n        self.register_buffer('pos', pos)\n\n    def forward(self, x):\n        batch_size, T, dim = x.shape\n        x = x + self.pos[:,:T]\n        return x\n\n#https://github.com/wzlxjtu/PositionalEncoding2D/blob/master/positionalembedding2d.py\nclass PositionEncode2D(nn.Module):\n    def __init__(self, dim, width, height):\n        super().__init__()\n        assert (dim % 4 == 0)\n        self.width  = width\n        self.height = height\n\n        dim = dim//2\n        d = torch.exp(torch.arange(0., dim, 2) * -(math.log(10000.0) / dim))\n        position_w = torch.arange(0., width ).unsqueeze(1)\n        position_h = torch.arange(0., height).unsqueeze(1)\n        pos = torch.zeros(1, dim*2, height, width)\n\n        pos[0,      0:dim:2, :, :] = torch.sin(position_w * d).transpose(0, 1).unsqueeze(1).repeat(1,1, height, 1)\n        pos[0,      1:dim:2, :, :] = torch.cos(position_w * d).transpose(0, 1).unsqueeze(1).repeat(1,1, height, 1)\n        pos[0,dim + 0:   :2, :, :] = torch.sin(position_h * d).transpose(0, 1).unsqueeze(2).repeat(1,1, 1, width)\n        pos[0,dim + 1:   :2, :, :] = torch.cos(position_h * d).transpose(0, 1).unsqueeze(2).repeat(1,1, 1, width)\n        self.register_buffer('pos', pos)\n\n    def forward(self, x):\n        batch_size,C,H,W = x.shape\n        x = x + self.pos[:,:,:H,:W]\n        return x\n    \n\n# def triangle_mask(size):\n#     mask = 1- np.triu(np.ones((1, size, size)),k=1).astype('uint8')\n#     mask = torch.autograd.Variable(torch.from_numpy(mask))\n#     return mask\n\n'''\ntriangle_mask(10)\n\nmask\narray([[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n        [1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n        [1, 1, 1, 0, 0, 0, 0, 0, 0, 0],\n        [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],\n        [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],\n        [1, 1, 1, 1, 1, 1, 0, 0, 0, 0],\n        [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],\n        [1, 1, 1, 1, 1, 1, 1, 1, 0, 0],\n        [1, 1, 1, 1, 1, 1, 1, 1, 1, 0],\n        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]], dtype=uint8)\n'''\n\n# ------------------------------------------------------\nclass FeedForward(nn.Module):\n    def __init__(self, dim, ff_dim=2048, dropout=0.1):\n        super().__init__()\n        self.linear1 = nn.Linear(dim, ff_dim)\n        self.linear2 = nn.Linear(ff_dim, dim)\n        self.dropout = nn.Dropout(dropout)\n\n    def forward(self, x):\n        x = self.dropout(F.relu(self.linear1(x)))\n        x = self.linear2(x)\n        return x\n\n#layer normalization\nclass Norm(nn.Module):\n    def __init__(self, dim, eps=1e-6):\n        super().__init__()\n        self.alpha = nn.Parameter(torch.ones(dim))\n        self.bias  = nn.Parameter(torch.zeros(dim))\n        self.eps   = eps\n    def forward(self, x):\n        #return x\n        z = (x - x.mean(dim=-1, keepdim=True)) / (x.std(dim=-1, keepdim=True) + self.eps)\n        x = self.alpha*z + self.bias\n        return x\n\nclass MultiHeadAttention(nn.Module):\n    def __init__(self, dim, num_head, dropout=0.1):\n        super().__init__()\n\n        self.dim = dim\n        self.d_k = dim // num_head\n        self.num_head = num_head\n        self.dropout = dropout\n\n        self.q = nn.Linear(dim, dim)\n        self.v = nn.Linear(dim, dim)\n        self.k = nn.Linear(dim, dim)\n        self.out = nn.Linear(dim, dim)\n\n    def attention(self, q, k, v, mask):\n        score = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)  # torch.Size([8, 4, 10, 10]) = batch_size, num_head, LqxLk\n        if mask is not None:\n            mask = mask.unsqueeze(1)\n            #print(score.min())\n            score = score.masked_fill(mask == 0, -6e4) #-65504\n#             score = score.masked_fill(mask == 0, -half('inf'))\n            # https://github.com/NVIDIA/apex/issues/93\n            # How to use fp16 training with masked operations\n\n        score = F.softmax(score, dim=-1)\n\n        if self.dropout > 0:\n            score = F.dropout(score, self.dropout, training=self.training)\n\n        value = torch.matmul(score, v)\n        return value\n\n\n    def forward(self, q, k, v, mask=None):\n        batch_size, T, dim = q.shape\n\n        # perform linear operation and split into h heads\n        k = self.k(k).reshape(batch_size, -1, self.num_head, self.d_k)\n        q = self.q(q).reshape(batch_size, -1, self.num_head, self.d_k)\n        v = self.v(v).reshape(batch_size, -1, self.num_head, self.d_k)\n\n        # transpose to get dimensions batch_size * num_head * T * d_k\n        k = k.transpose(1, 2)\n        q = q.transpose(1, 2)\n        v = v.transpose(1, 2)\n\n        # calculate attention using function we will define next\n        value = self.attention(q, k, v, mask)\n\n        # concatenate heads and put through final linear layer\n        value = value.transpose(1, 2).contiguous().reshape(batch_size, -1, self.dim)\n        value = self.out(value)\n        return value\n\n\n#---\nclass TransformerEncodeLayer(nn.Module):\n    def __init__(self, dim, ff_dim, num_head, dropout=0.1):\n        super().__init__()\n        self.norm1 = Norm(dim)\n        self.norm2 = Norm(dim)\n\n        self.attn = MultiHeadAttention(dim, num_head, dropout=0.1)\n        self.ff   = FeedForward(dim, ff_dim)\n\n        self.dropout1 = nn.Dropout(dropout)\n        self.dropout2 = nn.Dropout(dropout)\n\n    def forward(self, x, x_mask):\n        x1 = self.norm1(x)\n        x1 = self.attn(x1, x1, x1, x_mask) #self-attention\n        x   = x + self.dropout1(x1)\n\n        x2 = self.norm2(x)\n        x2 = self.ff(x2)\n        x  = x + self.dropout2(x2)\n        return x\n\nclass TransformerEncode(nn.Module):\n    def __init__(self, dim, ff_dim, num_head, num_layer):\n        super().__init__()\n        self.num_layer = num_layer\n\n        self.layer = nn.ModuleList([\n            TransformerEncodeLayer(dim, ff_dim, num_head) for i in range(num_layer)\n        ])\n        self.norm = Norm(dim)\n\n    def forward(self, x, x_mask):\n        for i in range(self.num_layer):\n            x = self.layer[i](x, x_mask)\n        x = self.norm(x)\n        return x\n\n#---\nclass TransformerDecodeLayer(nn.Module):\n    def __init__(self, dim, ff_dim, num_head, dropout=0.1):\n        super().__init__()\n        self.norm1 = Norm(dim)\n        self.norm2 = Norm(dim)\n        self.norm3 = Norm(dim)\n\n        self.attn1 = MultiHeadAttention(dim, num_head, dropout=0.1)\n        self.attn2 = MultiHeadAttention(dim, num_head, dropout=0.1)\n        self.ff = FeedForward(dim, ff_dim)\n\n        self.dropout1 = nn.Dropout(dropout)\n        self.dropout2 = nn.Dropout(dropout)\n        self.dropout3 = nn.Dropout(dropout)\n\n    def forward(self, x, mem, x_mask, mem_mask):\n        x1 = self.norm1(x)\n        x1 = self.attn1(x1, x1, x1, x_mask)  # self-attention\n        x  = x + self.dropout1(x1)\n\n        if mem is not None:\n            x2 = self.norm2(x)\n            x2 = self.attn2(x2, mem, mem, mem_mask)  # encoder input\n            x  = x + self.dropout2(x2)\n\n        x3 = self.norm3(x)\n        x3 = self.ff(x3)\n        x  = x + self.dropout3(x3)\n        return x\n\n    def forward_last(self, x_last, x_cache, mem, mem_mask):\n\n        x_last_norm = self.norm1(x_last)\n        x1 = torch.cat([x_cache, x_last_norm], 1)\n        x_cache = x1.clone() # update\n\n        x1 = self.attn1(x_last_norm, x1, x1)\n        x_last  = x_last + x1\n\n        if mem is not None:\n            x2 = self.norm2(x_last)\n            x2 = self.attn2(x2, mem, mem, mem_mask)\n            x_last = x_last + x2\n\n\n        x3 = self.norm3(x_last)\n        x3 = self.ff(x3)\n        x_last = x_last + x3\n\n        return x_last, x_cache\n\n\n# https://github.com/alexmt-scale/causal-transformer-decoder/blob/master/causal_transformer_decoder/model.py\nclass TransformerDecode(nn.Module):\n    def __init__(self, dim, ff_dim, num_head, num_layer):\n        super().__init__()\n        self.num_layer = num_layer\n\n        self.layer = nn.ModuleList([\n            TransformerDecodeLayer(dim, ff_dim, num_head) for i in range(num_layer)\n        ])\n        self.norm = Norm(dim)\n\n    def forward(self, x, mem, x_mask=None, mem_mask=None):\n\n        for i in range(self.num_layer):\n            x = self.layer[i](x, mem, x_mask, mem_mask)\n\n        x = self.norm(x)\n        return x\n\n    def forward_last(self, x_last, x_cache, mem,  mem_mask=None):\n        batch_size,t,dim = x_last.shape\n        assert(t==1)\n\n        for i in range(self.num_layer):\n            x_last, x_cache[i] = self.layer[i].forward_last(x_last, x_cache[i], mem, mem_mask)\n\n        x_last = self.norm(x_last)\n        return x_last, x_cache\n\n    \n### Decoder with pre transformer    \nclass DecoderWithTransformer(nn.Module):\n    def __init__(self,image_dim,text_dim,decoder_dim,ff_dim,vocab_size,num_layer=3,num_head=8,max_length=300,dropout=0.5):\n        \n        \"\"\"\n        :param image_dim: input size of image network \n        :param text_dim: input size of text network \n        :param decoder_dim: input size of decoder network \n        :param ff_dim: input size of forwardfeed network \n        :param vocab_size: total number of characters used in training\n        :param num_layer: decoder layer numbers\n        :param num_head : head numbers \n        :param max_length : text max lenght \n        :param dropout: dropout rate\n        \"\"\"\n        super(DecoderWithTransformer, self).__init__()\n        self.image_dim = image_dim\n        self.text_dim = text_dim\n        self.decoder_dim = decoder_dim\n        self.ff_dim = ff_dim\n        self.vocab_size = vocab_size\n        self.num_layer = num_layer\n        self.num_head = num_head\n        self.max_length = max_length\n        self.dropout = dropout\n        \n        \n        self.image_encode = nn.Identity()\n        \n        self.text_pos   = PositionEncode1D(text_dim,max_length)\n        self.token_embed = nn.Embedding(vocab_size, text_dim)\n        self.text_decode = TransformerDecode(decoder_dim, ff_dim, num_head, num_layer)\n\n        #---\n        self.logit  = nn.Linear(decoder_dim, vocab_size)\n        self.dropout = nn.Dropout(p=0.5)\n\n        #----\n        # initialization\n        self.token_embed.weight.data.uniform_(-0.1, 0.1)\n        self.logit.bias.data.fill_(0)\n        self.logit.weight.data.uniform_(-0.1, 0.1)\n    \n    @torch.jit.unused\n    def forward(self, image_embed, token, length):\n        device = image_embed.device\n        batch_size = image_embed.size(0)\n        \n#         image_embed = image_embed[:,1:]\n#         enc_self_attns = image_embed[:,0]\n        #---\n        image_embed = self.image_encode(image_embed)\n\n        length, sort_ind = length.squeeze(1).sort(dim=0, descending=True)\n        decode_lengths = (length).tolist()\n        \n        image_embed = image_embed[sort_ind]\n        token = token[sort_ind]\n        max_length = max(decode_lengths)\n        \n        text_embed = self.token_embed(token)\n        text_embed = self.text_pos(text_embed)\n\n        text_mask = 1 - np.triu(np.ones((batch_size, max_length, max_length)), k=1).astype(np.uint8)\n\n        text_mask = torch.autograd.Variable(torch.from_numpy(text_mask)).to(device)\n        text_image_mask = None\n        \n        #----\n        # <todo> mask based on length of token?\n        # <todo> perturb mask as augmentation https://arxiv.org/pdf/2004.13342.pdf\n\n        x = self.text_decode(text_embed, image_embed, text_mask, text_image_mask)\n        logit = self.logit(x)\n        \n        return logit, token, length\n    \n\n    @torch.jit.export\n    def predict(self, image_embed, max_length, tokenizer):\n        ####\n        # predict with argmax\n        ####\n        image_dim = self.image_dim\n        text_dim = self.text_dim\n        decoder_dim = self.decoder_dim\n        num_layer = self.num_layer\n        num_head = self.num_head\n        ff_dim = self.ff_dim\n        vocab_size = self.vocab_size\n\n        #---------------------------------\n        device = image_embed.device\n        batch_size = image_embed.size(0)\n        \n#         image_embed = image_embed[:,1:]\n#         enc_self_attns = image_embed[:,0]\n\n        image_embed = self.image_encode(image_embed)\n\n        token = torch.full((batch_size, max_length), tokenizer.stoi['<pad>'],dtype=torch.long, device=device)\n        text_pos = self.text_pos.pos\n        token[:,0] = tokenizer.stoi['<sos>']\n\n        #-------------------------------------\n        eos = tokenizer.stoi['<eos>']\n        pad = tokenizer.stoi['<pad>']\n        # https://github.com/alexmt-scale/causal-transformer-decoder/blob/master/tests/test_consistency.py\n        # slow version\n#         for t in range(max_length-1):\n#             last_token = token [:,:(t+1)]\n#             text_embed = self.token_embed(last_token)\n#             text_embed = self.text_pos(text_embed).permute(1,0,2).contiguous() #text_embed + text_pos[:,:(t+1)] #\n        \n#             text_mask = np.triu(np.ones((t+1, t+1)), k=1).astype(np.uint8)\n#             text_mask = torch.autograd.Variable(torch.from_numpy(text_mask)==1).to(device)\n        \n#             x = self.text_decode(text_embed, image_embed, text_mask)\n#             x = x.permute(1,0,2).contiguous()\n        \n#             l = self.logit(x[:,-1])\n#             k = torch.argmax(l, -1)  # predict max\n#             token[:, t+1] = k\n#             if ((k == eos) | (k == pad)).all():  break\n\n        ## fast version\n        cache = [torch.empty((batch_size,0,decoder_dim), device=device) for i in range(num_layer)]\n        for t in range(max_length-1):\n            #last_token = token [:,:(t+1)]\n            #text_embed = self.token_embed(last_token)\n            #text_embed = self.text_pos(text_embed) #text_embed + text_pos[:,:(t+1)] #\n\n            last_token = token[:, t]\n            text_embed = self.token_embed(last_token)\n            text_embed = text_embed + text_pos[:,t] #\n            text_embed = text_embed.reshape(batch_size,1,text_dim)\n\n            x, cache = self.text_decode.forward_last(text_embed[:, -1:], cache, image_embed)\n            x = x.reshape(batch_size,decoder_dim)\n\n            l = self.logit(x)\n            k = torch.argmax(l, -1)  # predict max\n            token[:, t+1] = k\n\n            if ((k == eos) | (k == pad)).all():  break\n\n        predictions = token[:, 1:]\n        return predictions        ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **PL Model**","metadata":{}},{"cell_type":"code","source":"class BMSPLModel(pl.LightningModule):\n    def __init__(self, CFG, tokenizer):\n        super().__init__()\n        self.save_hyperparameters()\n        self.CFG = CFG\n        self.tokenizer = tokenizer\n        \n        ### Encoder\n        self.encoder = Encoder(CFG.encoder_name, pretrained=True)\n                \n        print(f'vocab_size is {len(tokenizer)}')\n        ### Decoder\n        if CFG.decoder_mode == 'lstm': ### decoder with LSTM attention\n            self.decoder = DecoderWithAttention(\n                attention_dim=CFG.attention_dim,\n                embed_dim=CFG.embed_dim,\n                decoder_dim=CFG.decoder_dim,\n                vocab_size=len(tokenizer),\n                dropout=CFG.dropout,\n                encoder_dim=CFG.encoder_dim)\n            \n        elif CFG.decoder_mode == 'transformer': ### decoder with transformer\n            self.decoder = DecoderWithTransformer(\n                image_dim=CFG.image_dim,\n                text_dim=CFG.text_dim,\n                decoder_dim=CFG.decoder_dim,\n                ff_dim=CFG.ff_dim,\n                vocab_size=len(tokenizer),\n                num_layer=CFG.num_layer,\n                num_head=CFG.num_head,\n                max_length=CFG.max_length,\n                dropout=CFG.dropout)\n        else:\n            print(f'Decoder mode {CFG.decoder_mode} not supported, please select decoder mode in [lstm,transformer]')\n    \n            ### loss\n        self.criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.stoi[\"<pad>\"])\n    \n    def forward(self, x):\n        features = self.encoder(x)\n        predictions = self.decoder.predict(features, self.CFG.max_length, self.tokenizer)\n        return predictions\n        \n    def get_scheduler(self,optimizer):\n        if self.CFG.scheduler=='ReduceLROnPlateau':\n            scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=self.CFG.factor, \\\n                    patience=self.CFG.patience, verbose=True, eps=self.CFG.eps)\n        elif self.CFG.scheduler=='CosineAnnealingLR':\n            scheduler = CosineAnnealingLR(optimizer, T_max=self.CFG.T_max, eta_min=self.CFG.min_lr, last_epoch=-1)\n        elif CFG.scheduler=='CosineAnnealingWarmRestarts':\n            scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=self.CFG.T_0, T_mult=1, eta_min=self.CFG.min_lr, last_epoch=-1)\n        return scheduler\n    \n    def configure_optimizers(self):\n        encoder_optimizer = Adam(self.encoder.parameters(), lr=self.CFG.encoder_lr, weight_decay=self.CFG.weight_decay, amsgrad=False)\n        encoder_scheduler = self.get_scheduler(encoder_optimizer)\n        \n        decoder_optimizer = Adam(self.decoder.parameters(), lr=self.CFG.decoder_lr, weight_decay=self.CFG.weight_decay, amsgrad=False)\n        decoder_scheduler = self.get_scheduler(decoder_optimizer)\n        return [encoder_optimizer, decoder_optimizer], [encoder_scheduler, decoder_scheduler]\n    \n    @property\n    def automatic_optimization(self) -> bool:\n        return False\n    \n    def training_step(self, batch, batch_idx, optimizer_idx):\n        images, labels, label_lengths = batch\n        \n        features = self.encoder(images)        \n        if self.CFG.decoder_mode == 'lstm':\n            predictions, caps_sorted, decode_lengths, alphas, sort_ind = self.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 = self.criterion(predictions, targets)\n        elif self.CFG.decoder_mode == 'transformer':\n            predictions, caps_sorted, decode_lengths = self.decoder(features, labels, label_lengths)\n            targets = caps_sorted[:, 1:]\n            decode_lengths = [l - 1 for l in decode_lengths]\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 = self.criterion(predictions, targets)\n        else:\n            print('please select decoder mode in [lstm,transformer]')\n        \n        \n        if self.CFG.gradient_accumulation_steps > 1:\n            loss = loss / CFG.gradient_accumulation_steps\n        self.manual_backward(loss)\n        \n        opt_encoder, opt_decoder = self.optimizers()\n        accumulate_gradient_batches = batch_idx % self.CFG.gradient_accumulation_steps == 0\n        ### accumulate gradient batches\n        if accumulate_gradient_batches:\n            opt_encoder.step()\n            opt_decoder.step()\n            opt_encoder.zero_grad()\n            opt_decoder.zero_grad()\n    \n        tqdm_dict = {'train_loss': loss}\n        output = OrderedDict({\n            'loss': loss,\n            'progress_bar': tqdm_dict,\n            'log': tqdm_dict\n        })\n        return output\n    \n    def validation_step(self, batch, batch_idx):\n        images, text_labels = batch\n        \n        features = self.encoder(images)\n        predictions = self.decoder.predict(features, self.CFG.max_length, self.tokenizer)\n        if CFG.decoder_mode == 'lstm':\n            predicted_sequence = torch.argmax(predictions.detach().cpu(), -1).numpy()\n        elif CFG.decoder_mode == 'transformer':\n            predicted_sequence = predictions.data.cpu().numpy()\n\n        _text_preds = tokenizer.predict_to_inchi(predicted_sequence)\n        return {'text_preds': _text_preds, 'text_labels': text_labels}\n\n\n    def get_score(self,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\n\n    def validation_epoch_end(self, outputs):\n        text_labels = np.concatenate([x['text_labels'] for x in outputs])\n        text_preds = np.concatenate([x['text_preds'] for x in outputs])\n#         text_preds = [f\"InChI=1S/{text}\" for text in text_preds]\n        \n        # scoring\n        score = self.get_score(text_labels, text_preds)\n        print(f'Epoch: {self.current_epoch}, Score: {score}')\n        print(f\"labels: {text_labels[:5]}\")\n        print(f\"preds {text_preds[:5]}\")\n        \n        tqdm_dict = {'score': score}\n        output = OrderedDict({\n            'score': score,\n            'progress_bar': tqdm_dict,\n            'log': tqdm_dict\n        })\n        \n        return output","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Train Loop**","metadata":{}},{"cell_type":"code","source":"def train_loop(CFG, folds,fold, tokenizer):\n    print(f'=============== fold: {fold} training =============')\n    print(f'Training with {CFG.decoder_mode} decoder, params batch_size={CFG.batch_size*CFG.gpus}, encoder_lr={CFG.encoder_lr}, decoder_lr={CFG.decoder_lr}, epochs={CFG.epochs}')\n    trn_idx = folds[folds['fold'] != fold].index\n    val_idx = folds[folds['fold'] == fold].index\n\n    train_folds = folds.loc[trn_idx].reset_index(drop=True)\n    valid_folds = folds.loc[val_idx].reset_index(drop=True)\n    valid_labels = valid_folds['InChI'].values\n\n    train_dataset = TrainDataset(CFG, train_folds, tokenizer, transform=get_train_transforms(CFG))\n    valid_dataset = TestDataset(CFG, valid_folds, transform=get_val_transforms(CFG))\n\n\n    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)\n\n    train_dataloader = DataLoader(train_dataset, batch_size=CFG.batch_size, num_workers=CFG.num_workers,\n                          drop_last=True, shuffle=True, pin_memory=True, collate_fn=bms_collate)\n\n    valid_loader = DataLoader(valid_dataset, batch_size=CFG.batch_size, num_workers=CFG.num_workers,\n                          drop_last=False, shuffle=False, pin_memory=True)\n    \n    model = BMSPLModel(CFG, tokenizer)\n    \n    checkpoint_callback = pl.callbacks.ModelCheckpoint(\n        dirpath='./',\n        filename='{fold}_{epoch:02d}_{score:.4f}',\n        save_top_k=1, \n        verbose=True,\n        monitor='score', \n        mode='min')\n\n#     logger = TensorBoardLogger(\n#         save_dir = CFG.save_dir,\n#         # version=1,\n#         name='lightning_logs'\n#     )\n\n#     early_stop_callback = EarlyStopping(\n#         monitor='score',\n#         min_delta=100000.0,\n#         patience=3,\n#         verbose=False,\n#         mode='min',\n#     )\n    \n    trainer = pl.Trainer(\n        gpus=CFG.gpus,\n        precision=CFG.precision,\n        max_epochs=CFG.epochs,\n        num_sanity_val_steps=1 if CFG.debug else 0,\n        checkpoint_callback=checkpoint_callback,\n        check_val_every_n_epoch = 5,\n#         accelerator='ddp', #for multi gpus\n        sync_batchnorm=True,\n#         logger=logger,\n#         callbacks=[early_stop_callback],\n    )\n    \n    trainer.fit(model, train_dataloader=train_dataloader, val_dataloaders=valid_loader)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# **Run**","metadata":{}},{"cell_type":"code","source":"if __name__ == '__main__':\n    pl.seed_everything(seed=CFG.seed)\n    \n    if CFG.debug:\n#         CFG.epochs = 1\n        folds = folds.sample(n=CFG.batch_size*10, random_state=CFG.seed).reset_index(drop=True)\n    print(folds.groupby(['fold']).size())\n    \n    for fold in range(CFG.n_fold):\n        if fold in CFG.trn_fold:\n            train_loop(CFG, folds, fold, tokenizer)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}