{"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":"

Bristol-Myers Squibb – Molecular Translation

\n

Image Captioning Tutorial With TPU Basics - Inference

\n
CREATED BY: DARIEN SCHETTLER
\n\n
\n\n---\n\n**Learn how to use Tensorflow, TFRecords and TPUs to create an Image Captioning Pipeline That Will Yield The Following Performance - All Within 1 9HR Session**\n\n- ~25 Levenshtein Distance On CV (EfficientNetB3 + 1024/512 LSTM) \n- ~TBD Levenshtein Distance on Public LB (EfficientNetB3 + 1024/512 LSTM)\n- ~5 Levenshtein Distance on CV (EfficientNetB7 + 1024/512 LSTM)\n- ~TBD Levenshtein Distance on Public LB (EfficientNetB7 + 1024/512 LSTM)\n\nI really, really tried to make this notebook as understandable as possible. I hope it helps!\n\n
\n\n***NOTE - All of these models still had room for improvement. More finetuning could easily produce better scores***\n\n
\n\n---\n\n
\n\n
\n
🛑🛑🛑   CAUTION: THIS NOTEBOOK IS A WORK IN PROGRESS   🛑🛑🛑

\n
\n\n
\n\n\n---\n\n
\n\n
\n 🌒🌓🌔🌕🌖🌗🌘   MY OTHER NOTEBNOOKS

Here are the links to my other notebooks. Note that at this time the inference notebook is not open source. It will be soon.

\n - Inference [this] Notebook
\n - Training Notebook
\n - TF Record Generation [384x384]
\n - TF Record Generation [256x256]
\n - TF Record Generation [128x128]
\n
\n\n
\n\n---\n\n
\n\n
\n 🙏   CREDIT TO THE FOLLOWING NOTEBOOKS I USED IN CREATING THIS KERNEL:

If you liked this notebook please upvote these other notebooks. Without them I wouldn't have been able to make this!

- Awesome Notebook For Best Practices in Distributed Computing
- The Amazing Mark Wijkhuizen's TPU Training Notebook For This Competition
\n
\n\n
","metadata":{}},{"cell_type":"markdown","source":"

\n\n

\n\n

TABLE OF CONTENTS

\n\n---\n\n

0    IMPORTS

\n\n---\n\n

1    BACKGROUND INFORMATION

\n\n---\n\n

2    SETUP

\n\n---\n\n

3    HELPER FUNCTIONS

\n\n---\n\n

4    PREPARE THE DATASET

\n\n---\n\n

5    MODEL PREPARATION

\n\n---\n\n

6    MODEL TRAINING

\n\n---","metadata":{}},{"cell_type":"markdown","source":"
\n\n

0  IMPORTS    ⤒

","metadata":{}},{"cell_type":"code","source":"# Installs\n# print(\"\\n... PIP INSTALLS STARTING ...\\n\")\n# !pip install -q --upgrade pip\n# !pip install -q pydot\n# !pip install -q pydotplus\n# !apt-get -qq install graphviz\n# !pip install -q git+https://github.com/qubvel/classification_models.git\n# !pip install -q git+https://github.com/qubvel/efficientnet.git\n# print(\"\\n... PIP INSTALLS COMPLETE ...\\n\")\n\n\nprint(\"\\n... IMPORTS STARTING ...\\n\")\nprint(\"\\n\\tVERSION INFORMATION\")\n# Machine Learning and Data Science Imports\nimport tensorflow as tf; print(f\"\\t\\t– TENSORFLOW VERSION: {tf.__version__}\");\nimport tensorflow_addons as tfa; print(f\"\\t\\t– TENSORFLOW ADDONS VERSION: {tfa.__version__}\");\nimport pandas as pd; pd.options.mode.chained_assignment = None;\nimport numpy as np; print(f\"\\t\\t– NUMPY VERSION: {np.__version__}\");\n\n# Library used to easily calculate LD\nimport Levenshtein\n\n# Imports for Alternative Models\n# import classification_models\n# from classification_models.keras import Classifiers\n# import efficientnet.keras as efn\n\n# Built In Imports\nfrom kaggle_datasets import KaggleDatasets\nfrom collections import Counter\nfrom datetime import datetime\nfrom glob import glob\nimport warnings\nimport requests\nimport imageio\nimport IPython\nimport urllib\nimport zipfile\nimport pickle\nimport random\nimport shutil\nimport string\nimport math\nimport time\nimport gzip\nimport ast\nimport io\nimport os\nimport gc\nimport re\n\n# Visualization Imports\nfrom matplotlib.colors import ListedColormap\nimport matplotlib.patches as patches\nimport plotly.graph_objects as go\nimport matplotlib.pyplot as plt\nfrom tqdm.notebook import tqdm; tqdm.pandas();\nimport plotly.express as px\nimport seaborn as sns\nfrom PIL import Image\nimport matplotlib; print(f\"\\t\\t– MATPLOTLIB VERSION: {matplotlib.__version__}\");\nimport plotly\nimport PIL\nimport cv2\n\n\ndef seed_it_all(seed=7):\n \"\"\" Attempt to be Reproducible \"\"\"\n os.environ['PYTHONHASHSEED'] = str(seed)\n random.seed(seed)\n np.random.seed(seed)\n tf.random.set_seed(seed)\n\nprint(\"\\n\\n... IMPORTS COMPLETE ...\\n\")\n \nprint(\"\\n... SEEDING FOR DETERMINISTIC BEHAVIOUR ...\\n\")\nseed_it_all()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-06-03T23:31:02.621073Z","iopub.execute_input":"2021-06-03T23:31:02.621451Z","iopub.status.idle":"2021-06-03T23:31:02.638410Z","shell.execute_reply.started":"2021-06-03T23:31:02.621414Z","shell.execute_reply":"2021-06-03T23:31:02.637264Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n\n

1  BACKGROUND INFORMATION    ⤒

\n\n
PRIMARY TASK DESCRIPTION\n\n\n**Given an image, our goal is to generate a caption. In this case, that image is of a single molecule and the description/caption is the InChI string for that molecule.**\n\n---\n\n
\n\nSECONDARY TASK DESCRIPTION\n\nIn this notebook, we will go through, step by step, training models with TPUs in a custom way. The following steps will be covered:\n* Use **`tf.data.Dataset`** as input pipeline\n* Perform a custom training loop\n* Correctly define loss function\n* Gradient accumulation with TPUs
\n\n
\n\nMORE DETAIL ON IMAGE CAPTIONING\n\n\nDescription From a Tutorial I Used As Reference\n\n>Caption generation is a challenging artificial intelligence problem where a textual description must be generated for a given photograph.\n>\n>It requires both methods from computer vision to understand the content of the image and a language model from the field of natural language processing to turn the understanding of the image into words in the right order. Recently, deep learning methods have achieved state-of-the-art results on examples of this problem.\n>\n>Deep learning methods have demonstrated state-of-the-art results on caption generation problems. What is most impressive about these methods is a single end-to-end model can be defined to predict a caption, given a photo, instead of requiring sophisticated data preparation or a pipeline of specifically designed models.\n","metadata":{}},{"cell_type":"markdown","source":"
\n\n

2  SETUP    ⤒\n

","metadata":{}},{"cell_type":"markdown","source":"

2.1 ACCELERATOR DETECTION

\n\n---\n\nIn order to use **`TPU`**, we use **`TPUClusterResolver`** for the initialization which is necessary to connect to the remote cluster and initialize cloud TPUs. Let's go over two important points\n\n1. When using TPU on Kaggle, you don't need to specify arguments for **`TPUClusterResolver`**\n2. However, on **G**oogle **C**ompute **E**ngine (**GCE**), you will need to do the following:\n\n
\n\n```python\n# The name you gave to the TPU to use\nTPU_WORKER = 'my-tpu-name'\n\n# or you can also specify the grpc path directly\n# TPU_WORKER = 'grpc://xxx.xxx.xxx.xxx:8470'\n\n# The zone you chose when you created the TPU to use on GCP.\nZONE = 'us-east1-b'\n\n# The name of the GCP project where you created the TPU to use on GCP.\nPROJECT = 'my-tpu-project'\n\ntpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=TPU_WORKER, zone=ZONE, project=PROJECT)\n```\n\n
\n 🛑   WARNING:

- Although the Tensorflow documentation says it is the project name that should be provided for the argument `project`, it is actually the Project ID, that you should provide. This can be found on the GCP project dashboard page.
\n
\n\n
\n 📖   REFERENCES:

\n - Guide - Use TPUs
\n - Doc - TPUClusterResolver
\n\n
","metadata":{}},{"cell_type":"code","source":"print(f\"\\n... ACCELERATOR SETUP STARTING ...\\n\")\n\n# Detect hardware, return appropriate distribution strategy\ntry:\n # TPU detection. No parameters necessary if TPU_NAME environment variable is set. On Kaggle this is always the case.\n TPU = tf.distribute.cluster_resolver.TPUClusterResolver() \nexcept ValueError:\n TPU = None\n\nif TPU:\n print(f\"\\n... RUNNING ON TPU - {TPU.master()}...\")\n tf.config.experimental_connect_to_cluster(TPU)\n tf.tpu.experimental.initialize_tpu_system(TPU)\n strategy = tf.distribute.experimental.TPUStrategy(TPU)\nelse:\n print(f\"\\n... RUNNING ON CPU/GPU ...\")\n # Yield the default distribution strategy in Tensorflow\n # --> Works on CPU and single GPU.\n strategy = tf.distribute.get_strategy() \n\n# What Is a Replica?\n# --> A single Cloud TPU device consists of FOUR chips, each of which has TWO TPU cores. \n# --> Therefore, for efficient utilization of Cloud TPU, a program should make use of each of the EIGHT (4x2) cores. \n# --> Each replica is essentially a copy of the training graph that is run on each core and \n# trains a mini-batch containing 1/8th of the overall batch size\nN_REPLICAS = strategy.num_replicas_in_sync\n \nprint(f\"... # OF REPLICAS: {N_REPLICAS} ...\\n\")\n\nprint(f\"\\n... ACCELERATOR SETUP COMPLTED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:31:05.229182Z","iopub.execute_input":"2021-06-03T23:31:05.229751Z","iopub.status.idle":"2021-06-03T23:31:10.637819Z","shell.execute_reply.started":"2021-06-03T23:31:05.229703Z","shell.execute_reply":"2021-06-03T23:31:10.636672Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

2.2 COMPETITION DATA ACCESS

\n\n---\n\nTPUs read data must be read directly from **G**oogle **C**loud **S**torage **(GCS)**. Kaggle provides a utility library – **`KaggleDatasets`** – which has a utility function **`.get_gcs_path`** that will allow us to access the location of our input datasets within **GCS**.

\n\n
\n 📌   TIPS:

- If you have multiple datasets attached to the notebook, you should pass the name of a specific dataset to the `get_gcs_path()` function. In our case, the name of the dataset is the name of the directory the dataset is mounted within.

\n
","metadata":{}},{"cell_type":"code","source":"print(\"\\n... DATA ACCESS SETUP STARTED ...\\n\")\n\nif TPU:\n # Google Cloud Dataset path to training and validation images\n DATA_DIR = KaggleDatasets().get_gcs_path('bms-test-dataset-192x384')\nelse:\n # Local path to training and validation images\n DATA_DIR = \"/kaggle/input/bms-test-dataset-192x384\"\n \nprint(f\"\\n... DATA DIRECTORY PATH IS:\\n\\t--> {DATA_DIR}\")\nprint(f\"\\n... IMMEDIATE CONTENTS OF DATA DIRECTORY IS:\")\nfor file in tf.io.gfile.glob(os.path.join(DATA_DIR, \"*\")): print(f\"\\t--> {file}\")\n\nprint(\"\\n\\n... DATA ACCESS SETUP COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:31:10.641285Z","iopub.execute_input":"2021-06-03T23:31:10.641586Z","iopub.status.idle":"2021-06-03T23:31:11.082541Z","shell.execute_reply.started":"2021-06-03T23:31:10.641555Z","shell.execute_reply":"2021-06-03T23:31:11.081424Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

2.3 LEVERAGING MIXED PRECISION

\n\n---\n\nMixed precision is the use of both **`16-bit`** and **`32-bit`** floating-point types in a model during training to make it run faster and use less memory. By keeping certain parts of the model in the **`32-bit`** types for numeric stability, the model will have a lower step time and train equally as well in terms of the evaluation metrics such as accuracy. \n\nToday, most models use the **`float32`** dtype, which takes **`32`** bits of memory. However, there are two lower-precision dtypes, **`float16`** and **`bfloat16`**, each which take **`16`** bits of memory instead. Modern accelerators can run operations faster in the **`16-bit`** dtypes, as they have specialized hardware to run **`16-bit`** computations and **`16-bit`** dtypes can be read from memory faster.

\n\n**NVIDIA GPUs** can run operations in **`float16`** faster than in **`float32`**
\n**TPUs** can run operations **`bfloat16`** faster than in **`float32`**

\n\nTherefore, these lower-precision dtypes should be used whenever possible on those devices. However, variables and a few computations should still be in **`float32`** for numeric reasons so that the model trains to the same quality. \n\nThe Keras mixed precision API allows you to use a mix of either **`float16`** or **`bfloat16`** with **`float32`**, to get the performance benefits from **`float16/bfloat16`** and the numeric stability benefits from **`float32`**.

\n\n
\n 📖   DEFINITION:

- The term \"numeric stability\" refers to how a model's quality is affected by the use of a lower-precision dtype instead of a higher precision dtype. We say an operation is \"numerically unstable\" in float16 or bfloat16 if running it in one of those dtypes causes the model to have worse evaluation accuracy or other metrics compared to running the operation in float32.
\n
\n\n
\n 📖   REFERENCE:

- TF Mixed Precision Overview
\n
","metadata":{}},{"cell_type":"code","source":"print(f\"\\n... MIXED PRECISION SETUP STARTING ...\\n\")\nprint(\"\\n... SET TF TO OPERATE IN MIXED PRECISION – `bfloat16` – IF ON TPU ...\")\n\n# Set Mixed Precision Global Policy\n# ---> To use mixed precision in Keras, you need to create a `tf.keras.mixed_precision.Policy`\n# typically referred to as a dtype policy. \n# ---> Dtype policies specify the dtypes layers will run in\ntf.keras.mixed_precision.set_global_policy('float32')\n\n# target data type, bfloat16 when using TPU to improve throughput\nTARGET_DTYPE = tf.float32\nprint(f\"\\t--> THE TARGET DTYPE HAS BEEN SET TO {TARGET_DTYPE} ...\")\n\n# The policy specifies two important aspects of a layer: \n# 1. The dtype the layer's computations are done in\n# 2. The dtype of a layer's variables. \nprint(f\"\\n... TWO IMPORTANT ASPECTS OF THE GLOBAL MIXED PRECISION POLICY:\")\nprint(f'\\t--> COMPUTE DTYPE : {tf.keras.mixed_precision.global_policy().compute_dtype}')\nprint(f'\\t--> VARIABLE DTYPE : {tf.keras.mixed_precision.global_policy().variable_dtype}')\n\nprint(f\"\\n\\n... MIXED PRECISION SETUP COMPLTED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:31:11.084522Z","iopub.execute_input":"2021-06-03T23:31:11.084970Z","iopub.status.idle":"2021-06-03T23:31:11.094678Z","shell.execute_reply.started":"2021-06-03T23:31:11.084929Z","shell.execute_reply":"2021-06-03T23:31:11.093705Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

2.4 LEVERAGING XLA OPTIMIZATIONS

\n\n---\n\n\n**XLA** (Accelerated Linear Algebra) is a domain-specific compiler for linear algebra that can accelerate TensorFlow models with potentially no source code changes. **The results are improvements in speed and memory usage**.\n\n
\n\nWhen a TensorFlow program is run, all of the operations are executed individually by the TensorFlow executor. Each TensorFlow operation has a precompiled GPU/TPU kernel implementation that the executor dispatches to.\n\nXLA provides us with an alternative mode of running models: it compiles the TensorFlow graph into a sequence of computation kernels generated specifically for the given model. Because these kernels are unique to the model, they can exploit model-specific information for optimization.

\n\n
\n 🛑   WARNING:

- XLA can not currently compile functions where dimensions are not inferrable: that is, if it's not possible to infer the dimensions of all tensors without running the entire computation
\n
\n\n
\n 📌   NOTE:

- XLA compilation is only applied to code that is compiled into a graph (in TF2 that's only a code inside tf.function).
- The jit_compile API has must-compile semantics, i.e. either the entire function is compiled with XLA, or an errors.InvalidArgumentError exception is thrown)\n
\n\n
\n 📖   REFERENCE:

- XLA: Optimizing Compiler for Machine Learning
\n
","metadata":{}},{"cell_type":"code","source":"print(f\"\\n... XLA OPTIMIZATIONS STARTING ...\\n\")\n\nprint(f\"\\n... CONFIGURE JIT (JUST IN TIME) COMPILATION ...\\n\")\n# enable XLA optmizations (10% speedup when using @tf.function calls)\ntf.config.optimizer.set_jit(True)\n\nprint(f\"\\n... XLA OPTIMIZATIONS COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:31:11.096250Z","iopub.execute_input":"2021-06-03T23:31:11.096636Z","iopub.status.idle":"2021-06-03T23:31:11.107348Z","shell.execute_reply.started":"2021-06-03T23:31:11.096603Z","shell.execute_reply":"2021-06-03T23:31:11.105974Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

2.5 BASIC DATA DEFINITIONS & INITIALIZATIONS

\n\n---\n","metadata":{}},{"cell_type":"code","source":"# which_to_rotate = pd.read_csv(\"../input/bms-cnn-attn-lstm-tpu-tutorial-inference/sample_submission_w_extra_w_rotation.csv\")\n# pred_df = pd.read_csv(\"../input/bms-cnn-attn-lstm-tpu-tutorial-inference/prediction_confidences_for_rotation_analysis.csv\")\n# ss_df = pd.read_csv(\"../input/bms-molecular-translation/sample_submission.csv\")\n\n# pred_df[\"needs_rotation\"] = which_to_rotate[\"needs_rotation\"]\n# pred_df[\"best_inchi\"] = pred_df.progress_apply(lambda row: (row[\"w_rotated_inchi\"] if row[\"needs_rotation\"] else row[\"non_rotated_inchi\"]), axis=1)\n# ss_df.InChI=pred_df.best_inchi\n# ss_df.to_csv(\"submission.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:31:11.108966Z","iopub.execute_input":"2021-06-03T23:31:11.109643Z","iopub.status.idle":"2021-06-03T23:31:11.123905Z","shell.execute_reply.started":"2021-06-03T23:31:11.109599Z","shell.execute_reply":"2021-06-03T23:31:11.122764Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... BASIC DATA SETUP STARTING ...\\n\")\n\n# All the possible tokens in our InChI 'language'\nTOKEN_LIST = [\"\", \"InChI=1S/\", \"\", \"/c\", \"/h\", \"/m\", \"/t\", \"/b\", \"/s\", \"/i\"] +\\\n ['Si', 'Br', 'Cl', 'F', 'I', 'N', 'O', 'P', 'S', 'C', 'H', 'B', ] +\\\n [str(i) for i in range(167,-1,-1)] +\\\n [\"\\+\", \"\\(\", \"\\)\", \"\\-\", \",\", \"D\", \"T\"]\nprint(f\"\\n... TOKEN LIST:\")\nfor i, tok in enumerate(TOKEN_LIST): print(f\"\\t--> INTEGER-IDX = {i:<3} ––– STRING = {tok}\")\n\n# The start/end/pad tokens will be removed from the string when computing the Levenshtein distance\n# We want them as tf.constant's so they will operate properly within the @tf.function context\nSTART_TOKEN = tf.constant(TOKEN_LIST.index(\"InChI=1S/\"), dtype=tf.uint8)\nEND_TOKEN = tf.constant(TOKEN_LIST.index(\"\"), dtype=tf.uint8)\nPAD_TOKEN = tf.constant(TOKEN_LIST.index(\"\"), dtype=tf.uint8)\n\n# Prefixes and Their Respective Ordering/Format\n# -- ORDERING --> {c}{h/None}{b/None}{t/None}{m/None}{s/None}{i/None}{h/None}{t/None}{m/None}\nPREFIX_ORDERING = \"chbtmsihtm\"\nprint(f\"\\n... PREFIX ORDERING IS {PREFIX_ORDERING} ...\")\n\n# Paths to Respective Image Directories\nTEST_DIR = os.path.join(DATA_DIR, \"test_records\")\n\n# Get the Full Paths to The Individual TFRecord Files\nTEST_TFREC_PATHS = sorted(\n tf.io.gfile.glob(os.path.join(TEST_DIR, \"*.tfrec\")), \n key=lambda x: int(x.rsplit(\"_\", 2)[1]))\n\nprint(f\"\\n... TFRECORD INFORMATION:\")\nfor SPLIT, TFREC_PATHS in zip([\"TEST\",], [TEST_TFREC_PATHS,]):\n print(f\"\\t--> {len(TFREC_PATHS):<3} {SPLIT:<5} TFRECORDS\")\n\n# Paths to relevant CSV files containing training and submission information\n# SS_CSV_PATH = os.path.join(\"/kaggle/input\", \"bms-csv-files-with-extra-metadata\", \"sample_submission_w_extra.csv\")\nSS_CSV_PATH = \"../input/bms-molecular-translation/sample_submission.csv\"\nprint(f\"\\n... PATHS TO CSVS:\")\nprint(f\"\\t--> SS CSV : {SS_CSV_PATH}\")\n\n# When debug is true we use a smaller batch size and smaller model\nDEBUG=False\n\nprint(\"\\n\\n... BASIC DATA SETUP COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:31:11.125452Z","iopub.execute_input":"2021-06-03T23:31:11.126054Z","iopub.status.idle":"2021-06-03T23:31:11.232492Z","shell.execute_reply.started":"2021-06-03T23:31:11.126012Z","shell.execute_reply":"2021-06-03T23:31:11.231193Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

2.6 INITIAL DATAFRAME INSTANTIATION

\n\n---\n","metadata":{}},{"cell_type":"code","source":"print(\"\\n... INITIAL DATAFRAME INSTANTIATION STARTING ...\\n\")\n\n# Load the submission dataframe\nss_df = pd.read_csv(SS_CSV_PATH)\n\n# --- Distribution Information ---\nN_TEST = len(ss_df)\n\n# --- Batching Information ---\nBATCH_SIZE_DEBUG = 2\nREPLICA_BATCH_SIZE = 64\n\nif DEBUG:\n REPLICA_BATCH_SIZE = BATCH_SIZE_DEBUG\nOVERALL_BATCH_SIZE = REPLICA_BATCH_SIZE*N_REPLICAS\n\n# --- Required Amount of Padding To Ensure Fixed Batch Sizes (no dropping) ---\nREQUIRED_DATASET_PAD = OVERALL_BATCH_SIZE-N_TEST%OVERALL_BATCH_SIZE\n\n# --- Input Image Information ---\nIMG_SHAPE = (192,384,3)\n\n# --- Autocalculate Training/Validation Information ---\nTEST_STEPS = int(np.ceil(N_TEST/OVERALL_BATCH_SIZE))\n\n# --- Modelling Information ---\nATTN_EMB_DIM = 256\nN_RNN_UNITS = 768\n\nprint(f\"\\n... # OF TEST EXAMPLES : {N_TEST:<7} ...\\n\")\n\nprint(f\"\\n... REPLICA BATCH SIZE : {REPLICA_BATCH_SIZE} ...\")\nprint(f\"... OVERALL BATCH SIZE : {OVERALL_BATCH_SIZE} ...\\n\")\n\nprint(f\"\\n... AMOUNT OF PADDING TO MAKE FULL BATCHES : {REQUIRED_DATASET_PAD} ...\\n\")\n\nprint(f\"\\n... IMAGE SHAPE : {IMG_SHAPE} ...\\n\")\n\nprint(f\"\\n... TEST STEPS PER EPOCH : {TEST_STEPS:<5} ...\\n\")\n\nprint(f\"\\n... ATTENTION EMBEDDING DIMENSION : {ATTN_EMB_DIM:<5} ...\")\nprint(f\"... NUMBER OF UNITS IN LSTM : {N_RNN_UNITS:<5} ...\\n\")\n\nprint(\"\\n... SUBMISSION DATAFRAME ...\\n\")\ndisplay(ss_df.head(3))\n\nprint(\"\\n... INITIAL DATAFRAME INSTANTIATION COMPLETED...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:31:13.225052Z","iopub.execute_input":"2021-06-03T23:31:13.225410Z","iopub.status.idle":"2021-06-03T23:31:15.084684Z","shell.execute_reply.started":"2021-06-03T23:31:13.225382Z","shell.execute_reply":"2021-06-03T23:31:15.083222Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

2.7 USER INPUT VARIABLES

\n\n---\n","metadata":{}},{"cell_type":"code","source":"print(\"\\n... SPECIAL VARIABLE SETUP STARTING ...\\n\")\n\nENCODER_CKPT_PATH = \"/kaggle/input/new-bms-image-captioning-train-tpu/encoder_epoch_safety_save.h5\"\nDECODER_CKPT_PATH = \"/kaggle/input/new-bms-image-captioning-train-tpu/decoder_epoch_safety_save.h5\"\n\n# Paths to model weights to use in inference\nprint(f\"\\n... ENCODER MODEL INFERENCE WILL USE CHECKPOINT:\\n\\t-->{ENCODER_CKPT_PATH}\\n\")\nprint(f\"... DECODER MODEL INFERENCE WILL USE CHECKPOINT:\\n\\t-->{DECODER_CKPT_PATH}\\n\")\n \nprint(\"\\n... SPECIAL VARIABLE SETUP COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:32:25.514150Z","iopub.execute_input":"2021-06-03T23:32:25.514664Z","iopub.status.idle":"2021-06-03T23:32:25.520763Z","shell.execute_reply.started":"2021-06-03T23:32:25.514631Z","shell.execute_reply":"2021-06-03T23:32:25.519604Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n\n

\n 3  HELPER FUNCTION & CLASSESS    ⤒\n

","metadata":{}},{"cell_type":"markdown","source":"

3.1 GENERAL HELPER FUNCTIONS

\n\n---","metadata":{}},{"cell_type":"code","source":"def flatten_l_o_l(nested_list):\n \"\"\" Flatten a list of lists \"\"\"\n return [item for sublist in nested_list for item in sublist]\n\n\ndef tf_load_image(path, img_size=(192,384,3), invert=False):\n \"\"\" Load an image with the correct size and shape \n \n Args:\n path (tf.string): Path to the image to be loaded\n img_size (tuple, optional): Size to reshape image to (required for TPU)\n tile_to_3_channel (bool, optional): Whether to tile the single channel\n image to 3 channels which will be required for most off-the-shelf models\n invert (bool, optional): Whether or not to invert the background/foreground\n \n Returns:\n 3 channel tf.Constant image ready for training/inference\n \n \"\"\"\n img = decode_img(tf.io.read_file(path), img_size, n_channels=3, invert=invert) \n return img\n\n \ndef pad_to_square(a, constant=255):\n \"\"\" Pad a tensor array `a` evenly until it is a square \n \n Args:\n a (tf.constant array): Array to be padded to become square (3D)\n constant (int, optional): Value to use in padding\n \n Returns:\n square constant padded array\n \"\"\"\n h_src = tf.shape(a)[0]\n w_src = tf.shape(a)[1] \n if w_src>h_src: # pad height\n n_to_add = w_src-h_src\n top_pad = n_to_add//2\n bottom_pad = n_to_add-top_pad\n a = tf.pad(a, [(top_pad, bottom_pad), (0, 0), (0, 0)], mode='constant', constant_values=constant)\n elif h_src>w_src: # pad width\n n_to_add = h_src-w_src\n left_pad = n_to_add//2\n right_pad = n_to_add-left_pad\n a = tf.pad(a, [(0, 0), (left_pad, right_pad), (0, 0)], mode='constant', constant_values=constant)\n else:\n pass\n return a\n \n \ndef decode_image(image_data, resize_to=(384,384,3)):\n \"\"\" Function to decode the tf.string containing image information \n \n \n Args:\n image_data (tf.string): String containing encoded image data from tf.Example\n resize_to (tuple, optional): Size that we will reshape the tensor to (required for TPU)\n \n Returns:\n Tensor containing the resized single-channel image in the appropriate dtype\n \"\"\"\n image = tf.image.decode_png(image_data, channels=3)\n image = tf.reshape(image, resize_to)\n return tf.cast(image, TARGET_DTYPE)\n\n\ndef tokens_to_str(caption_tokens, discard_padding=True):\n \"\"\" Should convert a string of token ids to an InChI string \"\"\"\n if discard_padding:\n return \"\".join([int_2_char_lex[x] for x in caption_tokens if x!=len(int_2_char_lex)])\n else:\n return \"\".join([int_2_char_lex[x] for x in caption_tokens])\n \n \ndef evaluate(image, from_np=False):\n \"\"\" TBD \"\"\"\n attention_plot = np.zeros((MAX_LEN, fixed_encoder.output_shape[1]))\n hidden = tf.zeros((1, RNN_UNITS), tf.float32)\n\n if not from_np:\n temp_input = tf.expand_dims(tf_load_image(image, img_size=INPUT_SHAPE[:-1]), 0)\n img_tensor_val = fixed_encoder(temp_input)\n else:\n img_tensor_val=image\n \n features = trainable_encoder(img_tensor_val)\n dec = tf.ones((1, 1), tf.uint8)\n result = [int_2_char_lex[1],]\n\n for i in range(MAX_LEN-1):\n predictions, hidden, attention_weights = dec_model([dec, hidden, features])\n attention_plot[i] = tf.reshape(attention_weights, (-1, )).numpy()\n predicted_id = tf.random.categorical(predictions, 1)[0][0].numpy()\n result.append(int_2_char_lex[predicted_id])\n if int_2_char_lex[predicted_id] == '':\n return result, attention_plot\n\n dec = tf.expand_dims([predicted_id], 0)\n\n attention_plot = attention_plot[:len(result), :]\n return result, attention_plot\n\n\ndef plot_attention(image, result, attention_plot):\n \"\"\" TBD \"\"\" \n temp_image = np.array(Image.open(image))\n\n fig = plt.figure(figsize=(18, 14))\n\n len_result = len(result)\n for l in range(len_result):\n temp_att = np.resize(attention_plot[l], (8, 8))\n ax = fig.add_subplot(len_result//2, len_result//2, l+1)\n ax.set_title(result[l])\n img = ax.imshow(temp_image)\n ax.imshow(temp_att, cmap='gray', alpha=0.4, extent=img.get_extent())\n\n plt.tight_layout()\n plt.show()\n \ndef test_random_image(style=\"full\"):\n \"\"\" TBD \"\"\" \n rid = np.random.randint(0, len(val_subset_df))\n path = val_subset_df[\"img_path\"][rid]\n \n if style==\"full\":\n real_caption = val_subset_df[\"InChI\"][rid][:-1]\n else:\n real_caption = val_subset_df[\"InChI_chem\"][rid][:-1]\n\n result, attention_plot = evaluate(path)\n result = ''.join(result[:-1])\n print (f\"\\n\\tReal Caption : {real_caption}\")\n print (f\"\\tPrediction Caption : {result}\")\n print(f\"\\tLevenshtein Distance : {Levenshtein.distance(real_caption, result)}\\n\")\n plot_attention(path, result, attention_plot)\n \n \n# sparse tensors are required to compute the Levenshtein distance\ndef dense_to_sparse(dense):\n \"\"\" Convert a dense tensor to a sparse tensor \n \n Args:\n dense (Tensor): TBD\n \n Returns:\n A sparse tensor \n \"\"\"\n indices = tf.where(tf.ones_like(dense))\n values = tf.reshape(dense, (MAX_LEN*OVERALL_BATCH_SIZE,))\n sparse = tf.SparseTensor(indices, values, dense.shape)\n return sparse\n\ndef get_levenshtein_distance(preds, lbls):\n \"\"\" Computes the Levenshtein distance between the predictions and labels \n \n Args:\n preds (tensor): Batch of predictions\n lbls (tensor): Batch of labels\n \n Returns:\n The mean Levenshtein distance calculated across the batch\n \"\"\"\n preds = tf.where(tf.not_equal(lbls, END_TOKEN) & tf.not_equal(lbls, PAD_TOKEN), preds, 0)\n lbls = tf.where(tf.not_equal(lbls, END_TOKEN), lbls, 0)\n\n preds_sparse = dense_to_sparse(preds)\n lbls_sparse = dense_to_sparse(lbls)\n\n batch_distance = tf.edit_distance(preds_sparse, lbls_sparse, normalize=False)\n mean_distance = tf.math.reduce_mean(batch_distance)\n \n return mean_distance","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:32:31.083645Z","iopub.execute_input":"2021-06-03T23:32:31.084087Z","iopub.status.idle":"2021-06-03T23:32:31.116881Z","shell.execute_reply.started":"2021-06-03T23:32:31.084001Z","shell.execute_reply":"2021-06-03T23:32:31.115744Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n\n

4  PREPARE THE DATASET    ⤒

\n\nIn this section we prepare the **`tf.data.Datasets`** we will use for training and validation","metadata":{}},{"cell_type":"markdown","source":"

4.0 CREATE THE TOKENIZED INCHI LABELS (TO BE REMOVED LATER)

\n\n---\n\n\nThis step is optional and should be replaced with the InChI labels directly encoded into the TFRecord files. That being said, we will do it like so for now to allow for easy modification of the tokenization styles.","metadata":{}},{"cell_type":"code","source":"print(\"\\n... CREATING THE TOKEN MAPPINGS AND IMPORTANT CONSTANTS ...\\n\")\n\nprint(\"\\n... CREATE RELEVANT OBJECTS ... \")\ntok_2_int = {c.strip(\"\\\\\"):i for i,c in enumerate(TOKEN_LIST)}\nint_2_tok = {v:k for k,v in tok_2_int.items()}\n\nprint(\"... CREATE IMPORTANT CONSTANTS ...\\n\")\n# Max Length Was Determined Previously Using... \n# >>> MAX_LEN = train_df.InChI.progress_apply(lambda x: len(re.findall(\"|\".join(TOKEN_LIST), x))).max()+1\nMAX_LEN = 277\nVOCAB_LEN = len(int_2_tok)\n\nprint(\"\\n... FINISHED CREATING THE TOKEN MAPPINGS AND IMPORTANT CONSTANTS ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:32:35.082478Z","iopub.execute_input":"2021-06-03T23:32:35.082969Z","iopub.status.idle":"2021-06-03T23:32:35.090615Z","shell.execute_reply.started":"2021-06-03T23:32:35.082936Z","shell.execute_reply":"2021-06-03T23:32:35.089594Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

4.1 READ TFRECORD FILES - CREATE THE RAW DATASET(S)

\n\n---\n\nHere we will leverage **`tf.data.TFRecordDataset`** to read the TFRecord files.\n* The simplest way is to specify a list of filenames (paths) of TFRecord files.\n* It is a subclass of **`tf.data.Dataset`**.\n\nThis newly created raw dataset contains **`tf.train.Example`** messages, and when iterated over it, we get scalar string tensors.","metadata":{}},{"cell_type":"code","source":"print(\"\\n... CREATE TFRECORD RAW DATASETS STARTING ...\\n\")\n\n# Create tf.data.Dataset from filepaths for conversion later\nraw_test_ds = tf.data.TFRecordDataset(TEST_TFREC_PATHS, num_parallel_reads=None)\n\nprint(f\"\\n... THE RAW TF.DATA.TFRECORDDATASET OBJECT:\\n\\t--> {raw_test_ds}\\n\")\n\nprint(\"\\n... CREATE TFRECORD RAW DATASETS COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:32:38.502975Z","iopub.execute_input":"2021-06-03T23:32:38.503592Z","iopub.status.idle":"2021-06-03T23:32:38.523008Z","shell.execute_reply.started":"2021-06-03T23:32:38.503556Z","shell.execute_reply":"2021-06-03T23:32:38.521723Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

4.2 PARSE THE RAW DATASET(S)

\n\n---\n\n\nThe general recipe to parse the string tensors in the raw dataset looks something like this:\n\n
\n\n**STEP 1.** Create a description of the features. For example:\n\n```python\nfeature_description = { \n 'feature0': tf.io.FixedLenFeature([], tf.int64),\n 'feature1': tf.io.FixedLenFeature([], tf.string),\n 'feature2': tf.io.FixedLenFeature([], tf.float32),\n ...\n}\n```\n\n
\n\n**STEP 2.** Define a parsing function by using `tf.io.parse_single_example` and the defined feature description.\n```python\ndef _parse_function(example):\n \"\"\"\n Args:\n example: A string tensor representing a `tf.train.Example`.\n \"\"\"\n\n # Parse `example`.\n parsed_example = tf.io.parse_single_example(example, feature_description)\n\n return parsed_example\n```\n\n
\n\n**STEP 3.** Map the raw dataset by `_parse_function`.\n```python\ndataset = raw_dataset.map(_parse_function)\n```\n\n
\n\n---\n\n
\n\n**In the following cell, we apply the above recipe to our BMS tfrecord dataset.**\n\n\n
\n 📌   NOTE:

- The parsed images are `tf.string`, which are then decoded with `tf.image.decode_png` which is an alias for `tf.io.decode_png`
- The InChI strings and Image IDs will just be left as byte string tensors.\n
\n\n
\n 📖   REFERENCE:

\n - Tutorial - TFRecord and tf.Example
\n - TFRecordDataset Documentation
\n - Decoding PNGs Documentation
\n
\n","metadata":{}},{"cell_type":"code","source":"def decode(serialized_example, is_test=False, tokenized_inchi=True):\n \"\"\" Parses a set of features and label from the given `serialized_example`.\n \n It is used as a map function for `dataset.map`\n\n Args:\n serialized_example (tf.Example): A serialized example containing the\n following features:\n – 'image'\n – 'image_id'\n – 'inchi'\n is_test (bool, optional): Whether to allow for the InChI feature\n drop_id (bool, optional): Whether or not to drop the ID feature\n \n Returns:\n A decoded tf.data.Dataset object representing the tfrecord dataset\n \"\"\"\n feature_dict = {\n 'image': tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value=''),\n }\n \n if not is_test:\n if tokenized_inchi:\n feature_dict[\"inchi\"] = tf.io.FixedLenFeature(shape=[MAX_LEN], dtype=tf.int64, default_value=[0]*MAX_LEN)\n else:\n feature_dict[\"inchi\"] = tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value='')\n else:\n feature_dict[\"image_id\"] = tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value='')\n \n # Define a parser\n features = tf.io.parse_single_example(serialized_example, features=feature_dict)\n \n # Decode the tf.string\n image = decode_image(features['image'], resize_to=IMG_SHAPE)\n \n # Figure out the correct information to return\n if is_test:\n image_id = features[\"image_id\"] \n return image, image_id\n else:\n if tokenized_inchi:\n target = tf.cast(features[\"inchi\"], tf.uint8)\n else:\n target = features[\"inchi\"]\n return image, target","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:32:39.987409Z","iopub.execute_input":"2021-06-03T23:32:39.987764Z","iopub.status.idle":"2021-06-03T23:32:39.999647Z","shell.execute_reply.started":"2021-06-03T23:32:39.987729Z","shell.execute_reply":"2021-06-03T23:32:39.998536Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... DECODING RAW TFRECORD DATASETS STARTING ...\\n\")\n\n# Decode the tfrecords completely –– decode is our `_parse_function` (from recipe above)\ntest_ds = raw_test_ds.map(lambda x: decode(x, is_test=True))\nprint(f\"\\n... THE DECODED TF.DATA.TFRECORDDATASET OBJECT:\" \\\n f\"\\n\\t--> ((image, image_id))\" \\\n f\"\\n\\t--> {test_ds}\\n\")\n\nprint(\"\\n... 2 EXAMPLES OF IMAGES AND LABELS AFTER DECODING ...\")\nfor i, (img, image_id) in enumerate(test_ds.take(2)):\n print(f\"\\nIMAGE SHAPE : {img.shape}\")\n print(f\"\\nIMAGE ID : {image_id.numpy().decode()}\")\n\n plt.figure(figsize=(10,10))\n plt.imshow(img.numpy().astype(np.int64), cmap=\"gray\")\n plt.title(f\"{image_id.numpy().decode()}\")\n plt.show()\n\nprint(\"\\n... DECODING RAW TFRECORD DATASETS COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:32:41.259480Z","iopub.execute_input":"2021-06-03T23:32:41.259817Z","iopub.status.idle":"2021-06-03T23:32:44.217599Z","shell.execute_reply.started":"2021-06-03T23:32:41.259788Z","shell.execute_reply":"2021-06-03T23:32:44.216339Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

4.4 WORKING WITH `TF.DATA.DATASET` OBJECTS

\n\n---\n\nWith the above parsing methods defined, we can define how to load the dataset with more options and further apply shuffling, bacthing, etc. In particular the following methods and attributes are of special interest to us:\n* Use **`num_parallel_reads`** in **`tf.data.TFRecordDataset`** to read files in parallel.\n* Set **`tf.data.Options.experimental_deterministic=False`** and use it to get a new dataset that ignores the order of elements.\n* Use **`num_parallel_calls`** in **`tf.data.Dataset.map()`** method to have parallel processing.\n* Use **`tf.data.Dataset.prefetch()`** to allow later batches to be prepared while the current batch is being processed.\n* Use **`tf.data.AUTOTUNE`** to automatically determine parallelization argument values\n\nThe parallel processing and prefetching are particular important when working with TPU:\n* This is because a TPU can process batches very quickly\n* The dataset pipeline should be able to provide data for TPU efficiently, otherwise the TPU will be idle.\n\n**In the cell below we will create the functions and configuration template which will later be used to create our respective datasets**\n\n
\n 📖   REFERENCE:

\n - Guide - tf.data: Build TensorFlow Input Pipelines
\n - Guide - Better Performance With the tf.data API
\n - tf.data.Dataset Documentation
\n
","metadata":{}},{"cell_type":"code","source":"def load_dataset(filenames, is_test=False, ordered=False, tokenized_inchi=True):\n \"\"\"Read from TFRecords.\n \n For optimal performance, reading from multiple files at once and disregarding data order (if `ordered=False`).\n - If pulling InChI from TFRecords than order does not matter since we will \n be shuffling the data anyway (for training dataset).\n \n Args:\n filenames (list of strings): List of paths to that point to the respective TFRecord files\n is_test (bool, optional): Whether or not to include the image ID or label in the returned dataset\n ordered (bool, optional): Whether to ensured ordered results or maximize parallelization\n tokenized_inchi (bool, optional): Whether our dataset includes the tokenized inchi or we will be \n creating it from the caption numpy array\n \n Returns:\n Decoded tf.data.Dataset object\n \"\"\"\n\n options = tf.data.Options()\n if not ordered:\n # disable order, increase speed\n options.experimental_deterministic = False\n N_PARALLEL=tf.data.AUTOTUNE\n else:\n N_PARALLEL=None\n \n # If not-ordered, this will read in by automatically interleaving multiple tfrecord files.\n dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=N_PARALLEL)\n \n # If not-ordered, this will ensure that we use data as soon as it \n # streams in, rather than in its original order.\n dataset = dataset.with_options(options) \n \n # parse and return a dataset w/ the appropriate configuration\n dataset = dataset.map(\n lambda x: decode(x, is_test, tokenized_inchi),\n num_parallel_calls=N_PARALLEL,\n )\n \n return dataset\n\ndef get_dataset(filenames, batch_size, \n is_test=False, \n shuffle_buffer_size=1, \n repeat_dataset=True, \n preserve_file_order=False, \n drop_remainder=True,\n tokenized_inchi=True,\n external_inchi_dataset=None,\n test_padding=0):\n \"\"\" Get a tf.data.Dataset w/ the appropriate configuration\n \n Args:\n TBD\n test_padding (int, optional): Amount required to pad dataset to have only full batches\n \n Returns:\n TBD\n \n \"\"\"\n # Load the dataset\n dataset = load_dataset(filenames, is_test, preserve_file_order, tokenized_inchi)\n \n if test_padding!=0:\n pad_dataset = tf.data.Dataset.from_tensor_slices((\n tf.zeros((test_padding, *IMG_SHAPE), dtype=TARGET_DTYPE), # Fake Images\n tf.constant([\"000000000000\",]*test_padding, dtype=tf.string)) # Fake IDs\n )\n dataset = dataset.concatenate(pad_dataset)\n \n # If we are training than we will want to repeat the dataset. \n # We will determine the number of steps (or updates) later for 1 training epoch.\n if repeat_dataset:\n dataset = dataset.repeat()\n \n # If we need to add on manually the inchi\n if external_inchi_dataset is not None:\n # Zip the datasets and tile the 1 channel image to 3 channels & drop the old inchi value\n dataset = tf.data.Dataset.zip((dataset, external_inchi_dataset))\n dataset = dataset.map(lambda x,y: (tf.tile(tf.expand_dims(x[0], -1), tf.constant([1,1,3], tf.int32)), y))\n \n # Shuffling\n if shuffle_buffer_size!=1:\n dataset = dataset.shuffle(shuffle_buffer_size)\n \n # Batching\n dataset = dataset.batch(batch_size, drop_remainder=drop_remainder)\n \n # prefetch next batch while training (autotune prefetch buffer size)\n dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)\n \n return dataset","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:32:44.219143Z","iopub.execute_input":"2021-06-03T23:32:44.219449Z","iopub.status.idle":"2021-06-03T23:32:44.233836Z","shell.execute_reply.started":"2021-06-03T23:32:44.219421Z","shell.execute_reply":"2021-06-03T23:32:44.232666Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"####### ####### ####### ####### ####### ####### ####### #######\n\n# Template Configuration\nDS_TEMPLATE_CONFIG = dict(\n filenames=[],\n batch_size=1,\n is_test=False, \n shuffle_buffer_size=1, \n repeat_dataset=True, \n preserve_file_order=False, \n drop_remainder=True,\n tokenized_inchi=True,\n external_inchi_dataset=None,\n test_padding=0\n)\n\n####### ####### ####### ####### ####### ####### ####### #######\n\nTEST_DS_CONFIG = DS_TEMPLATE_CONFIG.copy()\nTEST_DS_CONFIG.update(dict(\n filenames=TEST_TFREC_PATHS,\n batch_size=OVERALL_BATCH_SIZE,\n is_test=True,\n repeat_dataset=False,\n drop_remainder=True,\n test_padding=REQUIRED_DATASET_PAD,\n))\n\n####### ####### ####### ####### ####### ####### ####### #######\n\ntest_ds = get_dataset(**TEST_DS_CONFIG)\n\nfor SPLIT, CONFIG in zip([\"TESTING\",], [TEST_DS_CONFIG]):\n print(f\"\\n... {SPLIT} CONFIGURATION:\")\n for k,v in CONFIG.items():\n if k==\"filenames\":\n print(f\"\\t--> {k:<23}: {[path.split('/', 4)[-1] for path in v[:2]]+['...']}\")\n else:\n print(f\"\\t--> {k:<23}: {v}\")\n\nprint(f\"... TESTING DATASET : {test_ds} ...\\n\")\n\nprint(\"\\n\\n ... SOME EXAMPLES ... \\n\\n\")\nfor x,y in test_ds.take(1):\n for i in range(REPLICA_BATCH_SIZE-3, REPLICA_BATCH_SIZE, 1):\n plt.figure(figsize=(12,12))\n plt.imshow(x[i].numpy().astype(np.int64))\n plt.title(f\"{y[i].numpy().decode()}\")\n plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:32:44.236198Z","iopub.execute_input":"2021-06-03T23:32:44.236631Z","iopub.status.idle":"2021-06-03T23:32:50.954372Z","shell.execute_reply.started":"2021-06-03T23:32:44.236587Z","shell.execute_reply":"2021-06-03T23:32:50.953178Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n\n

5  MODEL PREPERATION    ⤒

\n\nIn this section we prepare the models for training. We will be using a model architecture very similar to that found within the [**Show, Attend, and Tell Research Paper**](https://arxiv.org/pdf/1502.03044.pdf).\n\n
\n\n
\n \n
","metadata":{}},{"cell_type":"markdown","source":"

5.1 UNDERSTANDING THE MODELS - ENCODER

\n\n---\n\nWe will be leveraging an [**EfficientNet**](https://arxiv.org/pdf/1905.11946) model to act as the Encoder CNN in our network. Ideally I will update this to be an [**NFNet**](https://arxiv.org/pdf/2102.06171) or an [**EfficientNetV2**](https://arxiv.org/pdf/2104.00298) model in the near future... or perhaps [**ViT**](https://arxiv.org/pdf/2010.11929)??\n* On **TPU** we will use an **EfficientNetB3** model *(In my private training notebook I leverage EB7)*\n* On **GPU/CPU** we will use an **EfficientNetB0** model\n\n
\n\n***Basic View of EfficientNetB0 Architecture w/ 380x380x3 Input***\n
\n\n
\n\nOur encoder will create feature maps for each image which will in turn be passed to the decoder side of the network. EfficientNetB0 creates $1280$ feature maps with dimensions of $12\\cdot12$ pixels. These feature maps will be flattened by a reshape: $12\\cdot12 \\Rightarrow 144$. \n\n
\n\n**In the following cell, we will create a function to generate our encoder model.**\n\n\n
\n 📌   NOTE:

- For different encoder architectures we will have a different number of feature maps. i.e. If we utilized EfficientNetB3 we would have 1536 feature maps instead of the 1280 feature maps that EfficientNetB0 produces. \n
\n\n
\n 📖   REFERENCE:

\n - TF Tutorial – Transformer Model for Language Understanding
\n - TF Tutorial – Image Captioning
\n - TF Tutorial – Neural Machine Translation w/ Attention
\n
","metadata":{}},{"cell_type":"code","source":"print(\"\\n... ENCODER MODEL CREATION STARTING ...\\n\")\n\n# SAMPLE IMAGES\nSAMPLE_IMGS, SAMPLE_IDS = next(iter(test_ds.unbatch().batch(BATCH_SIZE_DEBUG)))\n\n# ENCODER_CONFIG\nUSE_QUBVEL = False\nPRETRAINED_WTS = None # or `noisy-student` is possible if using qubvel EfficientNet or `None`\nif TPU and not DEBUG:\n if USE_QUBVEL:\n BB_FN = efn.EfficientNetB5\n else:\n BB_FN = tf.keras.applications.EfficientNetB5\nelse:\n if USE_QUBVEL:\n BB_FN = efn.EfficientNetB5\n else:\n BB_FN = tf.keras.applications.EfficientNetB5\nPREPROCESSING_FN = tf.keras.applications.efficientnet.preprocess_input\n \n# This will be the dimension the network outputs flattened\nIMG_EMB_DIM = BB_FN(include_top=False, input_shape=IMG_SHAPE).output_shape[1:]\nIMG_EMB_DIM = (IMG_EMB_DIM[0]*IMG_EMB_DIM[1], IMG_EMB_DIM[2])\n\nclass Encoder(tf.keras.Model):\n def __init__(self, image_embedding_dim, preprocessing_fn, backbone_fn, image_shape, do_permute=False, pretrained_weights=\"imagenet\"):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n super(Encoder, self).__init__()\n \n self.image_embedding_dim = image_embedding_dim\n self.preprocessing_fn = preprocessing_fn\n self.encoder_backbone = backbone_fn(include_top=False, weights=None, input_shape=image_shape) \n self.reshape = tf.keras.layers.Reshape(self.image_embedding_dim, name='image_embedding')\n \n def call(self, x, training):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n x = self.preprocessing_fn(x)\n x = self.encoder_backbone(x, training=training)\n x = self.reshape(x, training=training)\n return x\n \n \n# Example enoder output\nwith tf.device('/CPU:0'):\n encoder = Encoder(IMG_EMB_DIM, PREPROCESSING_FN, BB_FN, IMG_SHAPE, do_permute=IMG_EMB_DIM[1]5.2 UNDERSTANDING THE MODELS - DECODER\n\n---\n\nWe will be leveraging a decoder nearly identical to that found in the [**Show, Attend, and Tell Research Paper**](https://arxiv.org/pdf/1502.03044). This paper introduces an archiecture for image captioning that incorporates attention (Soft and Hard Attention). We will be utilizing Bahdanau Attention (Soft Attention) within this notebook. The core components of our decoder will be:\n\n* **Attention Mechanism [(Partly Quoting From Mark Wijkhuizen's Notebook)](https://www.kaggle.com/markwijkhuizen/tensorflow-tpu-training-baseline-lb-16-92)**\n * The attention mechanism takes as input the hidden state from the LSTM, which is the LSTM state after the last predicted character, and encoder features. \n * The hidden LSTM state will differ each prediction iteration, but the encoder result remains the same. \n * Using this hidden LSTM state the attention mechanism learns which encoder features are important and which are not, given an LSTM state.\n * To make this idea of attention a bit less abstract, let us take the a random InChI string as an example...\n * **`C13H5F5N2/c14-7-3-6(5-19)1-2-10(7)20-13-11(17)8(15)4-9(16)12(13)18/h1-4,20H`**\n * Let's say the model has so far correctly predicted **`C13H5`**, the attention mechanism should now focus on features containing **`F`** atoms and disregard any feature maps that target **`C`** or **`H`** atoms. \n * The LSTM hidden state should tell the attention mechanism it has predicted **`C13H5`** so far and the attention mechanism will learn it has to focus on **`F`** atoms after **`C`** and **`H`** atoms are predicted.\n * The idea of 'focus' in this case refers to giving weight to certain parts of the feature maps where the attention should believes should be important.

\n\n* **Recurrent Mechanism – Embedding & LSTM/GRU [(Partly Quoting From Mark Wijkhuizen's Notebook)](https://www.kaggle.com/markwijkhuizen/tensorflow-tpu-training-baseline-lb-16-92)**\n * The Embedding layer will learn a representation to convert the chemical token/characters into a representative vector\n * The LSTMCell takes a concatenated context vector from the attention mechanism and an embedded token value as input. \n * The LSTMCell hidden and carry states are initialized using the encoder features.\n\n
\n\n***Basic Overview of the Image Captioning Process With Some Rudimentary Math***\n
\"attn_gif\"
\n\n
\n\n**In the following cell, we will create a function to generate our decoder model.**\n\n
\n 📌   NOTE:

- We decided to utilize an LSTM cell for the Recurrent portion of our decoder network. This decision is based on the concept that LSTMs can generally handle longer sequences better than GRUs. That being said, other options do exist. \n
\n\n
\n 📖   REFERENCE:

\n - TF Tutorial – Transformer Model for Language Understanding
\n - TF Tutorial – Image Captioning
\n - TF Tutorial – Neural Machine Translation w/ Attention
\n
","metadata":{}},{"cell_type":"code","source":"print(\"\\n... ATTENTION MECHANISM LAYER CREATION STARTING ...\\n\")\n\n# TO BE REPLACED W/ --> tfa.seq2seq.BahdanauAttention\nclass BahdanauAttention(tf.keras.layers.Layer):\n def __init__(self, attn_emb_depth):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n super(BahdanauAttention, self).__init__()\n self.W1 = tf.keras.layers.Dense(attn_emb_depth)\n self.W2 = tf.keras.layers.Dense(attn_emb_depth)\n self.V = tf.keras.layers.Dense(1)\n\n def call(self, hidden, features, training):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n # hidden shape == (batch_size, hidden size)\n # hidden_with_time_axis shape == (batch_size, 1, hidden size)\n # features shape == (batch_size, img_emb_size, hidden size)\n # we are doing this to broadcast addition along the time axis to calculate the score\n hidden_with_time_axis = tf.expand_dims(hidden, 1)\n\n # score shape == (batch_size, img_emb_size, 1)\n # we get 1 at the last axis because we are applying score to self.V\n # img_emb_size refers to the product of the image embedding width and height (144 for EB)\n # the shape of the tensor before applying self.V is (batch_size, img_emb_size, attention_emb_depth)\n score = self.V(\n tf.nn.tanh(self.W1(hidden_with_time_axis, training=training) + \\\n self.W2(features, training=training)), \n training=training\n )\n\n # attention_weights shape (before & after) == (batch_size, img_emb_size, 1)\n attention_weights = tf.nn.softmax(score, axis=1)\n\n # context_vector shape after multiplication is (batch_size, img_emb_size, 1)\n context_vector = attention_weights * features\n \n # context_vector shape after sum == (batch_size, hidden_size)\n context_vector = tf.reduce_sum(context_vector, axis=1) # this is axis 2 in some implementations?\n \n return context_vector, attention_weights\n \n\n# ATTN_EMB_DIM = 512\n# N_RNN_UNITS = 1024\nwith tf.device('/CPU:0'):\n attn_layer = BahdanauAttention(ATTN_EMB_DIM)\n context_vector, attn_weights = attn_layer(tf.zeros([BATCH_SIZE_DEBUG, IMG_EMB_DIM[0]]), img_embedding_batch)\n\nprint(f'\\n... Context Vector Shape : (batch_size, n_rnn_units) : {context_vector.shape} ...')\nprint(f'... Attention Weights Shape : (batch_size, img_emb_dim, 1) : {attn_weights.shape} ...\\n')\n\nprint(\"\\n... ATTENTION MECHANISM LAYER CREATION FINISHED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:33:03.631962Z","iopub.execute_input":"2021-06-03T23:33:03.632454Z","iopub.status.idle":"2021-06-03T23:33:03.678338Z","shell.execute_reply.started":"2021-06-03T23:33:03.632422Z","shell.execute_reply":"2021-06-03T23:33:03.677314Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... DECODER MODEL CREATION STARTING ...\\n\")\n\nclass Decoder(tf.keras.Model):\n def __init__(self, vocab_len, attn_emb_depth, img_emb_dim, n_rnn_units, dropout_rate=0.1):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n super(Decoder, self).__init__()\n \n # Basic Parameters\n self.vocab_len = vocab_len\n self.attn_emb_depth = attn_emb_depth\n self.img_emb_dim = img_emb_dim\n self.n_rnn_units = n_rnn_units\n self.dropout_rate = dropout_rate\n \n # Attention Mechanism\n self.attention_layer = BahdanauAttention(self.attn_emb_depth)\n \n # LSTM hidden and carry state initialization\n self.init_h_layer = tf.keras.layers.Dense(\n units=n_rnn_units, input_shape=[img_emb_dim], name='img_embedding__hidden_init_layer'\n )\n self.init_c_layer = tf.keras.layers.Dense(\n units=n_rnn_units, input_shape=[img_emb_dim], name='img_embedding__input_act_init_layer'\n )\n \n # Character Embedding Layer In The Decoder\n # - TODO ... `mask_zero=True` ???\n self.embedding_layer = tf.keras.layers.Embedding(vocab_len, n_rnn_units, )\n \n # The LSTM cell\n self.lstm_cell_layer = tf.keras.layers.LSTMCell(n_rnn_units, name='lstm_cell_layer')\n \n # Dropout Layer to Prevent Overfitting\n self.dropout_layer = tf.keras.layers.Dropout(self.dropout_rate, name='dropout_layer')\n \n # Fully Connected Prediction Layer \n self.fcn_layer = tf.keras.layers.Dense(units=vocab_len, input_shape=[n_rnn_units], dtype=tf.float32, name='fc_prediction_layer')\n\n def call(self, token, hidden, memory, img_emb, training):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n # img_emb shape == (batch_size, img_emb_shape, hidden_size)\n # --> We ignore the attention weights for now (_)\n context_vector, _ = self.attention_layer(hidden, img_emb, training=training)\n\n # shape after passing token through embedding == (batch_size, 1, n_rnn_units)\n # shape after passing through squeeze == (batch_size, n_rnn_units)\n tok_emb = tf.squeeze(self.embedding_layer(token, training=training), axis=1)\n \n # shape after concatenation == (batch_size, n_rnn_units + hidden_size)\n x = tf.concat((context_vector, tok_emb), axis=-1)\n\n # passing the concatenated vector to the LSTM cell\n # - also getting the new hidden (h) and memory (c) vectors\n _, (hidden_new, memory_new) = self.lstm_cell_layer(x, (hidden, memory), training=training)\n \n #### ##### ##### ####\n # MORE LSTM LAYERS? #\n #### ##### ##### ####\n \n # compute prediction logits and leverage dropout\n output = self.dropout_layer(hidden_new, training=training)\n output = self.fcn_layer(output, training=training)\n\n return output, hidden_new, memory_new\n \n def init_hidden_state(self, img_emb, training):\n mean_encoder_out = tf.math.reduce_mean(img_emb, axis=1)\n hidden = self.init_h_layer(mean_encoder_out, training=training) # (batch_size, n_rnn_units)\n memory = self.init_c_layer(mean_encoder_out, training=training)\n return hidden, memory\n\n\nwith tf.device('/CPU:0'):\n decoder = Decoder(VOCAB_LEN, ATTN_EMB_DIM, IMG_EMB_DIM[0], N_RNN_UNITS)\n h, c = decoder.init_hidden_state(img_embedding_batch[:BATCH_SIZE_DEBUG], training=False)\n pred_output, h, c = decoder(tf.random.uniform((BATCH_SIZE_DEBUG, 1)), h, c, img_embedding_batch)\n\nprint(f'\\n... Decoder Output Shape : (batch_size, vocab_len) : {pred_output.shape} ...\\n')\n\nprint(\"\\n... DECODER MODEL CREATION FINISHED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:33:03.679808Z","iopub.execute_input":"2021-06-03T23:33:03.680130Z","iopub.status.idle":"2021-06-03T23:33:04.673040Z","shell.execute_reply.started":"2021-06-03T23:33:03.680099Z","shell.execute_reply":"2021-06-03T23:33:04.671649Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

5.4 WRAP THE CONFIGURATION DETAILS IN A CLASS OBJECT FOR EASY ACCESS

\n\n---\n","metadata":{}},{"cell_type":"code","source":"class Config():\n def __init__(self,):\n self.encoder_config = {}\n self.decoder_config = {}\n def initialize_encoder_config(self, image_embedding_dim, preprocessing_fn, backbone_fn, image_shape, do_permute=False, pretrained_weights=\"imagenet\"):\n self.encoder_config = dict(\n image_embedding_dim=image_embedding_dim, \n preprocessing_fn=preprocessing_fn, \n backbone_fn=backbone_fn, \n image_shape=image_shape, \n do_permute=do_permute, \n pretrained_weights=pretrained_weights,\n )\n def initialize_decoder_config(self, vocab_len, attn_emb_depth, img_emb_dim, n_rnn_units, dropout_rate=0.1):\n self.decoder_config = dict(\n vocab_len=vocab_len, \n attn_emb_depth=attn_emb_depth, \n img_emb_dim=img_emb_dim, \n n_rnn_units=n_rnn_units, \n dropout_rate=dropout_rate,\n )\n \ninference_config = Config()\ninference_config.initialize_encoder_config(image_embedding_dim=IMG_EMB_DIM, \n preprocessing_fn=PREPROCESSING_FN, \n backbone_fn=BB_FN, \n image_shape=IMG_SHAPE, \n do_permute=IMG_EMB_DIM[1] {inference_config.encoder_config}\\n\")\nprint(f\"INFERENCE DECODER CONFIG:\\n\\t--> {inference_config.decoder_config}\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:33:04.677048Z","iopub.execute_input":"2021-06-03T23:33:04.677366Z","iopub.status.idle":"2021-06-03T23:33:04.688557Z","shell.execute_reply.started":"2021-06-03T23:33:04.677335Z","shell.execute_reply":"2021-06-03T23:33:04.687643Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

5.5 HOW TPU IMPACTS MODELS, METRICS, AND OPTIMIZERS

\n\nIn order to use TPU, or [**tensorflow distribute strategy**](https://www.tensorflow.org/api_docs/python/tf/distribute) in general, certain objects will have to be created inside the **strategy's scope**\n\n---\n\nHere is the rule of thumb:\n\n---\n\n* Anything that creates variables that will be used in a distributed way must be created inside **`strategy.scope()`**.\n* This includes, but is not limited to:\n - model creation\n - optimizer\n - metrics\n - sometimes, checkpoint restore\n - any custom code that creates distributed variables\n* Once a variable is created inside a strategy's scope, it captures the strategy's information, and **you can use it outside the strategy's scope.**\n* Unless using a high level API like **`model.fit()`**, defining something within the strategy's scope **WILL NOT automatically distribute the computation**. This will be discussed more in the section on training further down.\n\n---\n\nInside the scope, everything is defined in the same way it would be outside the distribution strategy. There is, however, a particularity about the loss function which we will discuss further down as well.\n\n**In the next cell, we instantiate the learning rate function, the loss object, and the model(s) inside the scope**\n\n
\n 📖   REFERENCE:

\n - TPUStrategy - Scope
\n - Tutorial - Custom Training With TPUs
\n
","metadata":{}},{"cell_type":"code","source":"print(\"\\n... INFERENCE PREPERATION STARTING ...\\n\")\n\ndef prepare_for_inference(encoder_config, decoder_config, encoder_wts=None, decoder_wts=None, verbose=0):\n \"\"\" Declare required objects under TPU session scope and return ready for training\n \n Args:\n encoder_config (dict): Keyword arguments mapped to desired values for encoder model instantiation\n decoder_config (dict): Keyword arguments mapped to desired values for decoder model instantiation \n encoder_wts (str, optional): Path to pretrained model weights for encoder\n decoder_wts (str, optional): Path to pretrained model weights for decoder\n verbose (bool, optional): Whether or not to print model information and plot lr schedule\n \n Returns:\n encoder - TBD\n decoder - TBD\n \n \"\"\"\n \n\n # Everything must be declared within the scope when leveraging the TPU strategy\n # - This will still function properly if scope is set to another type of accelerator\n with strategy.scope():\n # Instantiate the encoder model \n print(\"\\t--> CREATING & INITIALIZING ENCODER MODEL ...\")\n encoder = Encoder(**encoder_config)\n encoder.build(input_shape=[REPLICA_BATCH_SIZE, *IMG_SHAPE])\n initialization_batch = encoder(\n tf.ones((REPLICA_BATCH_SIZE,*IMG_SHAPE), dtype=TARGET_DTYPE), \n training=False,\n )\n \n # Instantiate the decoder model\n print(\"\\t--> CREATING & INITIALIZING DECODER MODEL...\")\n decoder = Decoder(**decoder_config)\n init_h, init_c = decoder.init_hidden_state(initialization_batch[:REPLICA_BATCH_SIZE], training=False)\n pred_output, h, c = decoder(\n tf.random.uniform((REPLICA_BATCH_SIZE, 1)), \n init_h, init_c, initialization_batch, \n training=False,\n )\n \n # Load weights after variable initialization\n if encoder_wts is not None:\n print(\"\\t--> LOADING ENCODER MODEL WEIGHTS ...\")\n encoder.load_weights(encoder_wts)\n if decoder_wts is not None:\n print(\"\\t--> LOADING DECODER MODEL WEIGHTS ...\")\n decoder.load_weights(decoder_wts)\n \n # Show the model architectures and plot the learning rate\n if verbose:\n print(\"\\n\\n... ENCODER MODEL SUMMARY...\\n\")\n print(encoder.summary())\n\n print(\"\\n\\n... DECODER MODEL SUMMARY...\\n\")\n print(decoder.summary())\n\n return encoder, decoder\n \n \nprint(\"\\n... GENERATING THE FOLLOWING:\")\n# Instantiate our required training components in the correct scope\nencoder, decoder = prepare_for_inference(\n encoder_config=inference_config.encoder_config,\n decoder_config=inference_config.decoder_config,\n encoder_wts=(ENCODER_CKPT_PATH if ENCODER_CKPT_PATH!=\"\" else None),\n decoder_wts=(DECODER_CKPT_PATH if DECODER_CKPT_PATH!=\"\" else None),\n)\n\nprint(\"\\n... INFERENCE PREPERATION FINISHED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:33:04.689947Z","iopub.execute_input":"2021-06-03T23:33:04.690463Z","iopub.status.idle":"2021-06-03T23:33:56.154639Z","shell.execute_reply.started":"2021-06-03T23:33:04.690425Z","shell.execute_reply":"2021-06-03T23:33:56.152956Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n\n

6  MODEL INFERENCE    ⤒

\n\nIn this section we will define the inference routine/step as well as the final inference loop that will execute everything we have worked on up until this point.","metadata":{}},{"cell_type":"markdown","source":"

6.1 INDIVIDUAL INFERENCE STEP

\n\n---\n\nWe use this function to represent the calculation occuring on one replica for one batch of images/labels. As this is for inference we want to pass back the predicted sequence.\n\n***WIP***","metadata":{}},{"cell_type":"code","source":"def beam_test_step(_image_batch, b_width=8, beta_factor=5):\n \"\"\" TBD \n \n Args:\n _image_batch (Image Tensor): Batch of images\n _b_width (int): The width of the beam to use\n \n Returns:\n The most likely sequence (argmax over final k beams)\n \"\"\"\n \n ##################################################################################\n # STEP 1: Get image embedding via the Encoder\n # ----- output shape -----> (REPLICA_BATCH_SIZE, IMG_EMB_DIM)\n ##################################################################################\n image_batch_embedding = encoder(_image_batch, training=False)\n ##################################################################################\n\n ##################################################################################\n # STEP 2: Initialize the hidden and memory objects for decoding\n # ----- output shapes are both -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH, N_RNN_UNITS)\n ##################################################################################\n hidden_batch, memory_batch = decoder.init_hidden_state(image_batch_embedding, training=False)\n hidden_batch = tf.repeat(tf.expand_dims(hidden_batch, axis=1), repeats=b_width, axis=1)\n memory_batch = tf.repeat(tf.expand_dims(memory_batch, axis=1), repeats=b_width, axis=1)\n ##################################################################################\n\n ##################################################################################\n # STEP 3: Initialize starting probabilities and sequence tokens. \n # These will get overwritten or appended to respectively.\n # ----- output shapes are both -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH, 1)\n ##################################################################################\n top_probabilities_batch = tf.concat([\n tf.ones((REPLICA_BATCH_SIZE, 1, 1), dtype=tf.float32)*45,\n tf.zeros((REPLICA_BATCH_SIZE, b_width-1, 1), dtype=tf.float32),\n ], axis=-2)\n top_k_prediction_seq_batch = tf.concat([\n tf.ones((REPLICA_BATCH_SIZE, 1, 1), dtype=tf.uint8),\n tf.zeros((REPLICA_BATCH_SIZE, b_width-1, 1), dtype=tf.uint8),\n ], axis=-2)\n ##################################################################################\n\n ##################################################################################\n # STEP 4: Loop the token level prediction number of times\n ##################################################################################\n for c_idx in tqdm(range(1, MAX_LEN), total=MAX_LEN-1): #MAX_LEN\n \n ##################################################################################\n # STEP 5: Initialize a list for every beam-level item we want to accumulate\n # ----- items are -----> probabilties, predictions (indices), hidden-states, memory-states\n ##################################################################################\n\n probs_batch_k_list = [] # cumulative\n preds_batch_k_list = [] # cumulative\n\n # LSTM understanding of the previous tokens/stuff\n hidden_batch_k_list = []\n memory_batch_k_list = []\n\n ##################################################################################\n # STEP 6: Loop the beam level predictions number of times\n ##################################################################################\n for k in range(b_width): \n \n ##################################################################################\n # STEP 6: Get the kth value for probabilities, token, hidden-state, and memory-state\n # ----- decoder_probabilities_batch_k shape -----> (REPLICA_BATCH_SIZE, 1)\n # ----- decoder_input_batch_k shape -----> (REPLICA_BATCH_SIZE, c_idx-1) -- it grows as we go\n # ----- hidden_batch_k shape -----> (REPLICA_BATCH_SIZE, N_RNN_UNITS)\n # ----- memory_batch_k shape -----> (REPLICA_BATCH_SIZE, N_RNN_UNITS)\n ##################################################################################\n decoder_probabilities_batch_k = top_probabilities_batch[:, k]\n decoder_input_batch_k = top_k_prediction_seq_batch[:, k, -1:]\n hidden_batch_k = hidden_batch[:, k]\n memory_batch_k = memory_batch[:, k]\n\n ##################################################################################\n # STEP 7: Perform inference with the kth values captured previously and get outputs\n # ----- prediction_batch_k shape -----> (REPLICA_BATCH_SIZE, VOCAB_LEN)\n # ----- output_hidden_batch_k shape -----> (REPLICA_BATCH_SIZE, N_RNN_UNITS)\n # ----- output_memory_batch_k shape -----> (REPLICA_BATCH_SIZE, N_RNN_UNITS)\n ##################################################################################\n prediction_batch_k, output_hidden_batch_k, output_memory_batch_k = \\\n decoder(decoder_input_batch_k, hidden_batch_k, memory_batch_k, image_batch_embedding, training=False)\n\n ##################################################################################\n # STEP 8: Get the top_k (k=BEAM_WIDTH) probabilties and their respective indices (tokens)\n # ----- raw_probs shape -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH)\n # ----- raw_indices shape -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH)\n ##################################################################################\n raw_probs, raw_indices = tf.math.top_k(prediction_batch_k, k=b_width)\n norm_probs = tf.nn.log_softmax(raw_probs/beta_factor) # beta factor determined off-cuff\n \n ##################################################################################\n # STEP 9: Use the indices captured above and the previous predictions to generate \n # the top k (BEAM_WIDTH) sequences for this beam.\n # ----- top_preds shape -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH, 1)\n ##################################################################################\n prep_k_seq_batch = tf.repeat(tf.expand_dims(top_k_prediction_seq_batch[:, k], axis=1), repeats=b_width, axis=1)\n prep_raw_indices_batch = tf.expand_dims(tf.cast(raw_indices, tf.uint8), axis=-1)\n prep_raw_indices_batch = tf.expand_dims(tf.where(tf.repeat(decoder_input_batch_k, repeats=b_width, axis=1)!=2, prep_raw_indices_batch[..., 0], 2), axis=-1)\n top_preds = tf.concat([\n prep_k_seq_batch,\n prep_raw_indices_batch\n ], axis=-1)\n\n ##################################################################################\n # STEP 10: Capture relevant values within the previously mentioned lists \n # as we will be using these lists as inputs into tf.concat\n # ----- NOTE ON SHAPES -----> hidden and memory batches have a dimension added and tripled\n # so that we can concat along it in the future\n ##################################################################################\n preds_batch_k_list.append(top_preds)\n probs_batch_k_list.append(norm_probs+decoder_probabilities_batch_k)\n hidden_batch_k_list.append(tf.repeat(tf.expand_dims(output_hidden_batch_k, axis=1), repeats=b_width, axis=1))\n memory_batch_k_list.append(tf.repeat(tf.expand_dims(output_memory_batch_k, axis=1), repeats=b_width, axis=1))\n\n \n ##################################################################################\n # STEP 11: After the beam looping is completed we concat the data along the 'beam-axis' (axis=1)\n # ----- preds_batch_k_arr shape -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH*BEAM_WIDTH, c_idx-1)\n # ----- probs_batch_k_arr shape -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH*BEAM_WIDTH, 1)\n # ----- hidden_batch_k_arr shape -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH*BEAM_WIDTH, 1)\n # ----- memory_batch_k_arr shape -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH*BEAM_WIDTH, 1)\n ##################################################################################\n preds_batch_k_arr = tf.concat(preds_batch_k_list, axis=1) \n probs_batch_k_arr = tf.concat(probs_batch_k_list, axis=1) \n hidden_batch_k_arr = tf.concat(hidden_batch_k_list, axis=1) \n memory_batch_k_arr = tf.concat(memory_batch_k_list, axis=1)\n\n ##################################################################################\n # STEP 12: Get the top k probabilities indices and use this to capture the top k items\n # for each arraay\n # ----- top_k_prediction_seq_batch shape -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH, c_idx-1)\n # ----- top_probabilities_batch shape -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH, 1)\n # ----- hidden_batch shape -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH, 1)\n # ----- memory_batch shape -----> (REPLICA_BATCH_SIZE, BEAM_WIDTH, 1)\n ##################################################################################\n filter_indices = tf.math.top_k(probs_batch_k_arr, k=b_width).indices\n\n top_k_prediction_seq_batch = tf.gather(preds_batch_k_arr, filter_indices, batch_dims=1)\n top_probabilities_batch = tf.expand_dims(tf.gather(probs_batch_k_arr, filter_indices, batch_dims=1), axis=-1)\n hidden_batch = tf.gather(hidden_batch_k_arr, filter_indices, batch_dims=1)\n memory_batch = tf.gather(memory_batch_k_arr, filter_indices, batch_dims=1)\n \n # Take the top most likely sequence and convert it into the appropriate datatype\n return tf.cast(tf.gather(top_k_prediction_seq_batch, tf.argmax(top_probabilities_batch, axis=1), batch_dims=1)[:, 0, :], tf.int32)\n\n@tf.function\ndef distributed_beam_test_step(_img_batch, _img_ids):\n per_replica_seqs = strategy.run(beam_test_step, args=(_img_batch,))\n predictions = strategy.gather(per_replica_seqs, axis=0)\n pred_ids = strategy.gather(_img_ids, axis=0)\n return predictions, pred_ids","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:33:56.157175Z","iopub.execute_input":"2021-06-03T23:33:56.157605Z","iopub.status.idle":"2021-06-03T23:33:56.191365Z","shell.execute_reply.started":"2021-06-03T23:33:56.157563Z","shell.execute_reply":"2021-06-03T23:33:56.189524Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def test_step(_image_batch):\n \"\"\" Forward pass (calculate gradients)\n \n Args:\n image_batch (): TBD\n inchi_batch (): TBD\n \n Returns:\n tbd\n \"\"\"\n \n # image_batch_embedding has shape --> (REPLICA_BATCH_SIZE, IMG_EMB_DIM)\n image_batch_embedding = encoder(_image_batch, training=False)\n\n # hidden and memory both have the shape --> (REPLICA_BATCH_SIZE, N_RNN_UNITS)\n hidden_batch, memory_batch = decoder.init_hidden_state(image_batch_embedding, training=False)\n\n # decoder_input and predictions_seq share the shape --> (REPLICA_BATCH_SIZE, 1)\n decoder_input_batch = tf.ones((REPLICA_BATCH_SIZE, 1), dtype=tf.uint8)\n predictions_seq_batch = tf.ones((REPLICA_BATCH_SIZE, 1), dtype=tf.uint8)\n confidence_over_batch = tf.zeros((REPLICA_BATCH_SIZE, 1), dtype=tf.float32)\n\n # Teacher forcing - feeding the target as the next input\n for c_idx in range(1, MAX_LEN):\n \n # passing enc_output to the decoder\n prediction_batch, hidden_batch, memory_batch = \\\n decoder(decoder_input_batch, hidden_batch, memory_batch, image_batch_embedding, training=False)\n \n confidence_over_batch += tf.reduce_max(prediction_batch, axis=-1, keepdims=True)\n \n # no teacher forcing, predicted char is next LSTMCell input\n decoder_input_batch = tf.cast(tf.expand_dims(tf.math.argmax(prediction_batch, axis=1, output_type=tf.int32), axis=1), tf.uint8)\n \n # Build the prediction sequence\n predictions_seq_batch = tf.concat([predictions_seq_batch, decoder_input_batch], axis=1)\n \n return predictions_seq_batch, confidence_over_batch\n\n \n@tf.function\ndef distributed_test_step(_img_batch, _img_ids):\n per_replica_seqs, per_replica_confs = strategy.run(test_step, args=(_img_batch,))\n predictions = strategy.gather(per_replica_seqs, axis=0)\n confidences = strategy.gather(per_replica_confs, axis=0)\n pred_ids = strategy.gather(_img_ids, axis=0)\n return predictions, confidences, pred_ids","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:33:56.195436Z","iopub.execute_input":"2021-06-03T23:33:56.195834Z","iopub.status.idle":"2021-06-03T23:33:56.211779Z","shell.execute_reply.started":"2021-06-03T23:33:56.195801Z","shell.execute_reply":"2021-06-03T23:33:56.210979Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

6.2 RAW INFERENCE LOOP

\n\n---\n\n***WIP***","metadata":{}},{"cell_type":"code","source":"DO_BEAM = False\n\n# To Store The Preds\nall_pred_arr = tf.zeros((1, MAX_LEN), dtype=tf.uint8)\nall_pred_ids = tf.zeros((1, 1), dtype=tf.string)\n\nif DO_BEAM:\n # Create an iterator\n dist_test_ds = iter(strategy.experimental_distribute_dataset(test_ds))\n for i in tqdm(range(TEST_STEPS), total=TEST_STEPS): \n img_batch, id_batch = next(dist_test_ds)\n preds, pred_ids = distributed_beam_test_step(img_batch, id_batch)\n all_pred_arr = tf.concat([all_pred_arr, tf.cast(preds, tf.uint8)], axis=0)\n all_pred_ids = tf.concat([all_pred_ids, tf.expand_dims(pred_ids, axis=-1)], axis=0)\nelse:\n # Create an iterator\n dist_test_ds = iter(strategy.experimental_distribute_dataset(test_ds))\n for i in tqdm(range(TEST_STEPS), total=TEST_STEPS): \n img_batch, id_batch = next(dist_test_ds)\n preds, _, pred_ids = distributed_test_step(img_batch, id_batch)\n all_pred_arr = tf.concat([all_pred_arr, tf.cast(preds, tf.uint8)], axis=0)\n all_pred_ids = tf.concat([all_pred_ids, tf.expand_dims(pred_ids, axis=-1)], axis=0)","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:33:56.213093Z","iopub.execute_input":"2021-06-03T23:33:56.213666Z","iopub.status.idle":"2021-06-03T23:45:30.421987Z","shell.execute_reply.started":"2021-06-03T23:33:56.213623Z","shell.execute_reply":"2021-06-03T23:45:30.420791Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

6.3 POST PROCESSING

\n\n---\n\n***WIP***","metadata":{}},{"cell_type":"code","source":"def arr_2_inchi(arr):\n inchi_str = ''\n for i in arr:\n c = int_2_tok.get(i)\n if c==\"\":\n break\n inchi_str += c\n return inchi_str\n\npred_df = pd.DataFrame({\n \"image_id\":[x[0].decode() for x in tqdm(all_pred_ids[1:-REQUIRED_DATASET_PAD].numpy(), total=N_TEST)], \n \"InChI\":[arr_2_inchi(pred_arr) for pred_arr in tqdm(all_pred_arr[1:-REQUIRED_DATASET_PAD].numpy(), total=N_TEST)]\n})\n\npred_df = pred_df.sort_values(by=\"image_id\").reset_index(drop=True)\npred_df.to_csv(\"submission.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2021-06-03T23:45:30.423715Z","iopub.execute_input":"2021-06-03T23:45:30.424176Z","iopub.status.idle":"2021-06-03T23:54:11.590006Z","shell.execute_reply.started":"2021-06-03T23:45:30.424127Z","shell.execute_reply":"2021-06-03T23:54:11.588983Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}