{"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":"# SUMMARY\n\nThis notebook builds on the great pipeline [introduced](https://www.kaggle.com/yasufuminakama/inchi-resnet-lstm-with-attention-starter) by Y. Nakama and [adapted to EfficientNets](https://www.kaggle.com/konradb/model-train-efficientnet) by Konrad Banachewicz. The notebook further extendens the pipeline by adding support for multi-layer LSTM in the decoder part. Most of the code changes are concentrated in the model class. Please credit the original authors for their contributions.\n\nThis is training notebook. Inference with multi-layer LSTM decoder is demonstarted [in this notebook](https://www.kaggle.com/kozodoi/efficientnet-multi-layer-lstm-inference).\n\n### References:\n\n- [starter notebook from Y. Nakama](https://www.kaggle.com/yasufuminakama/inchi-resnet-lstm-with-attention-starter)\n- [adapted notebook from Konrad](https://www.kaggle.com/yasufuminakama/inchi-resnet-lstm-with-attention-starter)\n- [PyTorch tutorial on image captioning](https://github.com/sgrvinod/a-PyTorch-Tutorial-to-Image-Captioning)\n- [two-layer RNN implementation](https://github.com/sgrvinod/a-PyTorch-Tutorial-to-Image-Captioning/pull/79)","metadata":{}},{"cell_type":"code","source":"import os\nfrom matplotlib import pyplot as plt\n\nOUTPUT_DIR = './'\nif not os.path.exists(OUTPUT_DIR):\n    os.makedirs(OUTPUT_DIR)\n    \nimport numpy as np\nimport pandas as pd\nimport torch\n\nimport sys\nsys.path.append('../input/timm-pytorch-image-models/pytorch-image-models-master')\n\nimport os\nimport gc\nimport re\nimport math\nimport time\nimport random\nimport shutil\nimport pickle\nfrom pathlib import Path\nfrom contextlib import contextmanager\nfrom collections import defaultdict, Counter\n\nimport scipy as sp\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\n\nimport Levenshtein\nfrom sklearn import preprocessing\nfrom sklearn.model_selection import StratifiedKFold, GroupKFold, KFold\n\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\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\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau\n\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\n\nimport timm\n\nimport warnings \nwarnings.filterwarnings('ignore')\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"papermill":{"duration":5.105933,"end_time":"2021-04-07T08:27:02.380708","exception":false,"start_time":"2021-04-07T08:26:57.274775","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"CFG class now includes a new parameter: `decoder_layers`. For illustration purposes, I am running a two-layer LSTM for 1 epoch on 100k images.","metadata":{}},{"cell_type":"code","source":"#  n_channels_dict = {'efficientnet-b0': 1280, 'efficientnet-b1': 1280, 'efficientnet-b2': 1408,\n#   'efficientnet-b3': 1536, 'efficientnet-b4': 1792, 'efficientnet-b5': 2048,\n#   'efficientnet-b6': 2304, 'efficientnet-b7': 2560}\n\n# This is not, to put it mildly, the most elegant solution ever - but I ran into some trouble \n# with checking the size of feature spaces programmatically inside the CFG definition.\n\nclass CFG:\n    debug          = True\n    apex           = False\n    max_len        = 275\n    print_freq     = 250\n    num_workers    = 4\n    model_name     = 'efficientnet_b2'\n    enc_size       = 1408\n    samp_size      = 100000\n    size           = 288\n    scheduler      = 'CosineAnnealingLR' \n    epochs         = 1 \n    T_max          = 4  \n    encoder_lr     = 1e-4\n    decoder_lr     = 4e-4\n    min_lr         = 1e-6\n    batch_size     = 32\n    weight_decay   = 1e-6\n    gradient_accumulation_steps = 1\n    max_grad_norm  = 10\n    attention_dim  = 256\n    embed_dim      = 512\n    decoder_dim    = 512\n    decoder_layers = 2     # number of LSTM layers\n    dropout        = 0.5\n    seed           = 42\n    n_fold         = 5\n    trn_fold       = 0 \n    train          = True\n    train_path     = '../input/bms-molecular-translation/train/'\n    prep_path      = '../input/preprocessed-stuff/'\n    prev_model     = './prp/muh_best.pth'","metadata":{"papermill":{"duration":0.022485,"end_time":"2021-04-07T08:27:10.713561","exception":false,"start_time":"2021-04-07T08:27:10.691076","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Functions","metadata":{"papermill":{"duration":0.011683,"end_time":"2021-04-07T08:27:10.737552","exception":false,"start_time":"2021-04-07T08:27:10.725869","status":"completed"},"tags":[]}},{"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 = ''\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\ntokenizer = torch.load(CFG.prep_path + 'tokenizer2.pth')\nprint(f\"tokenizer.stoi: {tokenizer.stoi}\")","metadata":{"papermill":{"duration":0.037989,"end_time":"2021-04-07T08:27:10.787242","exception":false,"start_time":"2021-04-07T08:27:10.749253","status":"completed"},"tags":[],"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\n\n\ndef init_logger(log_file=OUTPUT_DIR+'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()\n\n\ndef seed_torch(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n\nseed_torch(seed = CFG.seed)","metadata":{"papermill":{"duration":0.027729,"end_time":"2021-04-07T08:27:10.827442","exception":false,"start_time":"2021-04-07T08:27:10.799713","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ====================================================\n# Dataset\n# ====================================================\n\nclass 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['file_path'].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\n    \n\nclass TestDataset(Dataset):\n    def __init__(self, df, transform=None):\n        super().__init__()\n        self.df = df\n        self.file_paths = df['file_path'].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":{"papermill":{"duration":0.024936,"end_time":"2021-04-07T08:27:10.869493","exception":false,"start_time":"2021-04-07T08:27:10.844557","status":"completed"},"tags":[],"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":{"papermill":{"duration":0.021152,"end_time":"2021-04-07T08:27:10.902922","exception":false,"start_time":"2021-04-07T08:27:10.88177","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"####### CNN ENCODER\n\nclass Encoder(nn.Module):\n    def __init__(self, model_name = CFG.model_name, pretrained = False):\n        super().__init__()\n        self.cnn = timm.create_model(model_name, pretrained = pretrained)\n\n    def forward(self, x):\n        bs       = x.size(0)\n        features = self.cnn.forward_features(x)\n        features = features.permute(0, 2, 3, 1)\n        return features","metadata":{"papermill":{"duration":0.021209,"end_time":"2021-04-07T08:27:10.936685","exception":false,"start_time":"2021-04-07T08:27:10.915476","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The class `DecoderWithAttention` is updated to support a multi-layer LSTM.","metadata":{}},{"cell_type":"code","source":"####### RNN DECODER\n\n# attention module\nclass 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    \n# custom LSTM cell\ndef LSTMCell(input_size, hidden_size, **kwargs):\n    m = nn.LSTMCell(input_size, hidden_size, **kwargs)\n    for name, param in m.named_parameters():\n        if 'weight' in name or 'bias' in name:\n            param.data.uniform_(-0.1, 0.1)\n    return m\n\n\n# decoder\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, device, encoder_dim, dropout, num_layers):\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 num_layers: number of the LSTM layers\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.num_layers    = num_layers\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.ModuleList([LSTMCell(embed_dim + encoder_dim if layer == 0 else embed_dim, embed_dim) for layer in range(self.num_layers)]) # 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) for i in range(self.num_layers)]  # (batch_size, decoder_dim)\n        c = [self.init_c(mean_encoder_out) for i in range(self.num_layers)]\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        \n        # embedding transformed sequence for vector\n        embeddings = self.embedding(encoded_captions)  # (batch_size, max_caption_length, embed_dim)\n        \n        # Initialize LSTM state, initialize cell_vector and hidden_vector\n        prev_h, prev_c = self.init_hidden_state(encoder_out)  # (batch_size, decoder_dim)\n        \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, device = self.device)\n        alphas         = torch.zeros(batch_size, max(decode_lengths), num_pixels, device = self.device)\n        \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],\n                                                                prev_h[-1][:batch_size_t])\n            gate = self.sigmoid(self.f_beta(prev_h[-1][:batch_size_t]))  # gating scalar, (batch_size_t, encoder_dim)\n            attention_weighted_encoding = gate * attention_weighted_encoding\n\n            input = torch.cat([embeddings[:batch_size_t, t, :], attention_weighted_encoding], dim=1)\n            \n            for i, rnn in enumerate(self.decode_step):\n                # recurrent cell\n                h, c = rnn(input, (prev_h[i][:batch_size_t], prev_c[i][:batch_size_t])) # cell_vector and hidden_vector\n\n                # hidden state becomes the input to the next layer\n                input = self.dropout(h)\n\n                # save state for next time step\n                prev_h[i] = h\n                prev_c[i] = c\n                \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            \n        return predictions, encoded_captions, decode_lengths, alphas, sort_ind\n    \n    def predict(self, encoder_out, decode_lengths, tokenizer):\n        \n        # size variables\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        \n        # embed start tocken for LSTM input\n        start_tockens = torch.ones(batch_size, dtype = torch.long, device = self.device) * tokenizer.stoi['<sos>']\n        embeddings    = self.embedding(start_tockens)\n        \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, device = self.device)\n        \n        # predict sequence\n        end_condition = torch.zeros(batch_size, dtype=torch.long, device = self.device)\n        for t in range(decode_lengths):\n            awe, alpha = self.attention(encoder_out, h[-1])  # (s, encoder_dim), (s, num_pixels)\n            gate       = self.sigmoid(self.f_beta(h[-1]))    # gating scalar, (s, encoder_dim)\n            awe        = gate * awe\n            \n            input = torch.cat([embeddings, awe], dim=1)\n \n            for j, rnn in enumerate(self.decode_step):\n                at_h, at_c = rnn(input, (h[j], c[j]))  # (s, decoder_dim)\n                input = self.dropout(at_h)\n                h[j]  = at_h\n                c[j]  = at_c\n            \n            preds = self.fc(self.dropout(h[-1]))  # (batch_size_t, vocab_size)\n            predictions[:, t, :] = preds\n            end_condition |= (torch.argmax(preds, -1) == tokenizer.stoi[\"<eos>\"])\n            if end_condition.sum() == batch_size:\n                break\n            embeddings = self.embedding(torch.argmax(preds, -1))\n        \n        return predictions\n    \n    # beam search\n    def forward_step(self, prev_tokens, hidden, encoder_out, function):\n        \n        h, c = hidden\n        #h, c = h.squeeze(0), c.squeeze(0)\n        h, c = [hi.squeeze(0) for hi in h], [ci.squeeze(0) for ci in c]\n        \n        embeddings = self.embedding(prev_tokens)\n        if embeddings.dim() == 3:\n            embeddings = embeddings.squeeze(1)\n            \n        awe, alpha = self.attention(encoder_out, h[-1])  # (s, encoder_dim), (s, num_pixels)\n        gate       = self.sigmoid(self.f_beta(h[-1]))    # gating scalar, (s, encoder_dim)\n        awe        = gate * awe\n        \n        input = torch.cat([embeddings, awe], dim = 1)\n        for j, rnn in enumerate(self.decode_step):\n            at_h, at_c = rnn(input, (h[j], c[j]))  # (s, decoder_dim)\n            input = self.dropout(at_h)\n            h[j]  = at_h\n            c[j]  = at_c\n\n        preds = self.fc(self.dropout(h[-1]))  # (batch_size_t, vocab_size)\n\n        #hidden = (h.unsqueeze(0), c.unsqueeze(0))\n        hidden = [hi.unsqueeze(0) for hi in h], [ci.unsqueeze(0) for ci in c]\n        predicted_softmax = function(preds, dim = 1)\n        \n        return predicted_softmax, hidden, None","metadata":{"papermill":{"duration":0.067717,"end_time":"2021-04-07T08:27:11.017304","exception":false,"start_time":"2021-04-07T08:27:10.949587","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Helper functions\n\nclass 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\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))\n\n\ndef train_fn(train_loader, encoder, decoder, criterion, \n             encoder_optimizer, decoder_optimizer, epoch,\n             encoder_scheduler, decoder_scheduler, device):\n    \n    batch_time = AverageMeter()\n    data_time  = AverageMeter()\n    losses     = AverageMeter()\n    \n    # switch to train mode\n    encoder.train()\n    decoder.train()\n    \n    start = end = time.time()\n    global_step = 0\n    \n    for step, (images, labels, label_lengths) in enumerate(train_loader):\n        \n        # measure data loading time\n        data_time.update(time.time() - end)\n        \n        images        = images.to(device)\n        labels        = labels.to(device)\n        label_lengths = label_lengths.to(device)\n        batch_size    = images.size(0)\n        \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        \n        # record loss\n        losses.update(loss.item(), batch_size)\n        if CFG.gradient_accumulation_steps > 1:\n            loss = loss / CFG.gradient_accumulation_steps\n            \n        if CFG.apex:\n            with amp.scale_loss(loss, decoder_optimizer) as scaled_loss:\n                scaled_loss.backward()\n        else:\n            loss.backward()\n            \n        encoder_grad_norm = torch.nn.utils.clip_grad_norm_(encoder.parameters(), CFG.max_grad_norm)\n        decoder_grad_norm = torch.nn.utils.clip_grad_norm_(decoder.parameters(), CFG.max_grad_norm)\n        \n        if (step + 1) % CFG.gradient_accumulation_steps == 0:\n            encoder_optimizer.step()\n            decoder_optimizer.step()\n            encoder_optimizer.zero_grad()\n            decoder_optimizer.zero_grad()\n            global_step += 1\n            \n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n        if step % CFG.print_freq == 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                  #'Encoder LR: {encoder_lr:.6f}  '\n                  #'Decoder LR: {decoder_lr:.6f}  '\n                  .format(\n                   epoch+1, step, len(train_loader), \n                   batch_time        = batch_time,\n                   data_time         = data_time, \n                   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                   #encoder_lr=encoder_scheduler.get_lr()[0],\n                   #decoder_lr=decoder_scheduler.get_lr()[0],\n                   ))\n    return losses.avg\n\n\ndef valid_fn(valid_loader, encoder, decoder, tokenizer, criterion, device):\n    \n    batch_time = AverageMeter()\n    data_time  = AverageMeter()\n    \n    # switch to evaluation mode\n    encoder.eval()\n    decoder.eval()\n    \n    text_preds = []\n    start = end = time.time()\n    \n    for step, (images) in enumerate(valid_loader):\n        \n        # measure data loading time\n        data_time.update(time.time() - end)\n        \n        images     = images.to(device)\n        batch_size = images.size(0)\n        \n        with torch.no_grad():\n            features    = encoder(images)\n            predictions = decoder.predict(features, CFG.max_len, tokenizer)\n            \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        \n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n        if step % CFG.print_freq == 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), \n                   batch_time = batch_time,\n                   data_time  = data_time,\n                   remain     = timeSince(start, float(step+1)/len(valid_loader)),\n                   ))\n            \n    text_preds = np.concatenate(text_preds)\n    return text_preds","metadata":{"papermill":{"duration":0.039205,"end_time":"2021-04-07T08:27:11.070219","exception":false,"start_time":"2021-04-07T08:27:11.031014","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ====================================================\n# Train loop\n# ====================================================\ndef train_loop(folds, fold):\n\n    LOGGER.info(f\"========== fold: {fold} training ==========\")\n\n    # ====================================================\n    # loader\n    # ====================================================\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(train_folds, tokenizer, transform = get_transforms(data = 'train'))\n    valid_dataset = TestDataset(valid_folds, transform = get_transforms(data = 'valid'))\n\n    train_loader = DataLoader(train_dataset, \n                              batch_size  = CFG.batch_size, \n                              shuffle     = True, \n                              num_workers = CFG.num_workers, \n                              pin_memory  = True,\n                              drop_last   = True, \n                              collate_fn  = bms_collate)\n    valid_loader = DataLoader(valid_dataset, \n                              batch_size  = CFG.batch_size, \n                              shuffle     = False, \n                              num_workers = CFG.num_workers,\n                              pin_memory  = True, \n                              drop_last   = False)\n    \n    # ====================================================\n    # scheduler \n    # ====================================================\n    def get_scheduler(optimizer):\n        if CFG.scheduler=='ReduceLROnPlateau':\n            scheduler = ReduceLROnPlateau(optimizer, \n                                          mode     = 'min', \n                                          factor   = CFG.factor, \n                                          patience = CFG.patience, \n                                          verbose  = True, \n                                          eps      = CFG.eps)\n        elif CFG.scheduler=='CosineAnnealingLR':\n            scheduler = CosineAnnealingLR(optimizer, \n                                          T_max      = CFG.T_max, \n                                          eta_min    = CFG.min_lr, \n                                          last_epoch = -1)\n        elif CFG.scheduler=='CosineAnnealingWarmRestarts':\n            scheduler = CosineAnnealingWarmRestarts(optimizer, \n                                                    T_0        = CFG.T_0, \n                                                    T_mult     = 1, \n                                                    eta_min    = CFG.min_lr, \n                                                    last_epoch = -1)\n        return scheduler\n\n    # ====================================================\n    # model & optimizer\n    # ====================================================\n\n#    states = torch.load(CFG.prev_model,  map_location=torch.device('cpu'))\n\n    encoder = Encoder(CFG.model_name, \n                      pretrained = True)\n#    encoder.load_state_dict(states['encoder'])\n    \n    encoder.to(device)\n    encoder_optimizer = Adam(encoder.parameters(), \n                             lr           = CFG.encoder_lr, \n                             weight_decay = CFG.weight_decay, \n                             amsgrad      = False)\n#    encoder_optimizer.load_state_dict(states['encoder_optimizer'])\n    encoder_scheduler = get_scheduler(encoder_optimizer)\n#    encoder_scheduler.load_state_dict(states['encoder_scheduler'])\n    \n    decoder = DecoderWithAttention(attention_dim = CFG.attention_dim, \n                                   embed_dim     = CFG.embed_dim, \n                                   encoder_dim   = CFG.enc_size,\n                                   decoder_dim   = CFG.decoder_dim,\n                                   num_layers    = CFG.decoder_layers,\n                                   vocab_size    = len(tokenizer), \n                                   dropout       = CFG.dropout, \n                                   device        = device)\n#    decoder.load_state_dict(states['decoder'])\n    decoder.to(device)\n    decoder_optimizer = Adam(decoder.parameters(), \n                             lr           = CFG.decoder_lr, \n                             weight_decay = CFG.weight_decay, \n                             amsgrad      = False)\n#    decoder_optimizer.load_state_dict(states['decoder_optimizer'])\n\n    decoder_scheduler = get_scheduler(decoder_optimizer)\n #   decoder_scheduler.load_state_dict(states['decoder_scheduler'])\n\n    # ====================================================\n    # loop\n    # ====================================================\n    criterion = nn.CrossEntropyLoss(ignore_index = tokenizer.stoi[\"<pad>\"])\n\n    best_score = np.inf\n    best_loss  = np.inf\n    \n    for epoch in range(CFG.epochs):\n        \n        start_time = time.time()\n        \n        # train\n        avg_loss = train_fn(train_loader, encoder, decoder, criterion, \n                            encoder_optimizer, decoder_optimizer, epoch, \n                            encoder_scheduler, decoder_scheduler, 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        \n        if isinstance(encoder_scheduler, ReduceLROnPlateau):\n            encoder_scheduler.step(score)\n        elif isinstance(encoder_scheduler, CosineAnnealingLR):\n            encoder_scheduler.step()\n        elif isinstance(encoder_scheduler, CosineAnnealingWarmRestarts):\n            encoder_scheduler.step()\n            \n        if isinstance(decoder_scheduler, ReduceLROnPlateau):\n            decoder_scheduler.step(score)\n        elif isinstance(decoder_scheduler, CosineAnnealingLR):\n            decoder_scheduler.step()\n        elif isinstance(decoder_scheduler, CosineAnnealingWarmRestarts):\n            decoder_scheduler.step()\n\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        if score < best_score:\n            best_score = score\n            LOGGER.info(f'Epoch {epoch+1} - Save Best Score: {best_score:.4f} Model')\n            torch.save({'encoder': encoder.state_dict(), \n                        'encoder_optimizer': encoder_optimizer.state_dict(), \n                        'encoder_scheduler': encoder_scheduler.state_dict(), \n                        'decoder': decoder.state_dict(), \n                        'decoder_optimizer': decoder_optimizer.state_dict(), \n                        'decoder_scheduler': decoder_scheduler.state_dict(), \n                        'text_preds': text_preds,\n                       },\n                        OUTPUT_DIR+f'{CFG.model_name}_fold{fold}_best.pth')","metadata":{"papermill":{"duration":0.038367,"end_time":"2021-04-07T08:27:11.123364","exception":false,"start_time":"2021-04-07T08:27:11.084997","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_train_file_path(image_id):\n\n    return CFG.train_path + \"{}/{}/{}/{}.png\".format(\n        image_id[0], image_id[1], image_id[2], image_id \n    )","metadata":{"papermill":{"duration":0.020522,"end_time":"2021-04-07T08:27:11.156788","exception":false,"start_time":"2021-04-07T08:27:11.136266","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# transformations\n\ndef get_transforms(*, data):\n    \n    if data == 'train':\n        return Compose([\n            Resize(CFG.size, CFG.size),\n            HorizontalFlip(p=0.5),                  \n            Transpose(p=0.5),\n            HorizontalFlip(p=0.5),\n            VerticalFlip(p=0.5),\n            ShiftScaleRotate(p=0.5),   \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(CFG.size, CFG.size),\n            Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n            ),\n            ToTensorV2(),\n        ])\n","metadata":{"papermill":{"duration":0.022257,"end_time":"2021-04-07T08:27:11.192091","exception":false,"start_time":"2021-04-07T08:27:11.169834","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data","metadata":{"papermill":{"duration":0.013404,"end_time":"2021-04-07T08:27:11.218463","exception":false,"start_time":"2021-04-07T08:27:11.205059","status":"completed"},"tags":[]}},{"cell_type":"code","source":"train = pd.read_pickle(CFG.prep_path + 'train2.pkl')\n\ntrain['file_path'] = train['image_id'].apply(get_train_file_path)\n\nprint(f'train.shape: {train.shape}')\n\n\nif CFG.debug:\n    CFG.epochs = 1\n    train = train.sample(n = CFG.samp_size, random_state = CFG.seed).reset_index(drop = True)","metadata":{"papermill":{"duration":12.663809,"end_time":"2021-04-07T08:27:23.895404","exception":false,"start_time":"2021-04-07T08:27:11.231595","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = TrainDataset(train, tokenizer, transform = get_transforms(data='train'))\n\nfolds = train.copy()\nFold = StratifiedKFold(n_splits = CFG.n_fold, shuffle = True, random_state = CFG.seed)\nfor n, (train_index, val_index) in enumerate(Fold.split(folds, folds['InChI_length'])):\n    folds.loc[val_index, 'fold'] = int(n)\nfolds['fold'] = folds['fold'].astype(int)","metadata":{"papermill":{"duration":0.02079,"end_time":"2021-04-07T08:27:23.931533","exception":false,"start_time":"2021-04-07T08:27:23.910743","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{"papermill":{"duration":0.013921,"end_time":"2021-04-07T08:27:24.004031","exception":false,"start_time":"2021-04-07T08:27:23.99011","status":"completed"},"tags":[]}},{"cell_type":"code","source":"train_loop(folds, CFG.trn_fold)","metadata":{"papermill":{"duration":39.171513,"end_time":"2021-04-07T08:28:03.18932","exception":false,"start_time":"2021-04-07T08:27:24.017807","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]}]}