{"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":"# About This Notebook\n\nThis notebook is based on https://www.kaggle.com/konradb/model-train-efficientnet & https://www.kaggle.com/konradb/model-infer-efficientnet, with a final score of 8.90 achieved in the BMS competition.","metadata":{}},{"cell_type":"markdown","source":"# Import Libraries","metadata":{}},{"cell_type":"code","source":"!pip install timm","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:29:24.164136Z","iopub.execute_input":"2022-02-17T04:29:24.164852Z","iopub.status.idle":"2022-02-17T04:29:34.652232Z","shell.execute_reply.started":"2022-02-17T04:29:24.164741Z","shell.execute_reply":"2022-02-17T04:29:34.651364Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport re\nimport cv2\nimport gc\nimport timm\nimport time\nimport math\nimport torch\nimport random\nimport Levenshtein\nimport numpy as np\nimport pandas as pd\nimport torch.nn as nn\nfrom tqdm.auto import tqdm\nimport matplotlib.pyplot as plt\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 sklearn.model_selection import StratifiedKFold, GroupKFold, KFold\nfrom logging import getLogger, INFO, FileHandler, Formatter, StreamHandler\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau\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","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:29:34.654349Z","iopub.execute_input":"2022-02-17T04:29:34.654855Z","iopub.status.idle":"2022-02-17T04:29:39.169276Z","shell.execute_reply.started":"2022-02-17T04:29:34.654817Z","shell.execute_reply":"2022-02-17T04:29:39.168474Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(device)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:29:39.170688Z","iopub.execute_input":"2022-02-17T04:29:39.171875Z","iopub.status.idle":"2022-02-17T04:29:39.218193Z","shell.execute_reply.started":"2022-02-17T04:29:39.171834Z","shell.execute_reply":"2022-02-17T04:29:39.217515Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import urllib.request\nimport ssl\n\nssl._create_default_https_context = ssl._create_unverified_context\nresponse = urllib.request.urlopen('https://www.python.org')\nprint(response.read().decode('utf-8'))","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:29:39.220175Z","iopub.execute_input":"2022-02-17T04:29:39.220640Z","iopub.status.idle":"2022-02-17T04:29:39.689997Z","shell.execute_reply.started":"2022-02-17T04:29:39.220601Z","shell.execute_reply":"2022-02-17T04:29:39.689228Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Read Input Data\n> Import the train dataframe containing image IDs, InChI strings, their actual lengths and parsed sequences.","metadata":{}},{"cell_type":"code","source":"# read the input data contained in the pickle file saved previously\ntrain_df = pd.read_pickle('../input/lstm-model/results-6/results-6/train.pkl')\ntrain_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:29:39.691369Z","iopub.execute_input":"2022-02-17T04:29:39.691802Z","iopub.status.idle":"2022-02-17T04:29:47.039103Z","shell.execute_reply.started":"2022-02-17T04:29:39.691763Z","shell.execute_reply":"2022-02-17T04:29:47.038362Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(train_df)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:29:47.040329Z","iopub.execute_input":"2022-02-17T04:29:47.040655Z","iopub.status.idle":"2022-02-17T04:29:47.048369Z","shell.execute_reply.started":"2022-02-17T04:29:47.040612Z","shell.execute_reply":"2022-02-17T04:29:47.047552Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Add File Paths\n> Make the process of reading the input data more efficient by storing paths to files in the train dataframe.","metadata":{}},{"cell_type":"code","source":"def get_file_path(image_id: str) -> str:\n \n \"\"\"\n This method returns the paths to train images by indexing into the overall directory\n and the image_id's components.\n \n :param image_id: ID of the image\n :type image_id: str\n :return: path to image\n :rtype: str\n \"\"\"\n \n # index into original train images if '-' is not present\n if '-' not in image_id:\n return '../input/bms-molecular-translation/train/{}/{}/{}/{}.png'.format(\n image_id[0], image_id[1], image_id[2], image_id\n )\n \n # otherwise, it's an augmented image so need another indexing way\n else:\n return '../input/augmented-set-of-chemical-structures/0. augmentations/0. augmentations/{}.png'.format(\n image_id\n )","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:29:47.049692Z","iopub.execute_input":"2022-02-17T04:29:47.050603Z","iopub.status.idle":"2022-02-17T04:29:47.057253Z","shell.execute_reply.started":"2022-02-17T04:29:47.050564Z","shell.execute_reply":"2022-02-17T04:29:47.056613Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# get file paths\ntrain_df['file_path'] = train_df['image_id'].apply(get_file_path)\ntrain_df.to_csv('./train_df.csv')","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:29:47.058489Z","iopub.execute_input":"2022-02-17T04:29:47.058890Z","iopub.status.idle":"2022-02-17T04:30:27.061992Z","shell.execute_reply.started":"2022-02-17T04:29:47.058854Z","shell.execute_reply":"2022-02-17T04:30:27.061230Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import the file back\ntrain_df = pd.read_csv('./train_df.csv')\n\n# display\ntrain_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:27.063287Z","iopub.execute_input":"2022-02-17T04:30:27.063540Z","iopub.status.idle":"2022-02-17T04:30:41.054108Z","shell.execute_reply.started":"2022-02-17T04:30:27.063493Z","shell.execute_reply":"2022-02-17T04:30:41.053382Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# limit to ±300K data sub-selected by size 200-350 (HxW)\nvalid_ids = pd.read_csv('../input/bmssmalldataset/new_dataset.csv')['image_id']\ntrain_df = train_df[train_df['image_id'].isin(valid_ids)]\nprint(train_df.shape)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:41.057487Z","iopub.execute_input":"2022-02-17T04:30:41.058074Z","iopub.status.idle":"2022-02-17T04:30:43.648070Z","shell.execute_reply.started":"2022-02-17T04:30:41.058035Z","shell.execute_reply":"2022-02-17T04:30:43.647275Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df.reset_index(inplace=True)\ntrain_df.head()","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:43.651796Z","iopub.execute_input":"2022-02-17T04:30:43.652433Z","iopub.status.idle":"2022-02-17T04:30:43.669533Z","shell.execute_reply.started":"2022-02-17T04:30:43.652390Z","shell.execute_reply":"2022-02-17T04:30:43.668596Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img_test = cv2.imread(train_df.loc[0, 'file_path'])\nplt.imshow(img_test)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:43.670638Z","iopub.execute_input":"2022-02-17T04:30:43.671422Z","iopub.status.idle":"2022-02-17T04:30:43.935992Z","shell.execute_reply.started":"2022-02-17T04:30:43.671387Z","shell.execute_reply":"2022-02-17T04:30:43.934383Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def change_fg_bg_colors(img: np.array) -> np.array:\n \"\"\"Change foreground to white and background to black.\n \n :param img: image array\n :type img: np.array\n :return: image with reverted colors\n :rtype: np.array\n \"\"\"\n recolored_img = cv2.subtract(255, img) \n \n return recolored_img","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:43.938721Z","iopub.execute_input":"2022-02-17T04:30:43.938915Z","iopub.status.idle":"2022-02-17T04:30:43.943595Z","shell.execute_reply.started":"2022-02-17T04:30:43.938891Z","shell.execute_reply":"2022-02-17T04:30:43.942694Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"recolored = change_fg_bg_colors(img_test)\nplt.imshow(recolored)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:43.945150Z","iopub.execute_input":"2022-02-17T04:30:43.945679Z","iopub.status.idle":"2022-02-17T04:30:44.161097Z","shell.execute_reply.started":"2022-02-17T04:30:43.945638Z","shell.execute_reply":"2022-02-17T04:30:44.160234Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Read the Tokenizer\n> Import the string to index mapping for each InChI token, saved previously.","metadata":{}},{"cell_type":"code","source":"class Tokenizer(object):\n \n def __init__(self):\n # string to integer mapping\n self.stoi = {}\n # integer to string mapping\n self.itos = {}\n \n def __len__(self) -> None:\n \n \"\"\"\n This method returns the length of token:index map.\n \n :return: length of map\n :rtype: int\n \"\"\"\n # return the length of the map\n return len(self.stoi)\n \n def fit_on_texts(self, texts: list) -> None:\n \n \"\"\"\n This method creates a vocabulary of all tokens contained in provided texts,\n and updates the mapping of token to index, and index to token.\n \n :param texts: list of texts\n :type texts: list\n \"\"\"\n \n # create a storage for all tokens\n vocab = set()\n \n # add tokens from each text to vocabulary\n for text in texts:\n vocab.update(text.split(' '))\n \n # sort the vocabulary in alphabetical order\n vocab = sorted(vocab)\n \n # add start, end and pad for sentence\n vocab.append('')\n vocab.append('')\n vocab.append('')\n \n # update the string to integer mapping, where integer is the index of the token\n for i, s in enumerate(vocab):\n self.stoi[s] = i\n \n # reverse the previous vocabulary to create integer to string mapping\n self.itos = {item[1]: item[0] for item in self.stoi.items()}\n \n def text_to_sequence(self, text: str) -> list:\n \n \"\"\"\n This method converts the given text to a list of its individual tokens,\n including start and end of string symbols.\n \n :param text: input textual data\n :type text: str\n :return: list of tokens\n :rtype: list\n \"\"\"\n \n # storage to append symbols to\n sequence = []\n \n # add the start of string symbol to storage\n sequence.append(self.stoi[''])\n \n # add each token in text to storage\n for s in text.split(' '):\n sequence.append(self.stoi[s])\n \n # add the end of string symbol to storage\n sequence.append(self.stoi[''])\n \n return sequence\n \n def texts_to_sequences(self, texts: list) -> list:\n \n \"\"\"\n This method converts each text in the provided list into sequences of characters.\n Each sequence is appended to a list and the said list is returned.\n \n :param texts: a list of input texts\n :type texts: list\n :return: a list of sequences\n :rtype: list\n \"\"\"\n \n # storage to append sequences to\n sequences = []\n \n # for each text do\n for text in texts:\n # convert the text to a list of characters\n sequence = self.text_to_sequence(text)\n # append the lists of characters to an aggregated list storage\n sequences.append(sequence)\n\n return sequences\n \n def sequence_to_text(self, sequence: list) -> str:\n \n \"\"\"\n This method converts the sequence of characters back into text.\n \n :param sequence: list of characters\n :type sequence: list\n :return: text\n :rtype: str \n \"\"\"\n # join the characters with no space in between\n return ''.join(list(map(lambda i: self.itos[i], sequence)))\n \n def sequences_to_texts(self, sequences: list) -> list:\n \n \"\"\"\n This method converts each provided sequence into text and returns all texts inside a list.\n \n :param sequences: list of character sequences\n :type sequences: list\n :return: list of texts\n :rtype: list\n \"\"\"\n \n # storage for texts\n texts = []\n \n # convert each sequence to text and append to storage\n for sequence in sequences:\n text = self.sequence_to_text(sequence)\n texts.append(text)\n\n return texts\n \n def predict_caption(self, sequence: list) -> str:\n \n \"\"\"\n This method predicts the caption by adding each symbol in sequence to a resulting string.\n This keeps happening up until the end of sentence or padding is met.\n \n :param sequence: list of characters\n :type sequence: list\n :return: image caption\n :rtype: string\n \"\"\"\n \n # storage for the final caption\n caption = ''\n \n # for each index in a sequence of symbols\n for i in sequence:\n # if symbol is the end of sentence or padding, break\n if i == self.stoi[''] or i == self.stoi['']:\n break\n # otherwise, add the symbol to the final caption\n caption += self.itos[i]\n \n return caption\n \n def predict_captions(self, sequences: list) -> list:\n \n \"\"\"\n This method predicts the captions for each sequence in a list of sequences.\n \n :param sequences: list of sequences\n :type sequences: list\n :return: list of final image captions\n :rtype: list\n \"\"\"\n \n # storage for captions\n captions = []\n \n # for each sequence, do\n for sequence in sequences:\n \n # predict the caption per sequence\n caption = self.predict_caption(sequence)\n \n # append to the storage of captions\n captions.append(caption)\n \n return captions\n\n# load the saved tokenizer and print its string to index mapping\ntokenizer = torch.load('../input/lstm-model/results-6/results-6/tokenizer.pth')\nprint(tokenizer.stoi)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.162879Z","iopub.execute_input":"2022-02-17T04:30:44.163279Z","iopub.status.idle":"2022-02-17T04:30:44.195131Z","shell.execute_reply.started":"2022-02-17T04:30:44.163233Z","shell.execute_reply":"2022-02-17T04:30:44.194198Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Setup Configurations\n> Set configurations needed for modelling and training in a separate class.","metadata":{}},{"cell_type":"code","source":"class CFG:\n \n \"\"\"\n Set configurations for modelling and training.\n \"\"\"\n \n debug = False\n apex = False\n max_len = 275\n \n print_freq = 10000\n num_workers = 4\n model_name = 'efficientnet_b1'\n enc_size = 1280\n samp_size = 1000\n \n size = 288\n \n scheduler = 'CosineAnnealingLR'\n epochs = 5\n T_max = 4\n \n encoder_lr = 1e-4\n decoder_lr = 4e-4\n min_lr = 1e-6\n \n batch_size = 32\n weight_decay = 1e-6\n \n gradient_accumulation_steps = 1\n max_grad_norm = 5\n \n attention_dim = 256\n embed_dim = 256\n decoder_dim = 512\n dropout = 0.5\n seed = 42\n n_fold = 5\n trn_fold = [1]\n train = True\n \n prev_model = '../input/lstm-model/efficientnet_b1_fold1_best.pth'","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.196750Z","iopub.execute_input":"2022-02-17T04:30:44.197287Z","iopub.status.idle":"2022-02-17T04:30:44.205877Z","shell.execute_reply.started":"2022-02-17T04:30:44.197245Z","shell.execute_reply":"2022-02-17T04:30:44.205009Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# if in debug mode\nif CFG.debug:\n \n # set number of epochs to 1\n CFG.epochs = 1\n \n # reduce the train set to a 1000 examples\n train_df = train_df.sample(n=CFG.samp_size, random_state=CFG.seed).reset_index(drop=True)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.209064Z","iopub.execute_input":"2022-02-17T04:30:44.209703Z","iopub.status.idle":"2022-02-17T04:30:44.217195Z","shell.execute_reply.started":"2022-02-17T04:30:44.209662Z","shell.execute_reply":"2022-02-17T04:30:44.216376Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utilities\n> This is a set of utility functions used throughout the computations.","metadata":{}},{"cell_type":"code","source":"def get_score(y_true: str, y_pred: str) -> float:\n \n \"\"\"\n This function computes the Levenstein distance between a true label and a prediction.\n This gets computed for all the provided data and an average score is then returned.\n \n :param y_true: true InChI label\n :type y_true: str\n :param y_pred: predicted InChI label\n :type y_pred: str\n :return: average Levenstein score\n :rtype: float\n \"\"\"\n \n # storage for all Levenstein scores\n scores = []\n \n # for each (true label, predicted label) pair, do\n for true, pred in zip(y_true, y_pred):\n \n # find Levenstein distance for the pair and append to storage\n score = Levenshtein.distance(true, pred)\n scores.append(score)\n \n # compute mean Levenstein distance\n avg_score = np.mean(scores)\n \n return avg_score","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.219484Z","iopub.execute_input":"2022-02-17T04:30:44.221086Z","iopub.status.idle":"2022-02-17T04:30:44.228222Z","shell.execute_reply.started":"2022-02-17T04:30:44.221054Z","shell.execute_reply":"2022-02-17T04:30:44.227439Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def init_logger(log_file: str ='./train.log'):\n \n \"\"\"\n Initialize the logger file for training.\n \n :param log_file: name of the logger file\n :type log_file: str\n :return: logger\n :rtype: object\n \"\"\"\n \n # make a reference to a logger instance\n logger = getLogger(__name__)\n \n # specify lowest-severity log message to be handled by the logger\n logger.setLevel(INFO)\n \n # send logging outputs to stream, i.e. sys.out, and format as message\n handler1 = StreamHandler()\n handler1.setFormatter(Formatter(\"%(message)s\"))\n \n # send logging outputs to a disk file and format as message\n handler2 = FileHandler(filename=log_file)\n handler2.setFormatter(Formatter(\"%(message)s\"))\n \n # add both handlers to the logger\n logger.addHandler(handler1)\n logger.addHandler(handler2)\n \n return logger\n\n# initialize the logger\nLOGGER = init_logger()","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.229457Z","iopub.execute_input":"2022-02-17T04:30:44.229952Z","iopub.status.idle":"2022-02-17T04:30:44.238096Z","shell.execute_reply.started":"2022-02-17T04:30:44.229911Z","shell.execute_reply":"2022-02-17T04:30:44.237262Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def seed_torch(seed: int=42) -> None:\n \n \"\"\"\n Seed torch with a specific seed number to ensure code consistency across runs.\n \n :param seed: seed number\n :type seed: int\n \"\"\"\n \n # set random seed\n random.seed(seed)\n \n # set environment seed\n os.environ['PYTHONHASHSEED'] = str(seed)\n \n # set numpy seed\n np.random.seed(seed)\n \n # set torch seeds\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.backends.cudnn.deterministic = True\n \n# seed torch with the seed from configs\nseed_torch(seed=CFG.seed)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.239589Z","iopub.execute_input":"2022-02-17T04:30:44.240041Z","iopub.status.idle":"2022-02-17T04:30:44.252730Z","shell.execute_reply.started":"2022-02-17T04:30:44.239992Z","shell.execute_reply":"2022-02-17T04:30:44.251835Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Cross-Validation Split\n> Here cross validation splits are created for the train dataframe.","metadata":{}},{"cell_type":"code","source":"# create a copy of the train dataframe to modify\nfolds = train_df.copy()\n\n# provide train/validation indices to split data into train/validation sets\nFold = StratifiedKFold(n_splits=CFG.n_fold, shuffle=True, random_state=CFG.seed)\n\n# for fold_number and (train_index, validation_index) in the splitted fold\nfor n, (train_idx, val_idx) in enumerate(Fold.split(folds, folds['InChI_length'])): # folds = separator, \n # len = maxsplit\n # assign fold number\n folds.loc[val_idx, 'fold'] = int(n)\n\n# convert fold number to integer\nfolds['fold'] = folds['fold'].astype(int)\n\n# print the size of each fold\nprint(folds.groupby(['fold']).size())","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.254373Z","iopub.execute_input":"2022-02-17T04:30:44.254908Z","iopub.status.idle":"2022-02-17T04:30:44.463994Z","shell.execute_reply.started":"2022-02-17T04:30:44.254865Z","shell.execute_reply":"2022-02-17T04:30:44.462198Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset\n> Return the input data of the following format: (image, label tensor, label length). This is needed to ensure that the LSTM-Attention model is working as intended.","metadata":{}},{"cell_type":"code","source":"class TrainDataset(Dataset):\n \n \"\"\"\n This class stores train dataset attributes and methods.\n \"\"\"\n \n def __init__(self, df, tokenizer, transform=None):\n \n \"\"\"\n Initialize train dataset attributes.\n \n :param df: train dataframe\n :type df: pd.DataFrame\n :param tokenizer: string tokenizer\n :type tokenizer: object\n :param transform: torch transformation\n :type transform: object\n \"\"\"\n \n # inherit from parent class\n super().__init__()\n \n # assign train dataframe\n self.df = df\n \n # assign tokenizer\n self.tokenizer = tokenizer\n \n # assign file paths\n self.file_paths = df['file_path'].values\n \n # assign tokenized labels\n self.labels = df['InChI_text'].values\n \n # assign transformations\n self.transform = transform\n \n def __len__(self) -> int:\n \n \"\"\"\n Return size of the train dataframe.\n \n :return: length of the dataframe\n :rtype: int\n \"\"\"\n return len(self.df)\n \n def __getitem__(self, idx: int) -> tuple:\n \n \"\"\"\n Get the image, tensor of its label and label length at the inputted index.\n \n :param idx: index of dataframe\n :type idx: int\n :return: image, tensored label and label length\n :rtype: tuple \n \"\"\"\n \n # get file path of the indexed item\n file_path = self.file_paths[idx]\n \n # read in the image using the file path\n image = cv2.imread(file_path)\n \n # revert black to white and white to black\n image = change_fg_bg_colors(image)\n \n # convert the image to RGB and float type\n image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n \n # if transform is specified\n if self.transform:\n \n # augment the image with the transform\n augmented = self.transform(image=image)\n \n # update the image reference to point to the transformed image\n image = augmented['image']\n \n # get the label of the indexed item\n label = self.labels[idx]\n \n # convert label to a sequence of symbols\n label = self.tokenizer.text_to_sequence(label)\n \n # get the length of the label and convert it to Tensor\n label_length = len(label)\n label_length = torch.LongTensor([label_length])\n \n return image, torch.LongTensor(label), label_length","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.465582Z","iopub.execute_input":"2022-02-17T04:30:44.466044Z","iopub.status.idle":"2022-02-17T04:30:44.485817Z","shell.execute_reply.started":"2022-02-17T04:30:44.465992Z","shell.execute_reply":"2022-02-17T04:30:44.485100Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TestDataset(Dataset):\n \n \"\"\"\n This class stores test dataset attributes and methods.\n \"\"\"\n \n def __init__(self, df, transform=None):\n \n \"\"\"\n Initialize test dataset attributes.\n \n :param df: test dataframe\n :type df: pd.DataFrame\n :param tokenizer: string tokenizer\n :type tokenizer: object\n :param transform: torch transformation\n :type transform: object\n \"\"\"\n \n # inherit from parent class\n super().__init__()\n \n # assign train dataframe\n self.df = df\n \n # assign file paths\n self.file_paths = df['file_path'].values\n \n # assign transformations\n self.transform = transform\n \n # assign fixed transformations\n self.fix_transform = Compose([Transpose(p=1), VerticalFlip(p=1)])\n \n def __len__(self) -> int:\n \n \"\"\"\n Return size of the train dataframe.\n \n :return: length of the dataframe\n :rtype: int\n \"\"\"\n return len(self.df)\n \n def __getitem__(self, idx: int) -> np.array:\n \n \"\"\"\n Get the image, tensor of its label and label length at the inputted index.\n \n :param idx: index of dataframe\n :type idx: int\n :return: image\n :rtype: array\n \"\"\"\n \n # get file path of the indexed item\n file_path = self.file_paths[idx]\n \n # read in the image using the file path\n image = cv2.imread(file_path)\n \n # revert black to white and white to black\n image = change_fg_bg_colors(image)\n \n # convert the image to RGB and float type\n image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n \n # get image shape\n h, w, _ = image.shape\n \n # if height exceeds width, fix-transform the image\n if h > w:\n image = self.fix_transform(image=image)['image']\n \n # if transform is specified\n if self.transform:\n \n # augment the image with the transform\n augmented = self.transform(image=image)\n \n # update the image reference to point to the transformed image\n image = augmented['image']\n\n \n return image","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.487230Z","iopub.execute_input":"2022-02-17T04:30:44.488818Z","iopub.status.idle":"2022-02-17T04:30:44.521809Z","shell.execute_reply.started":"2022-02-17T04:30:44.488782Z","shell.execute_reply":"2022-02-17T04:30:44.520715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def bms_collate(batch: tuple) -> tuple:\n \n \"\"\"\n Combine images, labels and label lengths per batch.\n \n :param batch: a collection of data points\n :type batch: tuple\n :return: stacked images, labels and label lengths, i.e. batch\n :rtype: tuple\n \"\"\"\n \n # initialize storages for images, labels and label lengths\n imgs, labels, label_lengths = [], [], []\n \n # for each data point, append image, labels and label lengths to respective storages\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 \n # pad each label sequence with the index value\n labels = pad_sequence(labels, batch_first=True, padding_value=tokenizer.stoi[\"\"])\n \n return torch.stack(imgs), labels, torch.stack(label_lengths).reshape(-1, 1)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.523311Z","iopub.execute_input":"2022-02-17T04:30:44.523670Z","iopub.status.idle":"2022-02-17T04:30:44.537189Z","shell.execute_reply.started":"2022-02-17T04:30:44.523550Z","shell.execute_reply":"2022-02-17T04:30:44.536324Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Transformations\n> Define basic torch transforms for the dataset, including resizing, normalizing and tensoring.","metadata":{}},{"cell_type":"code","source":"def get_transforms(*, data):\n \n \"\"\"\n Compose several transforms together, mainly resizing, normalizing with Imagenet weights and tensoring.\n \n :param data: image data\n :type data: np.array\n :return: transformed image data\n :rtype: np.array\n \"\"\"\n if data == 'train':\n return Compose(\n [\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 \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 ])","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.538823Z","iopub.execute_input":"2022-02-17T04:30:44.539136Z","iopub.status.idle":"2022-02-17T04:30:44.550467Z","shell.execute_reply.started":"2022-02-17T04:30:44.539040Z","shell.execute_reply":"2022-02-17T04:30:44.549498Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# define train dataset from the train dataframe, using the tokenizer and transformations\ntrain_ds = TrainDataset(train_df, tokenizer, transform=get_transforms(data='train'))","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.552049Z","iopub.execute_input":"2022-02-17T04:30:44.552365Z","iopub.status.idle":"2022-02-17T04:30:44.558193Z","shell.execute_reply.started":"2022-02-17T04:30:44.552267Z","shell.execute_reply":"2022-02-17T04:30:44.556871Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# display an example from the train dataset\nfor i in range(1):\n \n # get image, label and label length from the train dataset\n image, label, label_length = train_ds[i]\n # convert label sequence to text\n text = tokenizer.sequence_to_text(label.numpy())\n \n # transpose the image and show it\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":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.569663Z","iopub.execute_input":"2022-02-17T04:30:44.569920Z","iopub.status.idle":"2022-02-17T04:30:44.953966Z","shell.execute_reply.started":"2022-02-17T04:30:44.569887Z","shell.execute_reply":"2022-02-17T04:30:44.953238Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Modeling\n> Create encoder, attention and decoder classes for modeling.","metadata":{}},{"cell_type":"code","source":"def _inflate(tensor, times, dim):\n # repeat_dims = [1] * tensor.dim()\n # repeat_dims[dim] = times\n # return tensor.repeat(*repeat_dims)\n return torch.repeat_interleave(tensor, times, dim)\n\n\nclass TopKDecoder(torch.nn.Module):\n r\"\"\"\n Top-K decoding with beam search.\n\n Args:\n decoder_rnn (DecoderRNN): An object of DecoderRNN used for decoding.\n k (int): Size of the beam.\n\n Inputs: inputs, encoder_hidden, encoder_outputs, function, teacher_forcing_ratio\n - **inputs** (seq_len, batch, input_size): list of sequences, whose length is the batch size and within which\n each sequence is a list of token IDs. It is used for teacher forcing when provided. (default is `None`)\n - **encoder_hidden** (num_layers * num_directions, batch_size, hidden_size): tensor containing the features\n in the hidden state `h` of encoder. Used as the initial hidden state of the decoder.\n - **encoder_outputs** (batch, seq_len, hidden_size): tensor with containing the outputs of the encoder.\n Used for attention mechanism (default is `None`).\n - **function** (torch.nn.Module): A function used to generate symbols from RNN hidden state\n (default is `torch.nn.functional.log_softmax`).\n - **teacher_forcing_ratio** (float): The probability that teacher forcing will be used. A random number is\n drawn uniformly from 0-1 for every decoding token, and if the sample is smaller than the given value,\n teacher forcing would be used (default is 0).\n\n Outputs: decoder_outputs, decoder_hidden, ret_dict\n - **decoder_outputs** (batch): batch-length list of tensors with size (max_length, hidden_size) containing the\n outputs of the decoder.\n - **decoder_hidden** (num_layers * num_directions, batch, hidden_size): tensor containing the last hidden\n state of the decoder.\n - **ret_dict**: dictionary containing additional information as follows {*length* : list of integers\n representing lengths of output sequences, *topk_length*: list of integers representing lengths of beam search\n sequences, *sequence* : list of sequences, where each sequence is a list of predicted token IDs,\n *topk_sequence* : list of beam search sequences, each beam is a list of token IDs, *inputs* : target\n outputs if provided for decoding}.\n \"\"\"\n\n def __init__(self, decoder_rnn, k, decoder_dim, max_length, tokenizer):\n super(TopKDecoder, self).__init__()\n self.rnn = decoder_rnn\n self.k = k\n self.hidden_size = decoder_dim # self.rnn.hidden_size\n self.V = len(tokenizer)\n self.SOS = tokenizer.stoi[\"\"]\n self.EOS = tokenizer.stoi[\"\"]\n self.max_length = max_length\n self.tokenizer = tokenizer\n\n def forward(self, inputs=None, encoder_hidden=None, encoder_outputs=None, function=F.log_softmax,\n teacher_forcing_ratio=0, retain_output_probs=True):\n \"\"\"\n Forward rnn for MAX_LENGTH steps. Look at :func:`seq2seq.models.DecoderRNN.DecoderRNN.forward_rnn` for details.\n \"\"\"\n\n # inputs, batch_size, max_length = self.rnn._validate_args(inputs, encoder_hidden, encoder_outputs,\n # function, teacher_forcing_ratio)\n\n batch_size = encoder_outputs.size(0)\n max_length = self.max_length\n\n self.pos_index = (torch.LongTensor(range(batch_size)) * self.k).view(-1, 1).cuda()\n\n # Inflate the initial hidden states to be of size: b*k x h\n # encoder_hidden = self.rnn._init_state(encoder_hidden)\n if encoder_hidden is None:\n hidden = None\n else:\n if isinstance(encoder_hidden, tuple):\n # hidden = tuple([_inflate(h, self.k, 1) for h in encoder_hidden])\n hidden = tuple([h.squeeze(0) for h in encoder_hidden])\n hidden = tuple([_inflate(h, self.k, 0) for h in hidden])\n hidden = tuple([h.unsqueeze(0) for h in hidden])\n else:\n # hidden = _inflate(encoder_hidden, self.k, 1)\n raise RuntimeError(\"Not supported\")\n\n # ... same idea for encoder_outputs and decoder_outputs\n if True: # self.rnn.use_attention:\n inflated_encoder_outputs = _inflate(encoder_outputs, self.k, 0)\n else:\n inflated_encoder_outputs = None\n\n # Initialize the scores; for the first step,\n # ignore the inflated copies to avoid duplicate entries in the top k\n sequence_scores = torch.Tensor(batch_size * self.k, 1)\n sequence_scores.fill_(-float('Inf'))\n sequence_scores.index_fill_(0, torch.LongTensor([i * self.k for i in range(0, batch_size)]), 0.0)\n sequence_scores = sequence_scores.cuda()\n\n # Initialize the input vector\n input_var = torch.transpose(torch.LongTensor([[self.SOS] * batch_size * self.k]), 0, 1).cuda()\n\n # Store decisions for backtracking\n stored_outputs = list()\n stored_scores = list()\n stored_predecessors = list()\n stored_emitted_symbols = list()\n stored_hidden = list()\n\n for i in range(0, max_length):\n\n # Run the RNN one step forward\n log_softmax_output, hidden, _ = self.rnn.forward_step(input_var, hidden,\n inflated_encoder_outputs, function=function)\n # If doing local backprop (e.g. supervised training), retain the output layer\n if retain_output_probs:\n stored_outputs.append(log_softmax_output)\n\n # To get the full sequence scores for the new candidates, add the local scores for t_i to the predecessor scores for t_(i-1)\n sequence_scores = _inflate(sequence_scores, self.V, 1)\n sequence_scores += log_softmax_output.squeeze(1)\n scores, candidates = sequence_scores.view(batch_size, -1).topk(self.k, dim=1)\n\n # Reshape input = (bk, 1) and sequence_scores = (bk, 1)\n input_var = (candidates % self.V).view(batch_size * self.k, 1)\n sequence_scores = scores.view(batch_size * self.k, 1)\n\n # Update fields for next timestep\n predecessors = (candidates // self.V + self.pos_index.expand_as(candidates)).view(batch_size * self.k, 1)\n if isinstance(hidden, tuple):\n hidden = tuple([h.index_select(1, predecessors.squeeze()) for h in hidden])\n else:\n hidden = hidden.index_select(1, predecessors.squeeze())\n\n # Update sequence scores and erase scores for end-of-sentence symbol so that they aren't expanded\n stored_scores.append(sequence_scores.clone())\n eos_indices = input_var.data.eq(self.EOS)\n if eos_indices.nonzero().dim() > 0:\n sequence_scores.data.masked_fill_(eos_indices, -float('inf'))\n\n # Cache results for backtracking\n stored_predecessors.append(predecessors)\n stored_emitted_symbols.append(input_var)\n stored_hidden.append(hidden)\n\n # Do backtracking to return the optimal values\n output, h_t, h_n, s, l, p = self._backtrack(stored_outputs, stored_hidden,\n stored_predecessors, stored_emitted_symbols,\n stored_scores, batch_size, self.hidden_size)\n\n # Build return objects\n decoder_outputs = [step[:, 0, :] for step in output]\n if isinstance(h_n, tuple):\n decoder_hidden = tuple([h[:, :, 0, :] for h in h_n])\n else:\n decoder_hidden = h_n[:, :, 0, :]\n metadata = {}\n metadata['inputs'] = inputs\n metadata['output'] = output\n metadata['h_t'] = h_t\n metadata['score'] = s\n metadata['topk_length'] = l\n metadata['topk_sequence'] = p\n metadata['length'] = [seq_len[0] for seq_len in l]\n metadata['sequence'] = [seq[0] for seq in p]\n return decoder_outputs, decoder_hidden, metadata\n\n def _backtrack(self, nw_output, nw_hidden, predecessors, symbols, scores, b, hidden_size):\n \"\"\"Backtracks over batch to generate optimal k-sequences.\n\n Args:\n nw_output [(batch*k, vocab_size)] * sequence_length: A Tensor of outputs from network\n nw_hidden [(num_layers, batch*k, hidden_size)] * sequence_length: A Tensor of hidden states from network\n predecessors [(batch*k)] * sequence_length: A Tensor of predecessors\n symbols [(batch*k)] * sequence_length: A Tensor of predicted tokens\n scores [(batch*k)] * sequence_length: A Tensor containing sequence scores for every token t = [0, ... , seq_len - 1]\n b: Size of the batch\n hidden_size: Size of the hidden state\n\n Returns:\n output [(batch, k, vocab_size)] * sequence_length: A list of the output probabilities (p_n)\n from the last layer of the RNN, for every n = [0, ... , seq_len - 1]\n\n h_t [(batch, k, hidden_size)] * sequence_length: A list containing the output features (h_n)\n from the last layer of the RNN, for every n = [0, ... , seq_len - 1]\n\n h_n(batch, k, hidden_size): A Tensor containing the last hidden state for all top-k sequences.\n\n score [batch, k]: A list containing the final scores for all top-k sequences\n\n length [batch, k]: A list specifying the length of each sequence in the top-k candidates\n\n p (batch, k, sequence_len): A Tensor containing predicted sequence\n \"\"\"\n\n lstm = isinstance(nw_hidden[0], tuple)\n\n # initialize return variables given different types\n output = list()\n h_t = list()\n p = list()\n # Placeholder for last hidden state of top-k sequences.\n # If a (top-k) sequence ends early in decoding, `h_n` contains\n # its hidden state when it sees EOS. Otherwise, `h_n` contains\n # the last hidden state of decoding.\n if lstm:\n state_size = nw_hidden[0][0].size()\n h_n = tuple([torch.zeros(state_size).cuda(), torch.zeros(state_size).cuda()])\n else:\n h_n = torch.zeros(nw_hidden[0].size()).cuda()\n l = [[self.max_length] * self.k for _ in range(b)] # Placeholder for lengths of top-k sequences\n # Similar to `h_n`\n\n # the last step output of the beams are not sorted\n # thus they are sorted here\n sorted_score, sorted_idx = scores[-1].view(b, self.k).topk(self.k)\n # initialize the sequence scores with the sorted last step beam scores\n s = sorted_score.clone()\n\n batch_eos_found = [0] * b # the number of EOS found\n # in the backward loop below for each batch\n\n t = self.max_length - 1\n # initialize the back pointer with the sorted order of the last step beams.\n # add self.pos_index for indexing variable with b*k as the first dimension.\n t_predecessors = (sorted_idx + self.pos_index.expand_as(sorted_idx)).view(b * self.k)\n while t >= 0:\n # Re-order the variables with the back pointer\n current_output = nw_output[t].index_select(0, t_predecessors)\n if lstm:\n current_hidden = tuple([h.index_select(1, t_predecessors) for h in nw_hidden[t]])\n else:\n current_hidden = nw_hidden[t].index_select(1, t_predecessors)\n current_symbol = symbols[t].index_select(0, t_predecessors)\n # Re-order the back pointer of the previous step with the back pointer of\n # the current step\n t_predecessors = predecessors[t].index_select(0, t_predecessors).squeeze()\n\n # This tricky block handles dropped sequences that see EOS earlier.\n # The basic idea is summarized below:\n #\n # Terms:\n # Ended sequences = sequences that see EOS early and dropped\n # Survived sequences = sequences in the last step of the beams\n #\n # Although the ended sequences are dropped during decoding,\n # their generated symbols and complete backtracking information are still\n # in the backtracking variables.\n # For each batch, everytime we see an EOS in the backtracking process,\n # 1. If there is survived sequences in the return variables, replace\n # the one with the lowest survived sequence score with the new ended\n # sequences\n # 2. Otherwise, replace the ended sequence with the lowest sequence\n # score with the new ended sequence\n #\n eos_indices = symbols[t].data.squeeze(1).eq(self.EOS).nonzero()\n if eos_indices.dim() > 0:\n for i in range(eos_indices.size(0) - 1, -1, -1):\n # Indices of the EOS symbol for both variables\n # with b*k as the first dimension, and b, k for\n # the first two dimensions\n idx = eos_indices[i]\n b_idx = int(idx[0] // self.k)\n # The indices of the replacing position\n # according to the replacement strategy noted above\n res_k_idx = self.k - (batch_eos_found[b_idx] % self.k) - 1\n batch_eos_found[b_idx] += 1\n res_idx = b_idx * self.k + res_k_idx\n\n # Replace the old information in return variables\n # with the new ended sequence information\n t_predecessors[res_idx] = predecessors[t][idx[0]]\n current_output[res_idx, :] = nw_output[t][idx[0], :]\n if lstm:\n current_hidden[0][:, res_idx, :] = nw_hidden[t][0][:, idx[0], :]\n current_hidden[1][:, res_idx, :] = nw_hidden[t][1][:, idx[0], :]\n h_n[0][:, res_idx, :] = nw_hidden[t][0][:, idx[0], :].data\n h_n[1][:, res_idx, :] = nw_hidden[t][1][:, idx[0], :].data\n else:\n current_hidden[:, res_idx, :] = nw_hidden[t][:, idx[0], :]\n h_n[:, res_idx, :] = nw_hidden[t][:, idx[0], :].data\n current_symbol[res_idx, :] = symbols[t][idx[0]]\n s[b_idx, res_k_idx] = scores[t][idx[0]].data[0]\n l[b_idx][res_k_idx] = t + 1\n\n # record the back tracked results\n output.append(current_output)\n h_t.append(current_hidden)\n p.append(current_symbol)\n\n t -= 1\n\n # Sort and re-order again as the added ended sequences may change\n # the order (very unlikely)\n s, re_sorted_idx = s.topk(self.k)\n for b_idx in range(b):\n l[b_idx] = [l[b_idx][k_idx.item()] for k_idx in re_sorted_idx[b_idx, :]]\n\n re_sorted_idx = (re_sorted_idx + self.pos_index.expand_as(re_sorted_idx)).view(b * self.k)\n\n # Reverse the sequences and re-order at the same time\n # It is reversed because the backtracking happens in reverse time order\n output = [step.index_select(0, re_sorted_idx).view(b, self.k, -1) for step in reversed(output)]\n p = [step.index_select(0, re_sorted_idx).view(b, self.k, -1) for step in reversed(p)]\n if lstm:\n h_t = [tuple([h.index_select(1, re_sorted_idx).view(-1, b, self.k, hidden_size) for h in step]) for step in reversed(h_t)]\n h_n = tuple([h.index_select(1, re_sorted_idx.data).view(-1, b, self.k, hidden_size) for h in h_n])\n else:\n h_t = [step.index_select(1, re_sorted_idx).view(-1, b, self.k, hidden_size) for step in reversed(h_t)]\n h_n = h_n.index_select(1, re_sorted_idx.data).view(-1, b, self.k, hidden_size)\n s = s.data\n\n return output, h_t, h_n, s, l, p\n\n def _mask_symbol_scores(self, score, idx, masking_score=-float('inf')):\n score[idx] = masking_score\n\n def _mask(self, tensor, idx, dim=0, masking_score=-float('inf')):\n if len(idx.size()) > 0:\n indices = idx[:, 0]\n tensor.index_fill_(dim, indices, masking_score)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:44.959066Z","iopub.execute_input":"2022-02-17T04:30:44.961102Z","iopub.status.idle":"2022-02-17T04:30:45.210108Z","shell.execute_reply.started":"2022-02-17T04:30:44.961053Z","shell.execute_reply":"2022-02-17T04:30:45.209348Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Encoder(nn.Module):\n \n \"\"\"\n Encodes the image with 3 color channels into a smaller learned image.\n \"\"\"\n \n def __init__(self, model_name=CFG.model_name, pretrained=False):\n \n \"\"\"\n Initialize the encoder with CNN equal to the chosen model and set pretrained parameter.\n \n :param model_name: name of the model\n :type model_name: str\n :param pretrained: pretrained weights or not\n :type pretrained: Boolean\n \"\"\"\n \n # inherit attributes and methods from parent class\n super().__init__()\n \n # create a chosen CNN model\n self.cnn = timm.create_model(model_name, pretrained=pretrained)\n \n def forward(self, x):\n \n \"\"\"\n Propagate the input forward.\n \n :param x: image data\n :type x: np.array\n :return: image features\n :rtype: np.array \n \"\"\"\n \n # get batch size from image dimensions\n bs = x.size(0)\n \n # get image features from the CNN\n features = self.cnn.forward_features(x)\n \n # re-arrange the dimensions so that (bs, encoded_image_size, encoded_image_size, n_channels=2048)\n features = features.permute(0, 2, 3, 1)\n \n return features","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:45.211435Z","iopub.execute_input":"2022-02-17T04:30:45.212965Z","iopub.status.idle":"2022-02-17T04:30:46.102549Z","shell.execute_reply.started":"2022-02-17T04:30:45.212924Z","shell.execute_reply":"2022-02-17T04:30:46.101606Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Attention(nn.Module):\n \n \"\"\"\n Define Attention network to calculate the attention value.\n \"\"\"\n \n def __init__(self, encoder_dim, decoder_dim, attention_dim):\n \n \"\"\"\n Initialize the Attention network.\n \n :param encoder_dim: input size of the encoder network\n :type encoder_dim: int\n :param decoder_dim: input size of the decoder network\n :type decoder_dim: int\n :param attention_dim: input size of the attention network\n :type attention_dim: int\n \"\"\"\n \n # initiliaze and inherit from parent class\n super(Attention, self).__init__()\n \n # 1st linear layer to transform the encoded image\n self.encoder_att = nn.Linear(encoder_dim, attention_dim)\n \n # 1st linear layer to transform the decoder's output\n self.decoder_att = nn.Linear(decoder_dim, attention_dim)\n \n # linear layer to calculate values to be softmax-ed\n self.full_att = nn.Linear(attention_dim, 1)\n \n # define ReLU function\n self.relu = nn.ReLU()\n \n # define the softmax layer to calculate weights\n self.softmax = nn.Softmax(dim=1)\n \n def forward(self, encoder_out, decoder_hidden):\n \n \"\"\"\n Propagate the inputs forward.\n \n :param encoder_out: encoded images of dimension (batch_size, num_pixels, encoder_dim)\n :type encoder_out: tensor\n :param decoder_hidden: previous decoder output of dimension (batch_size, decoder_dim)\n :type decoder_hidden: tensor\n :return: attention-weighted encoding, weights\n :rtype: tensor of (batch_size, encoder_dim)\n \"\"\"\n \n # apply the linear layer to encoded images to get (bs, num_pixels, attention_dim)\n att1 = self.encoder_att(encoder_out)\n \n # apply the linear layer to transform the decoder's output to get (bs, attention_dim)\n att2 = self.decoder_att(decoder_hidden)\n \n # pass the sum of transformed encoded images and decoder outputs through a ReLU and apply a linear layer\n # gets (bs, num_pixels)\n att = self.full_att(self.relu(att1 + att2.unsqueeze(1))).squeeze(2)\n \n # pass the results through a softmax layer to get weights\n alpha = self.softmax(att)\n \n # apply the resulting (bs, num_pixels) weights to the encoder output and sum across dim=1\n attention_weighted_encoding = (encoder_out * alpha.unsqueeze(2)).sum(dim=1)\n \n return attention_weighted_encoding, alpha","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.105794Z","iopub.execute_input":"2022-02-17T04:30:46.106037Z","iopub.status.idle":"2022-02-17T04:30:46.117614Z","shell.execute_reply.started":"2022-02-17T04:30:46.106003Z","shell.execute_reply":"2022-02-17T04:30:46.116849Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DecoderWithAttention(nn.Module):\n \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=CFG.enc_size, dropout=CFG.dropout):\n \n \"\"\"\n Initialize the decoder with attention network.\n \n :param attention_dim: input size of the attention network\n :type attention_dim: int\n :param embed_dim: input size of the embedding network\n :type embed_dim: int\n :param decoder_dim: input size of the decoder network\n :type decoder_dim: int\n :param vocab_size: total number of characters used in training\n :type vocab_size: int\n :param encoder_dim: input size of the encoder network\n :type encoder_dim: int\n :param dropout: dropout rate\n :type dropout: float\n \"\"\"\n \n # inherit from parent class\n super(DecoderWithAttention, self).__init__()\n \n # set dimensions of the encoder, attention, embedder and decoder\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 \n # set vocabulary size, dropout rate and the used device\n self.vocab_size = vocab_size\n self.dropout = dropout\n self.device = device\n \n # set attention network, embedding network and dropout layer\n self.attention = Attention(encoder_dim, decoder_dim, attention_dim)\n self.embedding = nn.Embedding(vocab_size, embed_dim)\n self.dropout = nn.Dropout(p=self.dropout)\n \n # set the LSTM Cell decoder\n self.decode_step = nn.LSTMCell(embed_dim+encoder_dim, decoder_dim, bias=True)\n \n # linear layer to find the initial hidden state of LSTM Cell\n self.init_h = nn.Linear(encoder_dim, decoder_dim)\n \n # linear layer to find the initial cell state of LSTM Cell\n self.init_c = nn.Linear(encoder_dim, decoder_dim)\n \n # linear layer to create a sigmoid-activated gate\n self.f_beta = nn.Linear(decoder_dim, encoder_dim)\n self.sigmoid = nn.Sigmoid()\n \n # linear layer to find scores over vocabulary\n self.fc = nn.Linear(decoder_dim, vocab_size)\n \n # initialize some layers with uniform distribution\n self.init_weights()\n \n def init_weights(self):\n \n \"\"\"\n Initialize weights with uniform distribution for embedding and FC layers.\n \"\"\"\n \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 \n \"\"\"\n Load the embedding layer with pre-trained embeddings.\n \n :param embeddings: pre-trained embeddings\n \"\"\"\n \n self.embedding.weight = nn.Parameter(embeddings)\n \n def fine_tune_embeddings(self, fine_tune=True):\n \n \"\"\"\n Allow fine-tuning of the embedding layer.\n \n :param fine_tune: allow fine-tuning\n :type fine_tune: Boolean\n \"\"\"\n # loop over each embedding parameter to set the fine-tuning option\n for p in self.embedding.parameters():\n p.requires_grad = fine_tune\n \n def init_hidden_state(self, encoder_out):\n \n \"\"\"\n Create the initial hidden and cell state for decoder's LSTM based on encoded images.\n \n :param encoder_out: encoded images, of size (bs, num_pixels, encoder_dim)\n :type encoder_out: tensor\n :return: initial hidden and cell states\n :rtype: tuple\n \n \"\"\"\n \n # get the mean of the encoded image's dim=1\n mean_encoder_out = encoder_out.mean(dim=1)\n \n # initialize the hidden and cell states\n h = self.init_h(mean_encoder_out)\n c = self.init_c(mean_encoder_out)\n \n return h, c\n \n def forward(self, encoder_out, encoded_captions, caption_lengths):\n \n \"\"\"\n Propagate the inputs forward.\n \n :param encoder_out: encoded images of dimension (bs, num_pixels, encoder_dim)\n :type encoder_out: tensor\n :param encoded_captions: encoded captions of dimension (bs, max_caption_length)\n :type encoded_captions: tensor\n :param caption_lengths: caption lengths of dimension (bs, 1)\n :type caption_lengths: tensor\n :return: scores for vocabulary, sorted encoded captions, decode lengths, weights, sort indices\n :rtype: tuple\n \"\"\"\n batch_size = encoder_out.size(0)\n encoder_dim = encoder_out.size(-1)\n vocab_size = self.vocab_size\n \n # flatten the image\n encoder_out = encoder_out.view(batch_size, -1, encoder_dim)\n num_pixels = encoder_out.size(1)\n \n # sort input data by decreasing lengths\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 # embed the encoded captions\n embeddings = self.embedding(encoded_captions)\n \n # initialize LSTM state\n h, c = self.init_hidden_state(encoder_out)\n \n # decoding lengths are actual lengths - 1 because of \n decode_lengths = (caption_lengths - 1).tolist()\n \n # initialize tensors to hold word predictions scores and alphas (weights)\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 \n # at each time step, decode by \n for t in range(max(decode_lengths)):\n batch_size_t = sum([l>t for l in decode_lengths])\n \n # attention weighted encoder's output based on decoder's previous state output\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]))\n attention_weighted_encoding = gate * attention_weighted_encoding\n \n # generate a new word with previous word and 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])\n )\n preds = self.fc(self.dropout(h))\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 \"\"\"\n Predict captions.\n \n :param encoder_out: encoded images of dimension (bs, num_pixels, encoder_dim)\n :type encoder_out: tensor\n :param decode_lengths: lengths of decoded captions\n :type decode_lengths: tensor\n :param tokenizer: word tokenizer\n :type tokenizer: dict\n :return: predictions\n :rtype: str\n \"\"\"\n batch_size = encoder_out.size(0)\n encoder_dim = encoder_out.size(-1)\n vocab_size = self.vocab_size\n \n # flatten the image\n encoder_out = encoder_out.view(batch_size, -1, encoder_dim)\n num_pixels = encoder_out.size(1)\n \n # define start tokens and embed them\n start_tokens = torch.ones(batch_size, dtype=torch.long).to(self.device) * tokenizer.stoi[\"\"]\n embeddings = self.embedding(start_tokens)\n \n # initialize LSTM state\n h, c = self.init_hidden_state(encoder_out) # (batch_size, decoder_dim)\n \n # initialize tensors to hold predictions and final conditions\n predictions = torch.zeros(batch_size, decode_lengths, vocab_size).to(self.device)\n end_condition = torch.zeros(batch_size, dtype=torch.long).to(encoder_out.device)\n \n # at each time step decode by\n for t in range(decode_lengths):\n \n # applying attention to encoded image and decoder's previous state output\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 \n # generate a new word with previous word and 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 \n # if end of sentence is reached, stop predicting\n end_condition |= (torch.argmax(preds, -1) == tokenizer.stoi[\"\"])\n if end_condition.sum() == batch_size:\n break\n \n # embed the predictions\n embeddings = self.embedding(torch.argmax(preds, -1))\n \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":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.120281Z","iopub.execute_input":"2022-02-17T04:30:46.120498Z","iopub.status.idle":"2022-02-17T04:30:46.156273Z","shell.execute_reply.started":"2022-02-17T04:30:46.120474Z","shell.execute_reply":"2022-02-17T04:30:46.155457Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Helper Functions\n> This is a set of functions used as utilities.","metadata":{}},{"cell_type":"code","source":"class AverageMeter(object):\n \n \"\"\"\n Computes and stores average and current values.\n \"\"\"\n \n def __init__(self):\n \n \"\"\"\n Reset settings.\n \"\"\"\n self.reset()\n \n def reset(self):\n \n \"\"\"\n Set current, average values, sum and count values to zero.\n \"\"\"\n \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 \n \"\"\"\n Update current value, sum, count and average value.\n \"\"\"\n \n self.val = val\n self.sum += val*n\n self.count += n\n self.avg = self.sum / self.count","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.157898Z","iopub.execute_input":"2022-02-17T04:30:46.158496Z","iopub.status.idle":"2022-02-17T04:30:46.166354Z","shell.execute_reply.started":"2022-02-17T04:30:46.158458Z","shell.execute_reply":"2022-02-17T04:30:46.165576Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def asMinutes(s):\n \n \"\"\"\n Convert seconds to minutes.\n \n :param s: seconds\n :type s: float\n \"\"\"\n m = math.floor(s/60)\n s -= m * 60\n return '%d %ds'% (m, s)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.169504Z","iopub.execute_input":"2022-02-17T04:30:46.170081Z","iopub.status.idle":"2022-02-17T04:30:46.175279Z","shell.execute_reply.started":"2022-02-17T04:30:46.170000Z","shell.execute_reply":"2022-02-17T04:30:46.174470Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def timeSince(since, percent):\n \n \"\"\"\n Calculate time since.\n \n :param since: previous date\n :type since: time\n \"\"\"\n \n now = time.time()\n s = now - since\n es = s / (percent)\n rs = es - s\n \n return '%s (remain %s)' % (asMinutes(s), asMinutes(rs))","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.176694Z","iopub.execute_input":"2022-02-17T04:30:46.177061Z","iopub.status.idle":"2022-02-17T04:30:46.184129Z","shell.execute_reply.started":"2022-02-17T04:30:46.177022Z","shell.execute_reply":"2022-02-17T04:30:46.183250Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_fn(\n train_loader, encoder, decoder, \n criterion, encoder_optimizer, decoder_optimizer, \n epoch, encoder_scheduler, decoder_scheduler, device\n):\n \"\"\"\n Perform one epoch training.\n \n :param train_loader: data loader for training data\n :type train_loader: DataLoader\n :param encoder: encoder model\n :type encoder: Encoder\n :param decoder: decoder model\n :type decoder: Decoder\n :param criterion: loss layer\n :type criterion: Loss \n :param encoder_optimizer: optimizer for encoder\n :type encoder_optimizer: Optimizer\n :param decoder_optimizer: optimizer for decoder\n :type decoder_optimizer: Optimizer\n :param epoch: Epoch number\n :type epoch: int\n :param encoder_scheduler: Encoder scheduler\n :type encoder_scheduler: Encoder\n :param decoder_scheduler: Decoder scheduler\n :type decoder_scheduler: Decoder\n :param device: device selection\n :type device: Device\n :return: Average loss\n :rtype: float\n \"\"\"\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 start = end = time.time()\n global_step = 0\n \n # for index and inputs (imgs, labels and label lengths) in train dataset, do\n for step, (images, labels, label_lengths) in enumerate(train_loader):\n \n # update time step\n data_time.update(time.time() - end)\n \n # send inputs to device\n images = images.to(device)\n labels = labels.to(device)\n label_lengths = label_lengths.to(device)\n \n # set batch size\n batch_size = images.size(0)\n \n # encode images to get features\n features = encoder(images)\n \n # return scores for vocabulary, sorted encoded captions, decode lengths, weights, sort indices\n predictions, caps_sorted, decode_lengths, alphas, sort_ind = decoder(features, labels, label_lengths)\n \n # specify true targets (true y's) after \n targets = caps_sorted[:, 1:]\n \n # get predictions in a packed sequence\n predictions = pack_padded_sequence(predictions, decode_lengths, batch_first=True).data\n \n # get targets in a packed sequence\n targets = pack_padded_sequence(targets, decode_lengths, batch_first=True).data\n \n # compute loss between predictions and targets\n loss = criterion(predictions, targets)\n \n # record the loss\n losses.update(loss.item(), batch_size)\n \n # if optimizer steps are more than 1, divide loss by the number of those steps\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 # backward propagation of loss\n loss.backward()\n \n # perform gradient clipping to avoid exploding gradients\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 # update weights\n if (step + 1) % CFG.gradient_accumulation_steps == 0:\n \n # take a step based on gradients\n encoder_optimizer.step()\n decoder_optimizer.step()\n \n # clear all gradients from last 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 \n # print out the results per epoch\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 .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 \n return losses.avg\n\n\ndef valid_fn(valid_loader, encoder, decoder, tokenizer, criterion, device):\n \n \"\"\"\n Predict for validation set.\n \n :param valid_loader: Data Loader for validation data\n :type valid_loader: DataLoader\n :param encoder: Encoder\n :type encoder: Encoder\n :param decoder: Decoder\n :type decoder: Decoder\n :param tokenizer: Tokenizer\n :type tokenizer: Tokenizer\n :param criterion: Loss\n :type criterion: Loss \n :param device: device selection\n :type device: Device\n :return: predictions\n :rtype: list\n \"\"\"\n \n batch_time = AverageMeter()\n data_time = AverageMeter()\n \n # switch to evaluation mode\n encoder.eval()\n decoder.eval()\n \n # store predictions here\n text_preds = []\n \n start = end = time.time()\n \n # for each image in validation set\n for step, (images) in enumerate(valid_loader):\n \n # measure data loading time\n data_time.update(time.time() - end)\n \n # send images to device\n images = images.to(device)\n \n # specify batch size\n batch_size = images.size(0)\n \n # disable gradient calculation to avoid CUDA errors\n with torch.no_grad():\n \n # encode images\n features = encoder(images)\n \n # predict sequence using decoder\n predictions = decoder.predict(features, CFG.max_len, tokenizer)\n \n # choose the best predicted sequence\n predicted_sequence = torch.argmax(predictions.detach().cpu(), -1).numpy()\n \n # predict captions from the predicted sequence and append to storage\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 \n # print out the validation results\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), batch_time=batch_time,\n data_time=data_time,\n remain=timeSince(start, float(step+1)/len(valid_loader)),\n ))\n \n # concatenate preds into one string\n text_preds = np.concatenate(text_preds)\n \n return text_preds","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.185763Z","iopub.execute_input":"2022-02-17T04:30:46.186020Z","iopub.status.idle":"2022-02-17T04:30:46.212234Z","shell.execute_reply.started":"2022-02-17T04:30:46.185987Z","shell.execute_reply":"2022-02-17T04:30:46.211433Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Loop\n> Performs training.","metadata":{}},{"cell_type":"code","source":"def train_loop(folds, fold):\n \n \"\"\"\n Perform training in a loop.\n \"\"\"\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 \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, mode='min', factor=CFG.factor, patience=CFG.patience, verbose=True, eps=CFG.eps)\n elif CFG.scheduler=='CosineAnnealingLR':\n scheduler = CosineAnnealingLR(optimizer, T_max=CFG.T_max, eta_min=CFG.min_lr, last_epoch=-1)\n elif CFG.scheduler=='CosineAnnealingWarmRestarts':\n scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=CFG.T_0, T_mult=1, eta_min=CFG.min_lr, last_epoch=-1)\n return scheduler\n\n # ====================================================\n # model & optimizer\n # ====================================================\n states = torch.load(CFG.prev_model, map_location=torch.device('cpu'))\n encoder = Encoder(CFG.model_name, pretrained=True)\n encoder.load_state_dict(states['encoder'])\n encoder.to(device)\n encoder_optimizer = Adam(encoder.parameters(), lr=CFG.encoder_lr, weight_decay=CFG.weight_decay, 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 decoder_dim=CFG.decoder_dim,\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(), lr=CFG.decoder_lr, weight_decay=CFG.weight_decay, amsgrad=False)\n decoder_optimizer.load_state_dict(states['decoder_optimizer'])\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[\"\"])\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 './'+f'{CFG.model_name}_fold{fold}_best.pth')","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.214512Z","iopub.execute_input":"2022-02-17T04:30:46.215086Z","iopub.status.idle":"2022-02-17T04:30:46.239717Z","shell.execute_reply.started":"2022-02-17T04:30:46.215048Z","shell.execute_reply":"2022-02-17T04:30:46.239042Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Main Program - Train","metadata":{}},{"cell_type":"code","source":"def main():\n\n \"\"\"\n Prepare: 1.train 2.folds\n \"\"\"\n\n if CFG.train:\n # train\n for fold in range(CFG.n_fold):\n if fold in CFG.trn_fold:\n print(fold)\n train_loop(folds, fold)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.240669Z","iopub.execute_input":"2022-02-17T04:30:46.241376Z","iopub.status.idle":"2022-02-17T04:30:46.250558Z","shell.execute_reply.started":"2022-02-17T04:30:46.241339Z","shell.execute_reply":"2022-02-17T04:30:46.249838Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# if __name__ == '__main__':\n# main()","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.253636Z","iopub.execute_input":"2022-02-17T04:30:46.253991Z","iopub.status.idle":"2022-02-17T04:30:46.261164Z","shell.execute_reply.started":"2022-02-17T04:30:46.253963Z","shell.execute_reply":"2022-02-17T04:30:46.260393Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Testing (Inference)","metadata":{}},{"cell_type":"code","source":"# def get_test_file_path(image_id):\n\n# return \"../input/bms-molecular-translation/test/\" + \"{}/{}/{}/{}.png\".format(\n# image_id[0], image_id[1], image_id[2], image_id \n# )\n\n# test = pd.read_csv('../input/bms-molecular-translation/sample_submission.csv')\n# test['file_path'] = test['image_id'].apply(get_test_file_path)\n\n# print(f'test.shape: {test.shape}')","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.263801Z","iopub.execute_input":"2022-02-17T04:30:46.264112Z","iopub.status.idle":"2022-02-17T04:30:46.269378Z","shell.execute_reply.started":"2022-02-17T04:30:46.264066Z","shell.execute_reply":"2022-02-17T04:30:46.268335Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# def inference(test_loader, encoder, decoder, tokenizer, device):\n# encoder.eval()\n# decoder.eval()\n# text_preds = []\n \n# # k = 2\n# topk_decoder = TopKDecoder(decoder, 2, CFG.decoder_dim, CFG.max_len, tokenizer)\n \n# tk0 = tqdm(test_loader, total=len(test_loader))\n# for images in tk0:\n# images = images.to(device)\n# predictions = []\n# with torch.no_grad():\n# encoder_out = encoder(images)\n# batch_size = encoder_out.size(0)\n# encoder_dim = encoder_out.size(-1)\n# encoder_out = encoder_out.view(batch_size, -1, encoder_dim)\n# h, c = decoder.init_hidden_state(encoder_out)\n# hidden = (h.unsqueeze(0), c.unsqueeze(0))\n \n# decoder_outputs, decoder_hidden, other = topk_decoder(None, hidden, encoder_out)\n \n# for b in range(batch_size):\n# length = other['topk_length'][b][0]\n# tgt_id_seq = [other['topk_sequence'][di][b, 0, 0].item() for di in range(length)]\n# predictions.append(tgt_id_seq)\n# assert len(predictions) == batch_size\n \n# predictions = tokenizer.predict_captions(predictions)\n# predictions = ['InChI=1S/' + p.replace('', '') for p in predictions]\n# # print(predictions[0])\n# text_preds.append(predictions)\n# text_preds = np.concatenate(text_preds)\n# return text_preds","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.271864Z","iopub.execute_input":"2022-02-17T04:30:46.272691Z","iopub.status.idle":"2022-02-17T04:30:46.278125Z","shell.execute_reply.started":"2022-02-17T04:30:46.272640Z","shell.execute_reply":"2022-02-17T04:30:46.277432Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# states = torch.load(CFG.prev_model, map_location=torch.device('cpu'))\n\n# encoder = Encoder(CFG.model_name, pretrained=False)\n# encoder.load_state_dict(states['encoder'])\n# encoder.to(device)\n\n# decoder = DecoderWithAttention(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# device=device)\n\n# decoder.load_state_dict(states['decoder'])\n# decoder.to(device)\n\n# del states; gc.collect()\n\n# test_dataset = TestDataset(test, transform=get_transforms(data='valid'))\n# test_loader = DataLoader(test_dataset, batch_size= 256, shuffle=False, num_workers=CFG.num_workers)\n# predictions = inference(test_loader, encoder, decoder, tokenizer, device)\n\n# del test_loader, encoder, decoder, tokenizer; gc.collect()","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.279484Z","iopub.execute_input":"2022-02-17T04:30:46.279874Z","iopub.status.idle":"2022-02-17T04:30:46.287291Z","shell.execute_reply.started":"2022-02-17T04:30:46.279838Z","shell.execute_reply":"2022-02-17T04:30:46.286637Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# # submission\n# test['InChI'] = [f\"InChI=1S/{text}\" for text in predictions]\n# test[['image_id', 'InChI']].to_csv('./submission.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.289773Z","iopub.execute_input":"2022-02-17T04:30:46.290338Z","iopub.status.idle":"2022-02-17T04:30:46.296670Z","shell.execute_reply.started":"2022-02-17T04:30:46.290287Z","shell.execute_reply":"2022-02-17T04:30:46.296035Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission = pd.read_csv('../input/submission-beam-search/submission.csv')\nsubmission.head(2)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:46.297779Z","iopub.execute_input":"2022-02-17T04:30:46.298423Z","iopub.status.idle":"2022-02-17T04:30:51.600666Z","shell.execute_reply.started":"2022-02-17T04:30:46.298386Z","shell.execute_reply":"2022-02-17T04:30:51.599848Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission['InChI'] = submission['InChI'].str.replace('InChI=1S/InChI=1S/','InChI=1S/')\nsubmission.set_index('image_id', inplace=True)\nsubmission.head(2)","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:51.602507Z","iopub.execute_input":"2022-02-17T04:30:51.602939Z","iopub.status.idle":"2022-02-17T04:30:53.685028Z","shell.execute_reply.started":"2022-02-17T04:30:51.602900Z","shell.execute_reply":"2022-02-17T04:30:53.684322Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission.to_csv('./submission.csv')","metadata":{"execution":{"iopub.status.busy":"2022-02-17T04:30:53.686337Z","iopub.execute_input":"2022-02-17T04:30:53.687155Z","iopub.status.idle":"2022-02-17T04:31:03.698977Z","shell.execute_reply.started":"2022-02-17T04:30:53.687117Z","shell.execute_reply":"2022-02-17T04:31:03.698180Z"},"trusted":true},"execution_count":null,"outputs":[]}]}