\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
- The InChI strings and Image IDs will just be left as byte string tensors.\n
\n- 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 📖 REFERENCE:
\n - Tutorial - TFRecord and tf.Example
\n - TFRecordDataset Documentation
\n - Decoding PNGs Documentation
\n
\n","metadata":{"papermill":{"duration":0.074838,"end_time":"2021-06-15T18:00:55.78693","exception":false,"start_time":"2021-06-15T18:00:55.712092","status":"completed"},"tags":[],"id":"impressed-routine"}},{"cell_type":"code","source":"def decode(serialized_example, is_test=False, tokenized_inchi=True):\n \"\"\" \n Function to parses a set of features and label from the given `serialized_example`.\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=[],\n dtype=tf.string, \n default_value=''),\n }\n \n if not is_test:\n if tokenized_inchi:\n feature_dict[\"inchi\"] = tf.io.FixedLenFeature(shape=[MAX_LEN], \n dtype=tf.int64, \n default_value=[0]*MAX_LEN)\n else:\n feature_dict[\"inchi\"] = tf.io.FixedLenFeature(shape=[], \n dtype=tf.string, \n default_value='')\n else:\n feature_dict[\"image_id\"] = tf.io.FixedLenFeature(shape=[], \n dtype=tf.string, \n 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":{"papermill":{"duration":0.089822,"end_time":"2021-06-15T18:00:55.952706","exception":false,"start_time":"2021-06-15T18:00:55.862884","status":"completed"},"tags":[],"id":"alpine-lloyd","execution":{"iopub.status.busy":"2022-01-25T14:53:24.908128Z","iopub.execute_input":"2022-01-25T14:53:24.908724Z","iopub.status.idle":"2022-01-25T14:53:24.921484Z","shell.execute_reply.started":"2022-01-25T14:53:24.908684Z","shell.execute_reply":"2022-01-25T14:53:24.920846Z"},"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` \ntrain_ds = raw_train_ds.map(lambda x: decode(x, is_test=False))\nval_ds = raw_val_ds.map(lambda x: decode(x, is_test=False))\ntest_ds = raw_test_ds.map(lambda x: decode(x, is_test=True))\n\nprint(f\"\\n... THE DECODED TF.DATA.TFRECORDDATASET OBJECT:\" \\\n f\"\\n\\t--> ((image), (image_id - optional), (inchi))\" \\\n f\"\\n\\t--> {train_ds}\\n\")\n\nprint(\"\\n... 2 EXAMPLES OF IMAGES AND LABELS AFTER DECODING ...\")\nfor i, (img, inchi) in enumerate(train_ds.take(2)):\n print(f\"\\nIMAGE SHAPE : {img.shape}\")\n print(f\"IMAGE INCHI : {[int_2_tok[x] for x in inchi.numpy()]}\\n\")\n plt.figure(figsize=(10,10))\n plt.imshow(img.numpy().astype(np.int64), cmap=\"gray\")\n plt.title(f\"{''.join([int_2_tok[x] for x in inchi.numpy() if x!=0][:50])} ... [truncated]\")\n plt.show()\n\nprint(\"\\n... DECODING RAW TFRECORD DATASETS COMPLETED ...\\n\")","metadata":{"papermill":{"duration":0.922426,"end_time":"2021-06-15T18:00:56.950876","exception":false,"start_time":"2021-06-15T18:00:56.02845","status":"completed"},"tags":[],"id":"other-rough","outputId":"d0e75f86-5dae-48e6-9cdc-825602e95261","execution":{"iopub.status.busy":"2022-01-25T14:53:24.922729Z","iopub.execute_input":"2022-01-25T14:53:24.923386Z","iopub.status.idle":"2022-01-25T14:53:25.967522Z","shell.execute_reply.started":"2022-01-25T14:53:24.923339Z","shell.execute_reply":"2022-01-25T14:53:25.96656Z"},"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 - Tutorial - TFRecord and tf.Example
\n - TFRecordDataset Documentation
\n - Decoding PNGs Documentation
\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":{"papermill":{"duration":0.08202,"end_time":"2021-06-15T18:00:57.12714","exception":false,"start_time":"2021-06-15T18:00:57.04512","status":"completed"},"tags":[],"id":"fatty-shopper"}},{"cell_type":"code","source":"def load_dataset(filenames, is_test=False, ordered=False, tokenized_inchi=True):\n \"\"\"\n Function to read the dataset from TFRecords.\n For optimal performance, reading from multiple files at once and \n disregarding data order (if `ordered=False`).\n - If pulling InChI from TFRecords than order does not matter since we \n will be shuffling the data anyway (for training dataset).\n \n Args:\n filenames (list of strings): List of paths to that point to the \n respective TFRecord files\n is_test (bool, optional): Whether or not to include the image ID or \n label in the returned dataset\n ordered (bool, optional): Whether to ensured ordered results or \n maximize parallelization\n tokenized_inchi (bool, optional): Whether our dataset includes the \n tokenized inchi or we will be creating \n it from the caption numpy array\n \n Returns:\n dataset: 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, is_test=False, shuffle_buffer_size=1, \n repeat_dataset=True, preserve_file_order=False, \n drop_remainder=True, tokenized_inchi=True,\n external_inchi_dataset=None, test_padding=0):\n \"\"\" \n Function to get a tf.data.Dataset with the appropriate configuration.\n \n Args:\n filenames (list of strings): List of paths to that point to the respective TFRecord files\n batch_size (int): Batch size to be used during batching the dataset\n is_test (bool, optional): Whether or not to include the image ID or label in the returned dataset\n shuffle_buffer_size (int, optional): Number of elements from which the new dataset will be sampled\n repeat_dataset (bool, optional): Whether the dataset is to be repeated \n preserve_file_order (bool, optional): Whether to ensured ordered results or maximize parallelization\n drop_remainder (bool, optional): Whether the last batch should be dropped \n in case its size is smaller than desired\n tokenized_inchi (bool, optional): Whether our dataset includes the \n tokenized inchi or we will be creating \n it from the caption numpy array\n external_inchi_dataset (obj or None): None if no external inchi dataset is uploaded \n test_padding (int, optional): Amount required to pad dataset to have only full batches\n \n Returns:\n dataset: new tf.data.Dataset object\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":{"papermill":{"duration":0.097793,"end_time":"2021-06-15T18:00:57.307718","exception":false,"start_time":"2021-06-15T18:00:57.209925","status":"completed"},"tags":[],"id":"tough-criminal","execution":{"iopub.status.busy":"2022-01-25T14:53:25.968761Z","iopub.execute_input":"2022-01-25T14:53:25.96904Z","iopub.status.idle":"2022-01-25T14:53:25.986Z","shell.execute_reply.started":"2022-01-25T14:53:25.969003Z","shell.execute_reply":"2022-01-25T14:53:25.984954Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 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# Individual Respective Configurations\nTRAIN_DS_CONFIG = DS_TEMPLATE_CONFIG.copy()\nTRAIN_DS_CONFIG.update(dict(\n filenames=TRAIN_TFREC_PATHS,\n batch_size=OVERALL_BATCH_SIZE,\n shuffle_buffer_size=OVERALL_BATCH_SIZE*6,\n))\n\nVAL_DS_CONFIG = DS_TEMPLATE_CONFIG.copy()\nVAL_DS_CONFIG.update(dict(\n filenames=VAL_TFREC_PATHS,\n batch_size=OVERALL_BATCH_SIZE,\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\ntrain_ds = get_dataset(**TRAIN_DS_CONFIG)\nval_ds = get_dataset(**VAL_DS_CONFIG)\ntest_ds = get_dataset(**TEST_DS_CONFIG)\n\nfor SPLIT, CONFIG in zip([\"TRAINING\", \"VALIDATION\", \"TESTING\"], [TRAIN_DS_CONFIG, \n VAL_DS_CONFIG, \n 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\"\\n\\n... TRAINING DATASET : {train_ds} ...\")\nprint(f\"... VALIDATION DATASET : {val_ds} ...\")\nprint(f\"... TESTING DATASET : {test_ds} ...\\n\")\n\nprint(\"\\n\\n ... SOME VALIDATION EXAMPLES ... \\n\\n\")\nfor x,y in val_ds.take(1):\n for i in range(2):\n plt.figure(figsize=(12,12))\n plt.imshow(x[i].numpy().astype(np.int64))\n plt.title(f\"IMAGE INCHI : {''.join([int_2_tok[z] for z in y[i].numpy() if z not in [0,1,2]])}\\n\")\n plt.show()\n \nprint(\"\\n\\n ... SOME TESTING EXAMPLES ... \\n\\n\")\nfor x,y in test_ds.take(1):\n for i in range(2):\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":{"papermill":{"duration":12.083416,"end_time":"2021-06-15T18:01:09.476216","exception":false,"start_time":"2021-06-15T18:00:57.3928","status":"completed"},"tags":[],"id":"intelligent-lying","outputId":"e029274d-3f27-4387-85b3-f94f71d41330","execution":{"iopub.status.busy":"2022-01-25T14:53:25.989714Z","iopub.execute_input":"2022-01-25T14:53:25.99041Z","iopub.status.idle":"2022-01-25T14:53:36.737427Z","shell.execute_reply.started":"2022-01-25T14:53:25.990371Z","shell.execute_reply":"2022-01-25T14:53:36.736233Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 5. MODEL PREPERATION ","metadata":{"id":"KLLZLlt_2upb"}},{"cell_type":"markdown","source":"\nIn this section we prepare the models for training. More information on Vision Transformers can we found on the [Google-search Vision Transformer implementation](https://github.com/google-research/vision_transformer).\n\n\n - Guide - tf.data: Build TensorFlow Input Pipelines
\n - Guide - Better Performance With the tf.data API
\n - tf.data.Dataset Documentation
\n
\n\n
![](\"https://github.com/google-research/vision_transformer/raw/main/vit_figure.png\")
","metadata":{"papermill":{"duration":0.090282,"end_time":"2021-06-15T18:01:09.658503","exception":false,"start_time":"2021-06-15T18:01:09.568221","status":"completed"},"tags":[],"id":"authorized-insider"}},{"cell_type":"markdown","source":"## 5.1 UNDERSTANDING THE MODELS - ViT\n","metadata":{"id":"EFRKuTZZ3PlA"}},{"cell_type":"markdown","source":"### 5.1.1 ViT - Implement Patch Creation as a Layer","metadata":{"id":"nQpl3kTJ3UyJ"}},{"cell_type":"code","source":"class PatchCreator(tf.keras.layers.Layer):\n ''' Creates Patches for input images '''\n def __init__(self, patch_size):\n '''\n Args:\n patch_size (int): Size of the patches to be extracted from the \n input images\n \n Returns:\n None, this is initialization\n '''\n super(PatchCreator, self).__init__()\n self.patch_size = patch_size\n\n def call(self, images):\n ''' Calling function for patch creation class '''\n batch_size = tf.shape(images)[0]\n patches = tf.image.extract_patches(\n images=images,\n sizes=[1, self.patch_size, self.patch_size, 1],\n strides=[1, self.patch_size, self.patch_size, 1],\n rates=[1, 1, 1, 1],\n padding=\"VALID\",\n )\n patch_dims = patches.shape[-1]\n patches = tf.reshape(patches, [batch_size, -1, patch_dims])\n return patches\n","metadata":{"id":"EasOQO5jfi9K","execution":{"iopub.status.busy":"2022-01-25T14:53:36.739529Z","iopub.execute_input":"2022-01-25T14:53:36.739905Z","iopub.status.idle":"2022-01-25T14:53:36.749395Z","shell.execute_reply.started":"2022-01-25T14:53:36.739858Z","shell.execute_reply":"2022-01-25T14:53:36.748807Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Grab a demo image and label and define an arbitrary patch_size\ndemo_img, demo_lbl = next(iter(train_ds.unbatch().batch(1)))\ndemo_patch_size=16\n\n# Instantiate the PatchCreator layer and call on the demo image\nwith tf.device('/CPU:0'):\n patch_creator = PatchCreator(demo_patch_size)\n patches = patch_creator(demo_img)\n\nprint(f\"Image size: {IMG_SHAPE}\")\nprint(f\"Patch size: {(demo_patch_size, demo_patch_size)}\")\nprint(f\"Patches shape: {patches.shape}\")\nprint(f\"Patches per image: {patches.shape[1]}\")\nprint(f\"Elements per patch: {patches.shape[-1]}\")","metadata":{"id":"OeL9eZGHfouL","execution":{"iopub.status.busy":"2022-01-25T14:53:36.750297Z","iopub.execute_input":"2022-01-25T14:53:36.75097Z","iopub.status.idle":"2022-01-25T14:53:41.157016Z","shell.execute_reply.started":"2022-01-25T14:53:36.750939Z","shell.execute_reply":"2022-01-25T14:53:41.156085Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# PATCH LAYER FUNCTIONALITY\n\n# 1. Plot the original image\nprint(\"\\n\\n... ORIGINAL IMAGE ...\\n\")\nplt.figure(figsize=(18,9))\nplt.imshow(demo_img[0].numpy().astype(\"float32\")/255.)\nplt.axis('off')\nplt.tight_layout()\nplt.show()\n\n# 2. Plot the patches in the same shape/order as the original image\nprint(\"\\n\\n\\n... IMAGE PATCHES ...\\n\")\nplt.figure(figsize=(18,9))\nfor i, patch in enumerate(patches[0]):\n ax = plt.subplot(int(np.ceil(IMG_SHAPE[0]/demo_patch_size)), \n int(np.ceil(IMG_SHAPE[1]/demo_patch_size)), i + 1)\n patch_img = tf.reshape(patch, (demo_patch_size, demo_patch_size, 3))\n plt.imshow(patch_img.numpy().astype(\"float32\")/255.)\n plt.axis(\"off\")\nplt.tight_layout()\nplt.subplots_adjust(wspace=0.03, hspace=0.06)\nplt.show()","metadata":{"papermill":{"duration":19.595564,"end_time":"2021-06-15T18:01:29.709609","exception":false,"start_time":"2021-06-15T18:01:10.114045","status":"completed"},"tags":[],"id":"operational-maker","outputId":"9fbbfdad-e89e-49f2-fcd6-16c40ebb4676","execution":{"iopub.status.busy":"2022-01-25T14:53:41.158476Z","iopub.execute_input":"2022-01-25T14:53:41.160375Z","iopub.status.idle":"2022-01-25T14:53:55.472144Z","shell.execute_reply.started":"2022-01-25T14:53:41.160337Z","shell.execute_reply":"2022-01-25T14:53:55.471278Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5.1.2 ViT - Implement Patch Encoder as a Layer\n\n\n---\n\nThe **`PatchEncoder`** layer will linearly transform a patch by projecting it into a vector of size **`projection_dim`**. In addition, it adds a learnable position embedding to the projected vector.","metadata":{"papermill":{"duration":0.095873,"end_time":"2021-06-15T18:01:29.898822","exception":false,"start_time":"2021-06-15T18:01:29.802949","status":"completed"},"tags":[],"id":"labeled-disclosure"}},{"cell_type":"code","source":"class PatchEncoder(tf.keras.layers.Layer):\n def __init__(self, num_patches, projection_dim):\n '''\n Implements the PatchEncoder Block\n\n Args:\n num_patches (int): Number of input image patches\n projection_dim (int): Dimension of the output projection vector\n \n Returns:\n None; this is initialization\n '''\n super(PatchEncoder, self).__init__()\n \n self.num_patches = num_patches\n self.projection_dim = projection_dim\n \n self.dense_projection = tf.keras.layers.Dense(units=self.projection_dim)\n self.positions = tf.reshape(tf.range(start=0, limit=self.num_patches, delta=1), (self.num_patches,))\n self.position_embedding = tf.keras.layers.Embedding(input_dim=self.num_patches, \n output_dim=self.projection_dim)\n\n def call(self, patch):\n ''' Returns the encoded patches of image '''\n encoded = self.dense_projection(patch) + self.position_embedding(self.positions)\n return encoded\n ","metadata":{"id":"0QqRsohwl3pA","execution":{"iopub.status.busy":"2022-01-25T14:53:55.473589Z","iopub.execute_input":"2022-01-25T14:53:55.473822Z","iopub.status.idle":"2022-01-25T14:53:55.481901Z","shell.execute_reply.started":"2022-01-25T14:53:55.473794Z","shell.execute_reply":"2022-01-25T14:53:55.481193Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Define an arbitrary projection_dim\ndemo_projection_dim = 128\n\n# Instantiate the PatcheEncoder layer and call on the demo image\nwith tf.device('/CPU:0'):\n patch_encoder = PatchEncoder(num_patches=patches.shape[1], projection_dim=demo_projection_dim)\n encoded_patches = patch_encoder(patches)\n\nprint(f\"Number of Patches: {patches.shape[1]}\") \nprint(f\"Patch size: {(demo_patch_size, demo_patch_size)}\")\nprint(f\"Encoded patches shape: {encoded_patches.shape}\")\n\n# PATCH LAYER FUNCTIONALITY\n\n# 1. Plot the original image\nprint(\"\\n\\n... ORIGINAL IMAGE ...\\n\")\nplt.figure(figsize=(18,9))\nplt.imshow(demo_img[0].numpy().astype(\"float32\")/255.)\nplt.axis('off')\nplt.tight_layout()\nplt.show()\n\n# 2. Plot the patches in the same shape/order as the original image\nprint(\"\\n\\n\\n... IMAGE PATCHES ...\\n\")\nplt.figure(figsize=(18,9))\nfor i, patch in enumerate(patches[0]):\n w,h = int(np.ceil(IMG_SHAPE[0]/demo_patch_size)), int(np.ceil(IMG_SHAPE[1]/demo_patch_size))\n ax = plt.subplot(w,h,i+1)\n patch_img = tf.reshape(patch, (demo_patch_size, demo_patch_size, 3))\n plt.imshow(patch_img.numpy().astype(\"float32\")/255.)\n plt.axis(\"off\")\nplt.tight_layout()\nplt.subplots_adjust(wspace=0.03, hspace=0.06)\nplt.show()\n\nwith tf.device('/CPU:0'):\n enc_patches_rescaled = patch_creator(tf.expand_dims(tf.image.resize(tf.reshape(tf.reduce_mean(encoded_patches[0], axis=-1), (w, h, 1)), (IMG_SHAPE[0], IMG_SHAPE[1])), axis=0))\n\n# 3. Plot the image embedding patches\nprint(\"\\n\\n\\n... IMAGE EMBEDDING AS PATCH VISUALIZATION ...\\n\")\nplt.figure(figsize=(18,9))\nfor i, enc_patch in enumerate(enc_patches_rescaled[0]):\n ax = plt.subplot(w,h,i+1)\n enc_patch_img = tf.reshape(enc_patch, (demo_patch_size, demo_patch_size))\n plt.imshow(enc_patch_img.numpy().astype(\"float32\")/255., cmap=\"jet\")\n plt.axis(\"off\")\nplt.tight_layout()\nplt.subplots_adjust(wspace=0.03, hspace=0.06)\nplt.show()\n","metadata":{"papermill":{"duration":29.60816,"end_time":"2021-06-15T18:01:59.821889","exception":false,"start_time":"2021-06-15T18:01:30.213729","status":"completed"},"tags":[],"id":"about-memorial","outputId":"01d8323f-fba8-4c3e-f37e-252c3de5ec81","execution":{"iopub.status.busy":"2022-01-25T14:53:55.483214Z","iopub.execute_input":"2022-01-25T14:53:55.484068Z","iopub.status.idle":"2022-01-25T14:54:24.712818Z","shell.execute_reply.started":"2022-01-25T14:53:55.48402Z","shell.execute_reply":"2022-01-25T14:54:24.711974Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5.1.3 Build the ViT Model\n\n---\n\nThe ViT model consists of multiple Transformer blocks, which use the layers.MultiHeadAttention layer as a self-attention mechanism applied to the sequence of patches. The Transformer blocks produce a **[batch_size, num_patches, projection_dim]** tensor, which is processed via an classifier head with softmax to produce the final class probabilities output.\n\nUnlike the technique described in the paper, which prepends a learnable embedding to the sequence of encoded patches to serve as the image representation, all the outputs of the final Transformer block are reshaped with **`tf.keras.layers.Flatten()`** and used as the image representation input to the classifier head. Note that the **`tf.keras.layers.GlobalAveragePooling1D`** layer could also be used instead to aggregate the outputs of the Transformer block, especially when the number of patches and the projection dimensions are large.","metadata":{"papermill":{"duration":0.103216,"end_time":"2021-06-15T18:02:00.028634","exception":false,"start_time":"2021-06-15T18:01:59.925418","status":"completed"},"tags":[],"id":"tutorial-water"}},{"cell_type":"code","source":"class ViTEncoder(tf.keras.Model):\n ''' Creates the stack of encoders '''\n def __init__(self, patch_size=16, projection_dim=256, n_transformer_layers=8,\n n_heads=4, dropout=0.1, img_shape=IMG_SHAPE):\n \"\"\"\n Initilizes the varibales needed to create the model\n \n Args:\n patch_size (int): Size of the patches to be extracted from the input images\n projection_dim (int): Dimension of the output projection vector\n n_transformer_layers (int): Number of transformer layers in the architecture\n n_heads (int): Number of attention head to be used\n dropout (float): Dropout value to be used in MLP dropout layers\n img_shape (int): Shape of the input image\n\n Returns:\n None, this is a initialization \n \"\"\"\n super(ViTEncoder, self).__init__()\n \n # Layer Arguments\n self.patch_size = patch_size\n self.n_patches = tf.cast(tf.round((img_shape[0]/self.patch_size)*(img_shape[1]/self.patch_size)), tf.int32)\n self.img_shape = img_shape\n self.projection_dim = projection_dim\n self.n_transformer_layers = n_transformer_layers\n self.mlp_intermediate_units = [self.projection_dim*2, self.projection_dim,]\n self.n_heads = n_heads\n self.dropout = dropout\n \n # Layers\n self.patch_creator = PatchCreator(self.patch_size)\n self.patch_encoder = PatchEncoder(self.n_patches, self.projection_dim)\n self.ln_1_layer = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n self.ln_2_layer = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n self.add_1_layer = tf.keras.layers.Add()\n self.add_2_layer = tf.keras.layers.Add()\n self.mha_layer = tf.keras.layers.MultiHeadAttention(num_heads=self.n_heads, \n key_dim=self.projection_dim, \n dropout=self.dropout)\n self.intermediate_mlp_dense_layers = [tf.keras.layers.Dense(transformer_units, activation=tf.nn.gelu) \\\n for transformer_units in self.mlp_intermediate_units]\n self.intermediate_mlp_dropout_layers = [tf.keras.layers.Dropout(self.dropout) \\\n for transformer_units in self.mlp_intermediate_units] \n \n def call(self, x, training):\n \"\"\"Creating the Stack of Encoders in transformer\n \n Args:\n x (array) : Input token embedding\n training (bool): Whether to train or not \n \n Returns:\n Stack of encoder architecture\n \"\"\"\n patches = self.patch_creator(x, training=training)\n encoded = self.patch_encoder(patches, training=training)\n \n for _ in range(self.n_transformer_layers):\n \n # Layer Norm 1\n x1 = self.ln_1_layer(encoded, training=training)\n \n # Create a multi-head attention layer.\n attention_output = self.mha_layer(x1, x1, training=training)\n \n # Skip Connection 1\n x2 = self.add_1_layer([attention_output, encoded], training=training)\n \n # Layer Norm 2\n x3 = self.ln_2_layer(x2, training=training)\n \n # Intermediate MLP\n for i in range(len(self.mlp_intermediate_units)):\n x3 = self.intermediate_mlp_dense_layers[i](x3, training=training)\n x3 = self.intermediate_mlp_dropout_layers[i](x3, training=training)\n \n # Skip Connection 2\n encoded = self.add_2_layer([x3, x2], training=training)\n \n return encoded","metadata":{"id":"JXT9kwwZnrQB","execution":{"iopub.status.busy":"2022-01-25T14:54:24.714233Z","iopub.execute_input":"2022-01-25T14:54:24.714479Z","iopub.status.idle":"2022-01-25T14:54:24.732374Z","shell.execute_reply.started":"2022-01-25T14:54:24.71445Z","shell.execute_reply":"2022-01-25T14:54:24.731564Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with tf.device('/CPU:0'):\n ViT = ViTEncoder(patch_size=demo_patch_size, projection_dim=demo_projection_dim)\n demo_encoder_output= ViT(tf.ones((1, *IMG_SHAPE)))\n IMG_SEQ_LEN, IMG_EMB_DEPTH = demo_encoder_output.shape[1], demo_encoder_output.shape[2]\n \nprint(f\"Encoder Output Shape: {demo_encoder_output.shape}\") \nprint(f\"Output 'Sequence' Length: {IMG_SEQ_LEN}\")\nprint(f\"Output 'Sequence' Feature Depth: {IMG_EMB_DEPTH}\")\n\n# PATCH LAYER FUNCTIONALITY\n\n# 1. Plot the original image\nprint(\"\\n\\n... ORIGINAL IMAGE ...\\n\")\nplt.figure(figsize=(18,9))\nplt.imshow(demo_img[0].numpy().astype(\"float32\")/255.)\nplt.axis('off')\nplt.tight_layout()\nplt.show()\n\n# 2. Plot the patches in the same shape/order as the original image\nprint(\"\\n\\n\\n... IMAGE PATCHES ...\\n\")\nplt.figure(figsize=(18,9))\nfor i, patch in enumerate(patches[0]):\n w,h = int(np.ceil(IMG_SHAPE[0]/demo_patch_size)), int(np.ceil(IMG_SHAPE[1]/demo_patch_size))\n ax = plt.subplot(w,h,i+1)\n patch_img = tf.reshape(patch, (demo_patch_size, demo_patch_size, 3))\n plt.imshow(patch_img.numpy().astype(\"float32\")/255.)\n plt.axis(\"off\")\nplt.tight_layout()\nplt.subplots_adjust(wspace=0.03, hspace=0.06)\nplt.show()\n\nwith tf.device('/CPU:0'):\n enc_patches_rescaled = patch_creator(tf.expand_dims(tf.image.resize(tf.reshape(tf.reduce_mean(encoded_patches[0], axis=-1), (w, h, 1)), (IMG_SHAPE[0], IMG_SHAPE[1])), axis=0))\n\n# 3. Plot the image embedding patches\nprint(\"\\n\\n\\n... IMAGE EMBEDDING AS PATCH VISUALIZATION ...\\n\")\nplt.figure(figsize=(18,9))\nfor i, enc_patch in enumerate(enc_patches_rescaled[0]):\n ax = plt.subplot(w,h,i+1)\n enc_patch_img = tf.reshape(enc_patch, (demo_patch_size, demo_patch_size))\n plt.imshow(enc_patch_img.numpy().astype(\"float32\")/255., cmap=\"jet\")\n plt.axis(\"off\")\nplt.tight_layout()\nplt.subplots_adjust(wspace=0.03, hspace=0.06)\nplt.show()\n\nwith tf.device('/CPU:0'):\n tmp_img = tf.image.resize(tf.reshape(tf.reduce_mean(demo_encoder_output[0], axis=-1), (w, h, 1)), (IMG_SHAPE[0], IMG_SHAPE[1]))\n tmp_img = tf.cast(255*((tmp_img-tf.math.reduce_min(tmp_img))/(tf.math.reduce_max(tmp_img)-tf.math.reduce_min(tmp_img))), dtype=tf.uint8)\n enc_output_patches_rescaled = patch_creator(tf.expand_dims(tmp_img, axis=0))\n\n# 4. Plot the encoder output as patches\nprint(\"\\n\\n\\n... ENCODER OUTPUT AS PATCH VISUALIZATION ...\\n\")\nplt.figure(figsize=(18,9))\nfor i, enc_out_patch in enumerate(enc_output_patches_rescaled[0]):\n ax = plt.subplot(w,h,i+1)\n enc_out_patch_img = tf.reshape(enc_out_patch, (demo_patch_size, demo_patch_size))\n plt.imshow(enc_out_patch_img.numpy().astype(\"float32\")/255., cmap=\"jet\")\n plt.axis(\"off\")\nplt.tight_layout()\nplt.subplots_adjust(wspace=0.03, hspace=0.06)\nplt.show()\n\nprint(\"\\n\\n\\n\\n\\tSUMMARY\\n\")\nViT.summary()","metadata":{"papermill":{"duration":44.15573,"end_time":"2021-06-15T18:02:44.288255","exception":false,"start_time":"2021-06-15T18:02:00.132525","status":"completed"},"tags":[],"id":"fuzzy-raleigh","outputId":"bfa430fc-4729-4a00-dfb2-63446fc3e31e","execution":{"iopub.status.busy":"2022-01-25T14:54:24.733704Z","iopub.execute_input":"2022-01-25T14:54:24.734076Z","iopub.status.idle":"2022-01-25T14:55:09.965619Z","shell.execute_reply.started":"2022-01-25T14:54:24.734044Z","shell.execute_reply":"2022-01-25T14:55:09.964672Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 5.2 UNDERSTANDING THE MODELS - TRANSFORMER","metadata":{"id":"pkyINyBI3men"}},{"cell_type":"markdown","source":"### 5.2.0 TRANSFORMER - HYPERPARAMETERS","metadata":{"id":"z9s9RtXK3w8T"}},{"cell_type":"code","source":"D_MODEL = IMG_EMB_DEPTH\nN_PE_POS = 72\nD_FF = 1024\n\nprint(f\"\\n... THE INPUT 'SEQUENCE LENGTH' IS {IMG_SEQ_LEN} (output of image encoder - shape flattened) ...\")\nprint(f\"... THE INPUT 'EMBEDDING DEPTH' IS {IMG_EMB_DEPTH} (output of image encoder - # of channels) ...\")\nprint(f\"... THE NUMBER OF POSITIONAL ENCODING POSITIONS IS {N_PE_POS} (arbitray) ...\\n\")","metadata":{"papermill":{"duration":0.128982,"end_time":"2021-06-15T18:02:45.016313","exception":false,"start_time":"2021-06-15T18:02:44.887331","status":"completed"},"tags":[],"id":"gross-submission","outputId":"871343cd-0b5d-4932-ccc8-dc8aae7c25d2","execution":{"iopub.status.busy":"2022-01-25T14:55:09.969014Z","iopub.execute_input":"2022-01-25T14:55:09.969272Z","iopub.status.idle":"2022-01-25T14:55:09.97608Z","shell.execute_reply.started":"2022-01-25T14:55:09.969243Z","shell.execute_reply":"2022-01-25T14:55:09.975162Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5.2.1 TRANSFORMER - POSITIAL ENCODING\n\n---\n\nSince this model doesn't contain any recurrence or convolution, positional encoding is added to give the model some information about the relative position of the words in the sentence. \n\nThe positional encoding vector is added to the embedding vector. \n* Embeddings represent a token in a **`d-dimensional`** space where tokens (encoded vectors) with similar meaning (feature representation) will be closer to each other. \n\nBut the embeddings do not encode the relative position of words in a sentence (or in our case the localization of features as encoded by our **efficientnetv2 encoder model**).\n* So after adding the positional encoding, words (feature representations) will be closer to each other based on the ***similarity of their meaning and their position in the sentence (feature vector)***, in the **`d-dimensional`** space.\n\nSee the notebook on **[positional encoding](https://www.tensorflow.org/tutorials/text/transformer#positional_encoding)** to learn more about it. The formula for calculating the positional encoding is as follows:\n\n---\n\n$$\\Large{PE_{(pos, 2i)} = sin(pos / 10000^{2i / d_{model}})} $$\n$$\\Large{PE_{(pos, 2i+1)} = cos(pos / 10000^{2i / d_{model}})} $$\n\n---","metadata":{"papermill":{"duration":0.119338,"end_time":"2021-06-15T18:02:45.25713","exception":false,"start_time":"2021-06-15T18:02:45.137792","status":"completed"},"tags":[],"id":"radio-oliver"}},{"cell_type":"code","source":"def get_angles(pos, i, d_model):\n ''' Function to calculate the angle between feature representations ''' \n angle_rates = tf.constant(1, TARGET_DTYPE) / tf.math.pow(tf.constant(10000, TARGET_DTYPE), \n (tf.constant(2, dtype=TARGET_DTYPE) * tf.cast((i//2), \n TARGET_DTYPE))/d_model)\n return pos * angle_rates\n\ndef do_interleave(arr_a, arr_b):\n ''' Function to perform interleaving '''\n a_arr_tf_column = tf.range(arr_a.shape[1])*2 # [0 2 4 ...]\n b_arr_tf_column = tf.range(arr_b.shape[1])*2+1 # [1 3 5 ...]\n column_indices = tf.argsort(tf.concat([a_arr_tf_column,b_arr_tf_column],axis=-1))\n column, row = tf.meshgrid(column_indices,tf.range(arr_a.shape[0]))\n combine_indices = tf.stack([row,column],axis=-1)\n combine_value = tf.concat([arr_a,arr_b],axis=1)\n return tf.gather_nd(combine_value,combine_indices)\n\ndef positional_encoding_1d(position, d_model):\n ''' Function to calculate the positional encodings for 1-D data '''\n angle_rads = get_angles(tf.cast(tf.range(position)[:, tf.newaxis], TARGET_DTYPE),\n tf.cast(tf.range(d_model)[tf.newaxis, :], TARGET_DTYPE),\n d_model)\n \n # apply sin to even indices in the array; 2i\n sin_angle_rads = tf.math.sin(angle_rads[:, ::2])\n cos_angle_rads = tf.math.cos(angle_rads[:, 1::2])\n angle_rads = do_interleave(sin_angle_rads, cos_angle_rads)\n pos_encoding = angle_rads[tf.newaxis, ...]\n return pos_encoding\n\ndef np_positional_encoding_2d(row, col, d_model):\n assert d_model % 2 == 0\n row_pos = np.repeat(np.arange(row),col)[:,np.newaxis]\n col_pos = np.repeat(np.expand_dims(np.arange(col),0),row,axis=0).reshape(-1,1)\n\n angle_rads_row = get_angles(row_pos,np.arange(d_model//2)[np.newaxis,:],d_model//2).numpy()\n angle_rads_col = get_angles(col_pos,np.arange(d_model//2)[np.newaxis,:],d_model//2).numpy()\n \n angle_rads_row[:, 0::2] = np.sin(angle_rads_row[:, 0::2])\n angle_rads_row[:, 1::2] = np.cos(angle_rads_row[:, 1::2])\n angle_rads_col[:, 0::2] = np.sin(angle_rads_col[:, 0::2])\n angle_rads_col[:, 1::2] = np.cos(angle_rads_col[:, 1::2])\n pos_encoding = np.concatenate([angle_rads_row,angle_rads_col],axis=1)[np.newaxis, ...]\n return tf.cast(pos_encoding, dtype=TARGET_DTYPE)\n\ndef positional_encoding_2d(row, col, d_model):\n row_pos = tf.repeat(tf.range(row), col)[:, tf.newaxis]\n col_pos = tf.reshape(tf.repeat(tf.expand_dims(tf.range(col),0), row, axis=0), (-1, 1))\n\n angle_rads_row = get_angles(tf.cast(row_pos, tf.float32), tf.range(d_model//2)[tf.newaxis,:], d_model//2)\n angle_rads_col = get_angles(tf.cast(col_pos, tf.float32), tf.range(d_model//2)[tf.newaxis,:], d_model//2)\n\n sin_angle_rads_row = tf.math.sin(angle_rads_row[:, ::2])\n cos_angle_rads_row = tf.math.cos(angle_rads_row[:, 1::2])\n angle_rads_row = do_interleave(sin_angle_rads_row, cos_angle_rads_row)\n\n sin_angle_rads_col = tf.math.sin(angle_rads_col[:, ::2])\n cos_angle_rads_col = tf.math.cos(angle_rads_col[:, 1::2])\n angle_rads_col = do_interleave(sin_angle_rads_col, cos_angle_rads_col)\n \n pos_encoding = tf.concat([angle_rads_row,angle_rads_col],axis=1)[tf.newaxis, ...]\n return pos_encoding\n\npos_encoding = positional_encoding_1d(256, 512)\n\nprint(pos_encoding.shape)\n\nplt.figure(figsize=(6,4))\nplt.pcolormesh(tf.cast(pos_encoding[0], tf.float32), cmap='RdBu')\nplt.xlim((0, 512))\nplt.ylim((0, 256))\nplt.xlabel('Depth', fontweight=\"bold\")\nplt.ylabel('Position', fontweight=\"bold\")\nplt.title(\"Visualization of Positional Encoding\", fontweight=\"bold\")\nplt.colorbar()\nplt.show()","metadata":{"papermill":{"duration":0.478163,"end_time":"2021-06-15T18:02:45.854617","exception":false,"start_time":"2021-06-15T18:02:45.376454","status":"completed"},"tags":[],"id":"headed-copper","outputId":"4132ddf6-1004-4586-b5e1-becd886c6c57","execution":{"iopub.status.busy":"2022-01-25T14:55:09.977583Z","iopub.execute_input":"2022-01-25T14:55:09.977915Z","iopub.status.idle":"2022-01-25T14:55:10.450755Z","shell.execute_reply.started":"2022-01-25T14:55:09.977854Z","shell.execute_reply":"2022-01-25T14:55:10.449886Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5.2.2 TRANSFORMER - MASKING\n\n---\n\n\nMask all the pad tokens in the batch of sequence. It ensures that the model does not treat padding as the input. \n\nThe mask indicates where pad value **`0`** is present: \n* it outputs a **`1`** at those locations\n* it outputs a **`0`** otherwise.\n\n---\n\nThe **look-ahead mask** is used to mask the future tokens in a sequence. \n\nIn other words, the mask indicates which entries should not be used.\n* This means that to predict the third token, only the first and second tokens will be used. \n* Similarly to predict the fourth token, only the first, second and the third tokens will be used and so on.","metadata":{"papermill":{"duration":0.12377,"end_time":"2021-06-15T18:02:46.10383","exception":false,"start_time":"2021-06-15T18:02:45.98006","status":"completed"},"tags":[],"id":"horizontal-journalism"}},{"cell_type":"code","source":"def create_padding_mask(seq):\n ''' \n Function to add extra dimensions to add the padding to the attention logits.\n - (batch_size, 1, 1, seq_len)\n '''\n seq = tf.cast(tf.math.equal(seq, 0), TARGET_DTYPE)\n return seq[:, tf.newaxis, tf.newaxis, :]\n\ndef create_look_ahead_mask(size):\n ''' Function to create a look ahead mask '''\n mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n # (seq_len, seq_len)\n return tf.cast(mask, TARGET_DTYPE)\n\ndef create_mask(inp, tar):\n '''\n Function to combine the look-ahead and padding masks to be used in the 1st attention block in the decoder.\n It is used to pad and mask future tokens in the input received by the decoder.\n '''\n look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])\n dec_target_padding_mask = create_padding_mask(tar)\n combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)\n return tf.cast(combined_mask, TARGET_DTYPE)\n\nx = tf.constant([[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]])\nprint(f\"x --\\n{x}\\n\\n\\nPADDING MASK\\n{create_padding_mask(x)}\\n\")\nprint(\"\")\n\nx = tf.random.uniform((1, 5))\nprint(f\"x.shape[1] -- {x.shape[1]}\")\nprint(f\"\\n\\nLOOK-AHEAD MASK\\n{create_look_ahead_mask(x.shape[1])}\\n\")","metadata":{"papermill":{"duration":0.160001,"end_time":"2021-06-15T18:02:46.387914","exception":false,"start_time":"2021-06-15T18:02:46.227913","status":"completed"},"tags":[],"id":"strategic-switch","outputId":"ccf4cf91-105a-4b5b-9093-0edf69638b45","execution":{"iopub.status.busy":"2022-01-25T14:55:10.452237Z","iopub.execute_input":"2022-01-25T14:55:10.45311Z","iopub.status.idle":"2022-01-25T14:55:10.488298Z","shell.execute_reply.started":"2022-01-25T14:55:10.453063Z","shell.execute_reply":"2022-01-25T14:55:10.487289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5.2.3 TRANSFORMER - SCALED DOT-PRODUCT ATTENTION","metadata":{"id":"jEzYMNV14ES0"}},{"cell_type":"markdown","source":"\n\n---\n\nScaled dot-product attention is an attention mechanism where the dot products are scaled down by $\\sqrt{d_k}$. \n\n---\n\n
![\"scaled_dot_product_attention\"](\"https://www.tensorflow.org/images/tutorials/transformer/scaled_attention.png\")
\n\nThis layer (the MHA layer) first projects **`query`**, **`key`** and **`value`**. \n* These are (effectively) a list of tensors of length num_attention_heads, where the corresponding shapes are: \n * **`[batch_size, 1, key_dim]`**, **`[batch_size, 1, key_dim]`**, **`[batch_size, 1, value_dim]`**\n\nThen, the **`query`** and **`key`** tensors are **dot-producted and scaled** (see previous section). These values are softmaxed to obtain attention probabilities. The tensors are then interpolated by these probabilities, then concatenated back to a single tensor.\n\nFinally, the result tensor with the last dimension as value_dim can take an linear projection and return.\n\n---\n\n
\n\n
![\"multi-head](\"https://www.tensorflow.org/images/tutorials/transformer/multi_head_attention.png\")
\n\n\n**Multi-head attention consists of four parts:**\n* Linear layers and split into heads.\n* Scaled dot-product attention.\n* Concatenation of heads.\n* Final linear layer.\n\n---\n\nEach multi-head attention block gets three inputs;\n* **`Q` (query)**\n* **`K` (key)**\n* **`V` (value)**\n\nThese are put through linear (**`Dense`**) layers and split up into multiple heads. \n\nThe **`scaled_dot_product_attention`** defined above is applied to each head (broadcasted for efficiency). \n* An appropriate mask must be used in the attention step. \n* The attention output for each head is then concatenated (using **`tf.transpose`**, and **`tf.reshape`**) and put through a final **`Dense`** layer\n\nInstead of one single attention head, **`Q`**, **`K`**, and **`V`** are split into multiple heads because it allows the model to jointly attend to information at different positions from different representational spaces. \n\nAfter the split each head has a reduced dimensionality, so the total computation cost is the same as a single head attention with full dimensionality.\n\n---\n\nLet's create a **`MultiHeadAttention`** layer to try out. \n* At each location in the sequence, **`y`**, the **`MultiHeadAttention`** runs all **`8`** attention heads across all other locations in the sequence, returning a new vector of the same length at each location.","metadata":{"papermill":{"duration":0.123639,"end_time":"2021-06-15T18:02:47.216045","exception":false,"start_time":"2021-06-15T18:02:47.092406","status":"completed"},"tags":[],"id":"respective-soundtrack"}},{"cell_type":"code","source":"class MultiHeadAttention(tf.keras.layers.Layer):\n ''' MultiHead Attention Layer Component of the Transformer '''\n def __init__(self, d_model, num_heads):\n '''\n Args:\n d_model (int): Depth of the d-dimensional space used for positional encoding\n num_heads (int): The number of heads to use in the multi-head-attention block\n\n Returns:\n None, this is initialization \n '''\n super(MultiHeadAttention, self).__init__()\n self.num_heads = num_heads\n self.d_model = d_model\n assert d_model % self.num_heads == 0\n \n self.depth = d_model // self.num_heads\n self.wq = tf.keras.layers.Dense(d_model)\n self.wk = tf.keras.layers.Dense(d_model)\n self.wv = tf.keras.layers.Dense(d_model)\n \n self.dense = tf.keras.layers.Dense(d_model)\n \n def split_heads(self, x, batch_size):\n \"\"\" \n Function to split the last dimension into (num_heads, depth).\n Then we transpose the result such that the shape is \n - (batch_size, num_heads, seq_len, depth)\n \"\"\"\n x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n return tf.transpose(x, perm=[0, 2, 1, 3])\n \n def call(self, v, k, q, mask):\n batch_size = tf.shape(q)[0]\n \n # (batch_size, seq_len, d_model)\n q = self.wq(q) \n # (batch_size, seq_len, d_model)\n k = self.wk(k) \n # (batch_size, seq_len, d_model)\n v = self.wv(v) \n\n # (batch_size, num_heads, seq_len_q, depth)\n q = self.split_heads(q, batch_size) \n # (batch_size, num_heads, seq_len_k, depth)\n k = self.split_heads(k, batch_size) \n # (batch_size, num_heads, seq_len_v, depth)\n v = self.split_heads(v, batch_size) \n \n # scaled_attention.shape – (batch_size, num_heads, seq_len_q, depth)\n # attention_weights.shape – (batch_size, num_heads, seq_len_q, seq_len_k)\n scaled_attention, attention_weights = scaled_dot_product_attention(q, k, v, mask)\n \n # (batch_size, seq_len_q, num_heads, depth)\n scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3]) \n\n # (batch_size, seq_len_q, d_model)\n concat_attention = tf.reshape(scaled_attention, \n (batch_size, -1, self.d_model)) \n\n # (batch_size, seq_len_q, d_model)\n output = self.dense(concat_attention) \n \n return output, attention_weights\n\n# CUSTOM\ntemp_mha = MultiHeadAttention(d_model=512, num_heads=8)\ny = tf.random.uniform((1, 60, 512), dtype=TARGET_DTYPE) # (batch_size, encoder_sequence, d_model)\nout, attn = temp_mha(y, y, y, mask=None)\nprint(f\"Custom MHA Layer:\\n\\t-->{out[0,:2]}\\n\\t-->{(out.shape, attn.shape)}\\n\")\n\n# TF NATIVE\ntemp_mha = tf.keras.layers.MultiHeadAttention(8, 512)\nout, attn = temp_mha(y, y, y, attention_mask=None, return_attention_scores=True)\nprint(f\"TF MHA Layer:\\n\\t-->{out[0,:2]}\\n\\t-->{(out.shape, attn.shape)}\\n\")\n\ndel temp_mha","metadata":{"papermill":{"duration":0.428295,"end_time":"2021-06-15T18:02:47.767688","exception":false,"start_time":"2021-06-15T18:02:47.339393","status":"completed"},"tags":[],"id":"prompt-services","outputId":"6194201b-1fb5-4af1-8ac1-ce93b4fa3a9c","execution":{"iopub.status.busy":"2022-01-25T14:55:10.552213Z","iopub.execute_input":"2022-01-25T14:55:10.552452Z","iopub.status.idle":"2022-01-25T14:55:10.8476Z","shell.execute_reply.started":"2022-01-25T14:55:10.552419Z","shell.execute_reply":"2022-01-25T14:55:10.846653Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5.2.5 TRANSFORMER - POINT-WISE FEED FORWARD NEURAL NETWORK\n\n---\n\nPoint wise feed forward network consists of two fully-connected layers with a ReLU activation in-between.","metadata":{"papermill":{"duration":0.126863,"end_time":"2021-06-15T18:02:48.028264","exception":false,"start_time":"2021-06-15T18:02:47.901401","status":"completed"},"tags":[],"id":"structural-german"}},{"cell_type":"code","source":"def point_wise_feed_forward_network(d_model, dff):\n '''\n Args:\n d_model (int): Depth of the d-dimensional space used for positional encoding\n dff (int): Number of units to use in the feed-forward neural network \n \n Returns:\n Feedforward neural network \n '''\n return tf.keras.Sequential([\n \n # INNER LAYER\n # – (batch_size, seq_len, dff)\n tf.keras.layers.Dense(dff, activation='relu'), \n \n # OUTPUT \n # – (batch_size, seq_len, d_model)\n tf.keras.layers.Dense(d_model) \n ])\n\nsample_ffn = point_wise_feed_forward_network(512, D_FF)\nprint(\"\\nFFN INPUT & OUTPUT SHAPE: \" \\\n f\"{sample_ffn(tf.random.uniform((64, 50, 512), dtype=TARGET_DTYPE)).shape}\" \\\n \"\\n\\nFFN SUMMARY:\")\nprint(sample_ffn.summary())\n\ndel sample_ffn","metadata":{"papermill":{"duration":0.20341,"end_time":"2021-06-15T18:02:48.35655","exception":false,"start_time":"2021-06-15T18:02:48.15314","status":"completed"},"tags":[],"id":"turkish-kentucky","outputId":"4e0e9e4d-27a9-4d8f-dc95-a95394dcbc12","execution":{"iopub.status.busy":"2022-01-25T14:55:10.849087Z","iopub.execute_input":"2022-01-25T14:55:10.849489Z","iopub.status.idle":"2022-01-25T14:55:10.922863Z","shell.execute_reply.started":"2022-01-25T14:55:10.849441Z","shell.execute_reply":"2022-01-25T14:55:10.921761Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5.2.6 TRANSFORMER - ENCODER-DECODER NETWORK ARCHITECTURE OVERVIEW\n---\n\nThe transformer model follows the same general pattern as a standard [sequence to sequence with attention model](nmt_with_attention.ipynb). \n\n* The input sequennce ***(image embedding sequence in our case)*** is passed through **`N` encoder layers** that generates an output for each word/token in the sequence.\n* The **decoder** attends on the encoder's output and its own input (self-attention) to predict the next word/token. \n\n---\n\n
![\"transformer\"](\"https://www.tensorflow.org/images/tutorials/transformer/transformer.png\")
\n","metadata":{"papermill":{"duration":0.124383,"end_time":"2021-06-15T18:02:48.606746","exception":false,"start_time":"2021-06-15T18:02:48.482363","status":"completed"},"tags":[],"id":"dental-racing"}},{"cell_type":"markdown","source":"### 5.2.7 TRANSFORMER - ENCODER\n\n---\n\n**Each transformer encoder layer consists of sublayers:**\n\n1. **Multi-Head AAttention (with padding mask)**\n2. **Point-Wise Feed Forward Neural Networks**\n\nEach of these sublayers has a **residual connection** around it followed by a **layer normalization**. \n* Residual connections help in avoiding the vanishing gradient problem in deep networks.\n\nThe output of each sublayer is **`LayerNorm(x + Sublayer(x))`**. \n* The normalization is done on the **`d_model`** (last) axis. \n* There are **`N encoder` layers** in the ***transformer***.","metadata":{"papermill":{"duration":0.127458,"end_time":"2021-06-15T18:02:48.859196","exception":false,"start_time":"2021-06-15T18:02:48.731738","status":"completed"},"tags":[],"id":"incredible-billion"}},{"cell_type":"code","source":"class TransformerEncoderLayer(tf.keras.layers.Layer):\n def __init__(self, d_model, num_heads, dff, dropout_rate=0.1):\n \"\"\" \n Encoder Layer Component Of Transformer Encoder Block \n \n Args:\n d_model (int): Depth of the d-dimensional space used for positional encoding\n num_heads (int): The number of heads to use in the multi-head-attention block\n dff (int): Number of units to use in the feed-forward neural network\n drop_out_rate (float): Percentage of nodes to drop in a given layer\n \n Returns:\n None; This is an intiailization\n \"\"\"\n super(TransformerEncoderLayer, self).__init__()\n\n self.mha = tf.keras.layers.MultiHeadAttention(num_heads, key_dim=d_model,)\n self.ffn = point_wise_feed_forward_network(d_model, dff)\n\n self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n \n self.dropout1 = tf.keras.layers.Dropout(dropout_rate)\n self.dropout2 = tf.keras.layers.Dropout(dropout_rate)\n \n def call(self, x, training, mask=None):\n \"\"\" \n Call function for our encoder layer\n \n Args:\n x (array): Input token embeddinng\n training (bool): Whether or not to apply certain operations (i.e. disable/enable dropout)\n mask (tensor): None if no masks are used in the MultiHead Attention layer\n \n Returns:\n The encoded input sequence\n - shape --> (batch_size, input_seq_len, d_model)\n \"\"\"\n \n # returns --> (batch_size, input_seq_len, d_model)\n attn_output, _ = self.mha(x, x, x, mask, return_attention_scores=True) \n\n # Potentially unncessary by passing dropout1 to tf.keras.layers.MultiHeadAttention (if using tf MHA)\n attn_output = self.dropout1(attn_output, training=training)\n \n # Residual connection followed by layer normalization\n # – returns --> (batch_size, input_seq_len, d_model)\n out1 = self.layernorm1(x + attn_output, training=training) \n \n # Point-wise Feed Forward Step\n # – returns --> (batch_size, input_seq_len, d_model)\n ffn_output = self.ffn(out1, training=training) \n ffn_output = self.dropout2(ffn_output, training=training)\n \n\n # Residual connection followed by layer normalization\n # – returns --> (batch_size, input_seq_len, d_model)\n out2 = self.layernorm2(out1 + ffn_output, training=training) \n \n return out2\n\nsample_encoder_layer = TransformerEncoderLayer(D_MODEL, 8, D_FF)\nsample_encoder_layer_output = sample_encoder_layer(demo_encoder_output, training=False, mask=None)\ndel sample_encoder_layer\n\n# (batch_size, input_seq_len, d_model)\nsample_encoder_layer_output.shape ","metadata":{"papermill":{"duration":0.281209,"end_time":"2021-06-15T18:02:49.265027","exception":false,"start_time":"2021-06-15T18:02:48.983818","status":"completed"},"tags":[],"id":"cleared-democracy","outputId":"4c901366-7500-4169-986b-10b50d1bafcb","execution":{"iopub.status.busy":"2022-01-25T14:55:10.924364Z","iopub.execute_input":"2022-01-25T14:55:10.924621Z","iopub.status.idle":"2022-01-25T14:55:11.077138Z","shell.execute_reply.started":"2022-01-25T14:55:10.924591Z","shell.execute_reply":"2022-01-25T14:55:11.076014Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5.2.8 TRANSFORMER - DECODER LAYER COMPONENT\n\n---\n\n**Each transformer decoder layer consists of sublayers:**\n\n1. **Masked Multi-Head Attention (with look ahead mask and padding mask)**\n2. **Multi-Head Attention (with padding mask)** \n * **`V`** (value) and **`K`** (key) receive the ***encoder output*** as inputs. \n * **`Q`** (query) receives the ***output from the masked multi-head attention sublayer.***\n3. **Point-Wise Feed Forward Networks**\n\nEach of these sublayers has a **residual connection** around it followed by a **layer normalization**\n* The output of each sublayer is **`LayerNorm(x + Sublayer(x))`**\n* The normalization is done on the **`d_model`** (last) axis.\n\nThere are **`N` decoder layers** in the ***transformer***\n\nAs **`Q`** receives the output from decoder's first attention block, and **`K`** receives the encoder output, the attention weights represent the importance given to the decoder's input based on the encoder's output. \n* In other words, the decoder predicts the next word/token by looking at the encoder output and self-attending to its own output. \n* See the demonstration above in the scaled dot product attention section.","metadata":{"papermill":{"duration":0.130418,"end_time":"2021-06-15T18:02:49.522845","exception":false,"start_time":"2021-06-15T18:02:49.392427","status":"completed"},"tags":[],"id":"ancient-corpus"}},{"cell_type":"code","source":"class TransformerDecoderLayer(tf.keras.layers.Layer):\n def __init__(self, d_model, num_heads, dff, dropout_rate=0.1):\n \"\"\" \n Decoder Layer Component Of Transformer Block \n \n Args:\n d_model (int): Depth of the d-dimensional space used for positional encoding of image embedding\n num_heads (int): Number of heads to use in the multi-head-attention block\n dff (int): Number of units to use in the feed-forward neural network\n dropout_rate (float): Percentage of nodes to drop in a given layer\n \n Returns:\n None; This is an intiailization\n \"\"\"\n super(TransformerDecoderLayer, self).__init__()\n\n # WE COULD USE A CUSTOM DEFINED MHA MODEL BUT WE WILL USE TFA INSTEAD\n self.mha1 = MultiHeadAttention(d_model, num_heads)\n self.mha2 = MultiHeadAttention(d_model, num_heads)\n #\n # # Multi Head Attention Layers\n # self.mha1 = tf.keras.layers.MultiHeadAttention(num_heads, key_dim=d_model,)\n # self.mha2 = tf.keras.layers.MultiHeadAttention(num_heads, key_dim=d_model,)\n\n # Feed Forward NN\n self.ffn = point_wise_feed_forward_network(d_model, dff)\n \n # Layer Normalization Layers\n self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n \n # Dropout Layers\n self.dropout1 = tf.keras.layers.Dropout(dropout_rate)\n self.dropout2 = tf.keras.layers.Dropout(dropout_rate)\n self.dropout3 = tf.keras.layers.Dropout(dropout_rate)\n \n # enc_output.shape == (batch_size, input_seq_len, d_model)\n def call(self, x, enc_output, training, look_ahead_mask=None, padding_mask=None):\n \"\"\" \n Call function for our encoder layer\n \n Args:\n x (array): token embeddinng (batch_size, output_seq_len, embedding_dim)\n enc_output (array): token embeddinng (batch_size, output_seq_len, embedding_dim)\n training (bool): Whether or not to apply certain operations (i.e. disable/enable dropout)\n look_ahead_mask (array): None if no look ahead mask is used\n padding_mask (array): None if no padding mask is used \n \n Returns:\n The encoded input sequence\n - shape --> (batch_size, input_seq_len, d_model)\n \"\"\"\n \n attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)\n attn1 = self.dropout1(attn1, training=training)\n\n # Residual connection followed by layer normalization\n # – (batch_size, target_seq_len, d_model)\n out1 = self.layernorm1(attn1 + x, training=training)\n \n # Merging connection between encoder and decoder (MHA)\n # – (batch_size, target_seq_len, d_model)\n attn2, attn_weights_block2 = self.mha2(enc_output, enc_output, out1, padding_mask) \n attn2 = self.dropout2(attn2, training=training)\n \n # Residual connection followed by layer normalization\n # – (batch_size, target_seq_len, d_model)\n out2 = self.layernorm2(attn2 + out1, training=training) \n \n # (batch_size, target_seq_len, d_model)\n ffn_output = self.ffn(out2, training=training) \n ffn_output = self.dropout3(ffn_output, training=training)\n\n # Residual connection followed by layer normalization\n # – (batch_size, target_seq_len, d_model)\n out3 = self.layernorm3(ffn_output + out2, training=training) \n \n return out3, attn_weights_block1, attn_weights_block2\n\nsample_decoder_layer = TransformerDecoderLayer(D_MODEL, 8, D_FF)\nsample_decoder_layer_output, _, _ = sample_decoder_layer(tf.random.uniform((BATCH_SIZE_DEBUG, MAX_LEN, D_MODEL), dtype=TARGET_DTYPE), sample_encoder_layer_output, False, None, None)\ndel sample_decoder_layer\n\nsample_decoder_layer_output.shape # (batch_size, target_seq_len, d_model)","metadata":{"papermill":{"duration":1.096462,"end_time":"2021-06-15T18:02:50.743125","exception":false,"start_time":"2021-06-15T18:02:49.646663","status":"completed"},"tags":[],"id":"dated-costume","outputId":"59309c27-b8a8-444c-f258-e194f98f06e6","execution":{"iopub.status.busy":"2022-01-25T14:55:11.078756Z","iopub.execute_input":"2022-01-25T14:55:11.079312Z","iopub.status.idle":"2022-01-25T14:55:12.06488Z","shell.execute_reply.started":"2022-01-25T14:55:11.079266Z","shell.execute_reply":"2022-01-25T14:55:12.063833Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5.2.9 TRANSFORMER - ENCODER COMPONENT\n\n---\n\nThe **`TransformerEncoder`** consists of:\n1. Input Embedding\n2. Positional Encoding\n3. **`N`** encoder layers\n\n
\n\nThe input is put through an embedding which is summed with the positional encoding. \n* The output of this summation is the input to the encoder layers. The output of the encoder is the input to the decoder.","metadata":{"papermill":{"duration":0.125895,"end_time":"2021-06-15T18:02:50.997831","exception":false,"start_time":"2021-06-15T18:02:50.871936","status":"completed"},"tags":[],"id":"dominican-sampling"}},{"cell_type":"code","source":"class TransformerEncoder(tf.keras.layers.Layer):\n ''' Encoder Component of the Transformer Block '''\n def __init__(self, num_layers, d_model, num_heads, dff,\n maximum_position_encoding, dropout_rate=0.1):\n '''\n Args:\n num_layers (int): Number of encoder layers to be used\n d_model (int): Depth of the d-dimensional space used for positional encoding of image embedding\n num_heads (int): Number of heads to use in the multi-head-attention block\n dff (int): Number of units to use in the feed-forward neural network\n maximum_positional_encoding (int): Maximum limit for positional encoding\n dropout_rate (float): Percentage of nodes to drop in a given layer\n \n Returns:\n None; This is an intiailization\n '''\n super(TransformerEncoder, self).__init__()\n\n self.d_model = d_model\n self.num_layers = num_layers\n self.pos_encoding = positional_encoding_1d(maximum_position_encoding, self.d_model)\n self.enc_layers = [TransformerEncoderLayer(d_model, num_heads, dff, dropout_rate) for _ in range(num_layers)]\n self.dropout = tf.keras.layers.Dropout(dropout_rate)\n \n def call(self, x, training, mask=None):\n \"\"\"\n Sequence of Operations:\n 1. Embed the input data as a fixed length vector\n 2. Scale the fixed length vector by the square root of the \n input/output dimensionality\n 3. Introduce the position encoding into the data\n 4. Perform some amount of dropout\n 5. Pass our preprocessed input data into a stack of encoding layers\n along with the input mask\n \"\"\"\n\n # adding embedding and position encoding.\n # – (batch_size, input_seq_len, d_model)\n x += self.pos_encoding\n x = self.dropout(x, training=training)\n \n for i in range(self.num_layers):\n x = self.enc_layers[i](x, training, mask)\n \n # – (batch_size, input_seq_len, d_model)\n return x \n\nsample_encoder = TransformerEncoder(num_layers=2, \n d_model=D_MODEL, \n num_heads=8, \n dff=D_FF,\n maximum_position_encoding=IMG_SEQ_LEN)\n\nsample_encoder_output = sample_encoder(demo_encoder_output, training=False, mask=None)\ndel sample_encoder\n\nprint(sample_encoder_output.shape) # (batch_size, input_seq_len, d_model)","metadata":{"papermill":{"duration":0.426684,"end_time":"2021-06-15T18:02:51.550018","exception":false,"start_time":"2021-06-15T18:02:51.123334","status":"completed"},"tags":[],"id":"resistant-curtis","outputId":"b8a4cace-e228-4e32-e8b4-da6c7d32b32e","execution":{"iopub.status.busy":"2022-01-25T14:55:12.066569Z","iopub.execute_input":"2022-01-25T14:55:12.067424Z","iopub.status.idle":"2022-01-25T14:55:12.370864Z","shell.execute_reply.started":"2022-01-25T14:55:12.067365Z","shell.execute_reply":"2022-01-25T14:55:12.369986Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5.2.10 TRANSFORMER - DECODER COMPONENT\n\n---\n\n1. Output Embedding\n2. Positional Encoding\n3. **`N`** decoder layers\n\n
\n\nThe target is put through an **embedding** which is **summed with the positional encoding.**\n* The output of this summation is the input to the decoder layers. \n* The output of the decoder is the input to the final linear layer.","metadata":{"papermill":{"duration":0.126475,"end_time":"2021-06-15T18:02:51.805314","exception":false,"start_time":"2021-06-15T18:02:51.678839","status":"completed"},"tags":[],"id":"entire-distinction"}},{"cell_type":"code","source":"class TransformerDecoder(tf.keras.layers.Layer):\n ''' Decoder Layer Component for the Transformer '''\n def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size, \n maximum_position_encoding, dropout_rate=0.1):\n '''\n Args:\n num_layers (int): Number of decoder layers to be used\n d_model (int): Depth of the d-dimensional space used for positional encoding of image embedding\n num_heads (int): Number of heads to use in the multi-head-attention block\n dff (int): Number of units to use in the feed-forward neural network\n maximum_positional_encoding (int): Maximum limit for positional encoding\n dropout_rate (float): Percentage of nodes to drop in a given layer\n \n Returns:\n None; This is an intiailization\n '''\n super(TransformerDecoder, self).__init__()\n\n self.d_model = d_model\n self.num_layers = num_layers\n self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)\n self.pos_encoding = positional_encoding_1d(maximum_position_encoding, d_model)\n self.dec_layers = [TransformerDecoderLayer(d_model, num_heads, dff, dropout_rate) for _ in range(num_layers)]\n self.dropout = tf.keras.layers.Dropout(dropout_rate)\n \n def call(self, x, enc_output, training, look_ahead_mask=None, padding_mask=None):\n ''' Calling function for the decoder '''\n seq_len = tf.shape(x)[1]\n attention_weights = {}\n \n # adding embedding and position encoding.\n # – (batch_size, target_seq_len, d_model)\n x = self.embedding(x) \n x *= tf.math.sqrt(tf.cast(self.d_model, TARGET_DTYPE))\n x += self.pos_encoding[:, :seq_len, :]\n \n x = self.dropout(x, training=training)\n\n for i in range(self.num_layers):\n x, block1, block2 = self.dec_layers[i](x, enc_output, training, look_ahead_mask, padding_mask)\n attention_weights['decoder_layer{}_block1'.format(i+1)] = block1\n attention_weights['decoder_layer{}_block2'.format(i+1)] = block2\n \n # x.shape == (batch_size, target_seq_len, d_model)\n return x, attention_weights\n\n\nsample_decoder = TransformerDecoder(num_layers=2, d_model=D_MODEL, num_heads=8, \n dff=D_FF, target_vocab_size=VOCAB_LEN,\n maximum_position_encoding=MAX_LEN)\ntemp_input = tf.random.uniform((1, MAX_LEN), dtype=tf.int64, minval=0, maxval=VOCAB_LEN)\noutput, attn = sample_decoder(temp_input, \n enc_output=demo_encoder_output, \n training=False,\n look_ahead_mask=None, \n padding_mask=None)\ndel sample_decoder\noutput.shape, attn['decoder_layer2_block2'].shape","metadata":{"papermill":{"duration":1.081253,"end_time":"2021-06-15T18:02:53.012722","exception":false,"start_time":"2021-06-15T18:02:51.931469","status":"completed"},"tags":[],"id":"functional-health","outputId":"e56cce23-0e8d-4ae7-f4b4-2622140ecddd","execution":{"iopub.status.busy":"2022-01-25T14:55:12.372216Z","iopub.execute_input":"2022-01-25T14:55:12.372441Z","iopub.status.idle":"2022-01-25T14:55:13.346057Z","shell.execute_reply.started":"2022-01-25T14:55:12.372414Z","shell.execute_reply":"2022-01-25T14:55:13.345253Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 5.2.11 TRANSFORMER - PUT IT ALL TOGETHER\n\n---\n\n\nOur Transformer consists of the **transformer encoder**, **transformer decoder** and a **final linear layer**. \n* The input to the encoder is the output of our image encoder (i.e. output of EfficientNetV2)\n* The output of the decoder is the input to the linear layer and its output is returned.","metadata":{"papermill":{"duration":0.125082,"end_time":"2021-06-15T18:02:53.263494","exception":false,"start_time":"2021-06-15T18:02:53.138412","status":"completed"},"tags":[],"id":"inclusive-typing"}},{"cell_type":"code","source":"class Transformer(tf.keras.Model):\n \"\"\" Final Transformer Architecture Block \"\"\"\n def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size, \n pe_input, pe_target, dropout_rate=0.1):\n '''\n Args:\n num_layers (int): Number of decoder layers to be used\n d_model (int): Depth of the d-dimensional space used for positional encoding of image embedding\n num_heads (int): Number of heads to use in the multi-head-attention block\n dff (int): Number of units to use in the feed-forward neural network\n target_vocab_size (int): Vocabulary size of the target output\n pe_input (int): Input positional encodings\n pe_target (int): Target positional encoding\n dropout_rate (float): Percentage of nodes to drop in a given layer\n \n Returns:\n None; This is an intiailization\n '''\n \n super(Transformer, self).__init__()\n \n self.t_encoder = TransformerEncoder(num_layers, d_model, num_heads, \n dff, pe_input, dropout_rate)\n self.t_decoder = TransformerDecoder(num_layers, d_model, num_heads, dff, \n target_vocab_size, pe_target, dropout_rate)\n self.t_final_layer = tf.keras.layers.Dense(target_vocab_size)\n \n def call(self, t_inp, t_tar, training, enc_padding_mask=None, \n look_ahead_mask=None, dec_padding_mask=None):\n '''\n Args:\n t_inp: Input sequences \n t_tar: Target sequences \n training (bool): Whether training is to be done\n enc_padding_mask (array or None): Padding masks from the encoder\n look_ahead_mask: None if no look ahead masks to be used\n dec_padding_mask (array or None): Padding masks from the decoder\n\n Returns:\n final_output and attention_weights of the Transformer model\n '''\n # (batch_size, inp_seq_len, d_model)\n enc_output = self.t_encoder(t_inp, training, enc_padding_mask) \n \n # dec_output.shape == (batch_size, tar_seq_len, d_model)\n dec_output, attention_weights = self.t_decoder(t_tar, enc_output, \n training, look_ahead_mask, \n dec_padding_mask)\n\n # (batch_size, tar_seq_len, target_vocab_size)\n final_output = self.t_final_layer(dec_output) \n \n return final_output, attention_weights\n\n\n# sample_transformer = Transformer(num_layers=2, d_model=D_MODEL, \n# num_heads=8, dff=1024, \n# target_vocab_size=VOCAB_LEN, \n# pe_input=IMG_SEQ_LEN, pe_target=MAX_LEN)\n# fn_out, _ = sample_transformer(demo_encoder_output, SAMPLE_LBLS, training=False, \n# enc_padding_mask=None, \n# look_ahead_mask=None,\n# dec_padding_mask=None)\n\n# # (batch_size, tar_seq_len, target_vocab_size)\n# print(fn_out.shape)\n\n# del sample_transformer","metadata":{"papermill":{"duration":0.138285,"end_time":"2021-06-15T18:02:53.527724","exception":false,"start_time":"2021-06-15T18:02:53.389439","status":"completed"},"tags":[],"id":"minute-vertex","execution":{"iopub.status.busy":"2022-01-25T14:55:13.347312Z","iopub.execute_input":"2022-01-25T14:55:13.34753Z","iopub.status.idle":"2022-01-25T14:55:13.357822Z","shell.execute_reply.started":"2022-01-25T14:55:13.347504Z","shell.execute_reply":"2022-01-25T14:55:13.357191Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 5.3 CREATE A LEARNING RATE SCHEDULER\n---\n\nWe utiliize the learning rate scheduler from the \"Attention Is All You Need\" paper with some minor tweaks.","metadata":{"papermill":{"duration":0.126289,"end_time":"2021-06-15T18:02:53.778871","exception":false,"start_time":"2021-06-15T18:02:53.652582","status":"completed"},"tags":[],"id":"urban-comparison"}},{"cell_type":"code","source":"print(\"\\n... LEARNING RATE SCHEDULE CREATION STARTING ...\\n\")\n\n# Part of the Training Configuration\nEPOCHS = 30\nTOTAL_STEPS = TRAIN_STEPS*EPOCHS\n\n# Learning Rate Scheduler Configuration\nWARM_STEPS = (TRAIN_STEPS-1)*4 # Suuuuuper long ramp-up\n\nclass CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n ''' Creates a custom Learning Rate Scheduler '''\n def __init__(self, d_model, warmup_steps=4000):\n '''\n Args:\n d_model (int): Depth of the d-dimensional space used for positional encoding of image embedding\n warmup_steps (int): Number of steps for which the learning rate will ramp up to the desired peak learning rate value\n '''\n super(CustomSchedule, self).__init__()\n\n self.d_model = tf.cast(d_model, tf.float32)\n self.warmup_steps = warmup_steps\n\n def __call__(self, step):\n ''' Returns learning rate at different steps '''\n arg1 = tf.math.rsqrt(step)\n arg2 = step * (self.warmup_steps ** -1.5)\n\n return tf.math.rsqrt(self.d_model*1.75) * tf.math.minimum(arg1, arg2)\n \ntemp_learning_rate_schedule = CustomSchedule(D_MODEL, WARM_STEPS)\nplt.plot(temp_learning_rate_schedule(tf.range(TRAIN_STEPS*EPOCHS, dtype=tf.float32)))\nplt.ylabel(\"Learning Rate\")\nplt.xlabel(\"Train Step\")\nplt.show()\n \n# def lr_schedule_fn(step, total_steps, warm_lr_start, warm_steps, peak_lr_start, lr_final, n_epochs):\n# \"\"\" Function to generate the learning rate for a given step based on parameters\n \n# Args:\n# step (int): The current step for which to calculate the respective learning rate\n# total_steps (int): The total number of steps for the entire training regime\n# warm_lr_start (float): The starting learning rate prior to warmup\n# warm_steps (int): The number of steps for which the learning rate will ramp up\n# to the desired peak learning rate value (more steps will result in less\n# dramatic changes to existing weights... better for pretrained models)\n# peark_lr_start (float): The starting learning rate after warmup (peak value)\n# lr_final (float): The final learning rate to step down to by the end of training\n# n_epochs (int): The total number of epochs for the training regime\n \n# Returns:\n# The learning rate (float) to be used for a given step\n# \"\"\"\n \n# # exponential warmup\n# if step < warm_steps:\n# warmup_factor = (step / warm_steps) ** 2\n# lr_rate = warm_lr_start + (peak_lr_start - warm_lr_start) * warmup_factor \n \n# # staircase decay\n# else:\n# power = (step - warm_steps) // ((total_steps - warm_steps) / (n_epochs + 1))\n# decay_factor = ((peak_lr_start / lr_final) ** (1 / n_epochs)) ** power\n# lr_rate = peak_lr_start / decay_factor\n \n# return round(lr_rate, 8)\n\n\n# def plot_lr_schedule(lr_schedule, name=\"\"):\n# \"\"\" Plot the learning rate schedule over the course of training\n \n# Args:\n# lr_schedule (list of floats): The values to use for the LR over the\n# course of training\n# name (str, optional): A name for the LR schedule\n \n# Returns:\n# None; A plot of the how the learning rate changes over time will be displayed\n \n# \"\"\"\n# schedule_info = f'start: {lr_schedule[0]:.6f}, max: {max(lr_schedule):.6f}, final: {lr_schedule[-1]:.6f}'\n# plt.figure(figsize=(18,6))\n# plt.plot(lr_schedule)\n# plt.title(f\"Step Learning Rate Schedule {name+', ' if name else name}{schedule_info}\", size=16, fontweight=\"bold\")\n# plt.grid()\n# plt.show()\n \n# class LRS():\n# \"\"\" LEARNING RATE SCHEDULER OBJECT\"\"\"\n# def __init__(self, optimizer, lr_schedule):\n# self.opt = optimizer\n# self.lr_schedule = lr_schedule\n \n# # assign initial learning rate\n# self.lr = lr_schedule[0]\n# self.opt.learning_rate.assign(self.lr)\n \n# def step(self, step):\n# self.lr = self.lr_schedule[step]\n# # assign learning rate to optimizer\n# self.opt.learning_rate.assign(self.lr)\n \n# def get_counter(self):\n# return self.c\n \n# def get_lr(self):\n# return self.lr\n\n# # Create the Schedule and Plot\n# lr_schedule = [\n# lr_schedule_fn(step, TOTAL_STEPS, WARM_START_LR, WARM_STEPS, PEAK_START_LR, FINAL_LR, EPOCHS) \\\n# for step in range(TOTAL_STEPS)\n# ]\n# plot_lr_schedule(lr_schedule)\n\nprint(\"\\n... LEARNING RATE SCHEDULE CREATION FINISHED ...\\n\")","metadata":{"papermill":{"duration":0.308955,"end_time":"2021-06-15T18:02:54.21401","exception":false,"start_time":"2021-06-15T18:02:53.905055","status":"completed"},"tags":[],"id":"municipal-target","outputId":"51445c5f-72c6-4297-ddaf-d434cd581679","execution":{"iopub.status.busy":"2022-01-25T14:55:13.359361Z","iopub.execute_input":"2022-01-25T14:55:13.35962Z","iopub.status.idle":"2022-01-25T14:55:13.606745Z","shell.execute_reply.started":"2022-01-25T14:55:13.359593Z","shell.execute_reply":"2022-01-25T14:55:13.605849Z"},"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":{"papermill":{"duration":0.181274,"end_time":"2021-06-15T18:02:54.524739","exception":false,"start_time":"2021-06-15T18:02:54.343465","status":"completed"},"tags":[],"id":"available-moisture"}},{"cell_type":"code","source":"# Hyperparameters For ViT\nVIT_PATCH_SIZE=16\nVIT_PROJECTION_DIM=128\nVIT_N_TRANSFORMER_LAYERS=8\nVIT_N_HEADS=4\nVIT_DROPOUT=0.1\n\n# Hyperparameters For Transformer\nN_LAYERS = 4\nD_MODEL = IMG_EMB_DEPTH\nD_FF = 256\nN_HEADS = 4\nDROPOUT_RATE = 0.1\nPE_INPUT = IMG_SEQ_LEN\nPE_OUTPUT = MAX_LEN\nTARGET_V_SIZE = VOCAB_LEN\n\nclass Config():\n ''' Class to initialize the Encoder, Decoder and Learning Rate configurations '''\n def __init__(self,):\n self.encoder_config = {}\n self.transformer_config = {}\n self.lr_config = {}\n \n def initialize_encoder_config(self, patch_size, projection_dim, \n n_transformer_layers, n_heads, dropout, img_shape):\n self.encoder_config = dict(\n patch_size=patch_size, \n projection_dim=projection_dim, \n n_transformer_layers=n_transformer_layers, \n n_heads=n_heads, \n dropout=dropout, \n img_shape=img_shape)\n \n def initialize_transformer_config(self, vocab_len, n_transformer_layers, \n transformer_d_dff, transformer_n_heads, \n encoder_out_seq_len, encoder_out_depth, \n dropout_rate=0.1):\n self.transformer_config = dict(num_layers=n_transformer_layers, \n d_model=encoder_out_depth, \n num_heads=transformer_n_heads, \n dff=transformer_d_dff,\n target_vocab_size=vocab_len, \n pe_input=encoder_out_seq_len, \n pe_target=MAX_LEN, dropout_rate=0.1)\n \n def initialize_lr_config(self, warm_steps, n_epochs):\n self.lr_config = dict(\n warm_steps=warm_steps, \n n_epochs=n_epochs,\n )\n \ntraining_config = Config()\ntraining_config.initialize_transformer_config(vocab_len=VOCAB_LEN, \n n_transformer_layers=N_LAYERS,\n transformer_d_dff=D_FF,\n transformer_n_heads=N_HEADS,\n encoder_out_seq_len=IMG_SEQ_LEN, \n encoder_out_depth=IMG_EMB_DEPTH)\n\ntraining_config.initialize_encoder_config(patch_size=VIT_PATCH_SIZE, \n projection_dim=VIT_PROJECTION_DIM, \n n_transformer_layers=VIT_N_TRANSFORMER_LAYERS, \n n_heads=VIT_N_HEADS, \n dropout=VIT_DROPOUT, \n img_shape=IMG_SHAPE)\n\ntraining_config.initialize_lr_config(warm_steps=WARM_STEPS, n_epochs=EPOCHS,)\n\nprint(f\"\\nTRAINING ENCODER CONFIG:\\n\\t--> {training_config.encoder_config}\\n\")\nprint(f\"\\nTRAINING TRANSFORMER CONFIG:\\n\\t--> {training_config.transformer_config}\\n\")\nprint(f\"TRAINING LEARNING RATE CONFIG:\\n\\t--> {training_config.lr_config}\\n\")","metadata":{"papermill":{"duration":0.144498,"end_time":"2021-06-15T18:02:54.797447","exception":false,"start_time":"2021-06-15T18:02:54.652949","status":"completed"},"tags":[],"id":"expected-debut","outputId":"6a66c6a5-7f24-46b7-dfb1-52b15085fa17","execution":{"iopub.status.busy":"2022-01-25T14:55:13.608134Z","iopub.execute_input":"2022-01-25T14:55:13.608392Z","iopub.status.idle":"2022-01-25T14:55:13.624169Z","shell.execute_reply.started":"2022-01-25T14:55:13.608363Z","shell.execute_reply":"2022-01-25T14:55:13.623441Z"},"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","metadata":{"papermill":{"duration":0.127266,"end_time":"2021-06-15T18:02:55.052494","exception":false,"start_time":"2021-06-15T18:02:54.925228","status":"completed"},"tags":[],"id":"automated-technology"}},{"cell_type":"code","source":"print(\"\\n... TRAINING PREPERATION STARTING ...\\n\")\n\ndef prepare_for_training(lr_config, encoder_config, transformer_config, encoder_wts=None, transformer_wts=None, verbose=0):\n \"\"\" \n Declare required objects under TPU session scope and return ready for training\n \n Args:\n lr_config (dict): Keyword arguments mapped to desired values for lr schedule function\n encoder_config (dict): Keyword arguments mapped to desired values for encoder model instantiation \n transformer_config (dict): Keyword arguments mapped to desired values for transformer model instantiation \n encoder_wts (str, optional): Path to pretrained model weights for encoder\n transformer_wts (str, optional): Path to pretrained model weights for encoder\n verbose (bool, optional): Whether or not to print model information and plot lr schedule\n \n Returns:\n loss_fn - Loss function used during training\n metrics - Loss and Accuracy metrics for training and validation data\n optimizer - Optimizer used for training\n lr_scheduler - Learning rate scheduler \n encoder - Encoder model\n decoder - Decoder model\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 \n print(\"\\t--> CREATING LOSS FUNCTION ...\")\n # Declare the loss object\n # - Sparse categorical cross entropy loss is used as root loss\n loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n from_logits=True, reduction=tf.keras.losses.Reduction.NONE\n )\n \n def loss_fn(real, pred):\n # Convert to uint8\n mask = tf.math.not_equal(real, 0)\n loss_ = loss_object(real, pred)\n loss_ *= tf.cast(mask, dtype=loss_.dtype)\n\n # https://www.tensorflow.org/tutorials/distribute/custom_training#define_the_loss_function\n loss_ = tf.nn.compute_average_loss(loss_, global_batch_size=REPLICA_BATCH_SIZE)\n return loss_\n \n \n # def loss_fn(real, pred):\n # per_example_loss = loss_object(real, pred)\n # return tf.nn.compute_average_loss(per_example_loss, global_batch_size=OVERALL_BATCH_SIZE)\n \n # Declare the metrics\n # - Loss (train only) and sparse categorical accuracy will be used\n print(\"\\t--> CREATING METRICS ...\")\n metrics = {\n 'batch_loss':tf.keras.metrics.Mean(),\n 'train_loss': tf.keras.metrics.Mean(),\n 'train_acc': tf.keras.metrics.SparseCategoricalAccuracy(),\n 'val_loss': tf.keras.metrics.Mean(),\n 'val_acc': tf.keras.metrics.SparseCategoricalAccuracy(),\n 'val_lsd': tf.keras.metrics.Mean(), \n }\n \n \n print(\"\\t--> CREATING LEARNING RATE SCHEDULER ...\")\n # Declare the learning rate schedule (try this as actual lr schedule and list...)\n lr_scheduler = CustomSchedule(transformer_config[\"d_model\"], lr_config[\"warm_steps\"])\n \n print(\"\\t--> CREATING OPTIMIZER ...\")\n # Instiate an optimizer\n optimizer = tf.keras.optimizers.Adam(lr_scheduler)\n \n # Instantiate the encoder model \n print(\"\\t--> CREATING ENCODER MODEL ARCHITECTURE ...\")\n encoder = ViTEncoder(**encoder_config)\n initialization_batch = encoder(\n tf.ones(((REPLICA_BATCH_SIZE,)+encoder_config[\"img_shape\"]), dtype=TARGET_DTYPE), \n training=False,\n )\n \n # Instantiate the decoder model\n print(\"\\t--> CREATING TRANSFORMER MODEL ARCHITECTURE...\")\n transformer = Transformer(**transformer_config)\n transformer(initialization_batch, tf.random.uniform((REPLICA_BATCH_SIZE, 1)), training=False)\n \n if encoder_wts is not None:\n print(\"\\t--> LOADING ENCODER MODEL WEIGHTS ...\")\n encoder.load_weights(encoder_wts)\n\n \n if transformer_wts is not None:\n print(\"\\t--> LOADING TRANSFORMER MODEL WEIGHTS (WILL OVERWRITE ENCODER)...\")\n transformer.load_weights(transformer_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... TRANSFORMER MODEL SUMMARY...\\n\")\n print(transformer.summary())\n\n # print(\"\\n\\n... LR SCHEDULE PLOT...\\n\")\n # plot_lr_schedule(lr_schedule)\n \n return loss_fn, metrics, optimizer, lr_scheduler, encoder, transformer\n \n \nprint(\"\\n... GENERATING THE FOLLOWING:\")\n# Instantiate our required training components in the correct scope\nloss_fn, metrics, optimizer, lr_scheduler, encoder, transformer = \\\n prepare_for_training(lr_config=training_config.lr_config,\n encoder_config=training_config.encoder_config,\n transformer_config=training_config.transformer_config,\n encoder_wts=(ENCODER_CKPT_PATH if ENCODER_CKPT_PATH!=\"\" else None),\n transformer_wts=(TRANSFORMER_CKPT_PATH if TRANSFORMER_CKPT_PATH!=\"\" else None),\n verbose=1,)\n\nprint(\"\\n... TRAINING PREPERATION FINISHED ...\\n\")","metadata":{"papermill":{"duration":24.444162,"end_time":"2021-06-15T18:03:19.624233","exception":false,"start_time":"2021-06-15T18:02:55.180071","status":"completed"},"tags":[],"id":"dense-turtle","outputId":"b4231e09-85e4-4699-d135-57cdd63e3e87","execution":{"iopub.status.busy":"2022-01-25T14:55:13.625522Z","iopub.execute_input":"2022-01-25T14:55:13.625989Z","iopub.status.idle":"2022-01-25T14:55:38.246509Z","shell.execute_reply.started":"2022-01-25T14:55:13.625952Z","shell.execute_reply":"2022-01-25T14:55:38.245485Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 5.6 LOSS CLASSES AND REDUCTION\n\nIn order to accurately calculate loss when leveraging a TPU, we have to accumulate the losses that will be calculated across the individual replicas. Knowing this we are limited to using a **`reduction`** value of **`SUM`** or **`NONE`** as the default value and some of the other options will not work with TPU.\n\n---\n\nDuring training, when a batch is [**distributed to the replicas**](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#run), each replica receives a part of the batch and\ncalculates the loss values separately. We **SHOULD NOT** calculate the average of the per-example losses on the (partial) batch the replica recieves. \n\n**The intuition behind this is as follows:**\n\n---\n* The gradients calculated on each replica will be synced across the replicas\n * Therefore, they are summed before the optimizer applies the gradients to update the model's parameters\n* If we use the averaged per examples loss to compute the graident on each replica, the final graident applied by the optimizer will correspond to the sum of these averaged per-examples losses for respective replicas.\n * This is incorrect. The optimizer should apply the gradient obtained from the averaged per-examples loss **over the whole distributed batch**\n * It's worth noting that each replica may infact receive different number of examples. \n * Therefore it is impossible, in general, to obtain the averaged per example loss over the whole distributed batch from by simply dividing it by the number of replicas.\n\n---\n\n**Therefore, we can see that for each replica, we calculate the sum of per examples losses divided by the batch size of the whole distributed batch, which will give the optimizer the correct gradients to apply.**\n\n**EDIT**\n* **In this notebook, we have the option to use [*gradient accumulation*](https://arxiv.org/pdf/1710.02368)**\n* In ***gradient accumulation***, each replica receives several batches before the optimizer applies the graidents\n * we divide the sum of per examples losses by the update size (i.e. the number of examples used for one parameter update) rather than by the size of a single distributed batch.\n\n**In the following cell we will demonstrate, using dummy values and pretending we are distributing them, how to deal with the accumulation of the loss values across replicas.**","metadata":{"papermill":{"duration":0.142029,"end_time":"2021-06-15T18:03:19.908079","exception":false,"start_time":"2021-06-15T18:03:19.76605","status":"completed"},"tags":[],"id":"incident-separate"}},{"cell_type":"markdown","source":"## 5.7 DISTRIBUTE THE DATASETS ACROSS REPLICAS\n\nWith an input pipeline written using the [**tf.data.Dataset**](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) API, we can use [**strategy.experimental_distribute_dataset**](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#experimental_distribute_dataset) to turn it into a ***distributed dataset***, which produces **`per-replica`** values (which are objects of type [**PerReplica**](https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/distribute/values.py#L361)) when iterating over it. \n\nFor example, \n\n```python\n ds = (... something that is a `tf.data.Dataset` ...)\n dist_ds = strategy.experimental_distribute_dataset(ds)\n```\n\n**`dist_ds`** will now be distributed across all replicas.\n\n---\n\nThe distributed datasets (when working with TPU) contain objects of type [**tensorflow.python.distribute.values.PerReplica**](https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/distribute/values.py#L361), which is a subclass of [**tf.distribute.DistributedValues**](https://www.tensorflow.org/api_docs/python/tf/distribute/DistributedValues) that is the base class for representing distributed values.\n\nWhen iterating over the dataset we will still get a tuple containing two values. However, the tuple now contains **`PerReplica`** objects wheras before that tuple contained tensors representing the image and the label/id respectively.","metadata":{"papermill":{"duration":0.136056,"end_time":"2021-06-15T18:03:20.204669","exception":false,"start_time":"2021-06-15T18:03:20.068613","status":"completed"},"tags":[],"id":"tropical-singer"}},{"cell_type":"markdown","source":"## 5.8 DISTRIBUTED COMPUTATION & OPTIMIZING LOOPS\n\nFor each distributed batch (which contains **`PerReplica`** objects as discussed previously) produced by a distributed dataset, we use [**`strategy.run`**](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#run) to perform a distributed computation on different TPU replicas, each processes a part of the batch.\n\n\n---\n\nTo understand how **`strategy.run`** will execute across the replicas, we can look at an example:\n\n```python\n @tf.function\n def dist_step(dist_batch):\n strategy.run(replica_fn, args=dist_batch)\n \n for dist_batch in dist_ds:\n dist_step(dist_batch)\n```\n\nHere **`replica_fn`** is a function that is going to be run on each replica, and it should work with **tensors**, not with **`PerReplica`** objects.\n* You define the operations (for example, forward pass, compute loss values and gradients, etc.) to peform just like witout using TPU. \n\n---\n\nWhen working with **`TPU`**, either [**`strategy.run`**](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#run) has to be called inside [**`tf.function`**](https://www.tensorflow.org/api_docs/python/tf/function) or the replica function has to be annotated with [**`tf.function`**](https://www.tensorflow.org/api_docs/python/tf/function). \n\nFor example:\n\n```python\n @tf.function\n def replica_fn(batch):\n \n model(batch)\n ...\n \n for dist_batch in dist_ds:\n strategy.run(replica_fn, args=dist_batch)\n```\n\nThe above code snippet is a high level concept, and **`replica_fn`** doesn't necessary receive a single argument. \n* In our case, the original dataset yields tuples of tensors\n* A distributed batch is also a tuple of **`PerReplica`** objects and the **`replica_fn`** is actually receiving the unpacked version of a tuple of tensors as arguments.\n\n---\n\nIf a dataset yield a single tensor, you can do things like \n\n```python\n @tf.function\n def replica_fn(batch):\n \n tensor0 (, ... tensorN) = batch\n model(tensor0, ... tensorN)\n\n strategy.run(replica_fn, args=(dist_batch,))\n```\n\nwhere **`replica_fn`** expects a single tensor as arugment. Even if a dataset yields tuples of tensors, the above code still works, but **`replica_fn`** expects a single tuple of tensors as argument.\n\n---\n\nWe also have to discuss how to collect the returned values from [**`strategy.run`**](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#run).\n\nThe results of [**`strategy.run`**](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#run) are also \ndistributed values, just like the distributed batches it takes as inputs. \n* For each return value, we can use [strategy.experimental_local_results](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#experimental_local_results) to obtain a tuple of tensors from all replicas, and we can use [**`tf.concat`**](https://www.tensorflow.org/api_docs/python/tf/concat) to aggregate them into a single tensor.\n* We will use this method to collect the labels and model predictions\n\n---\n\nWe will need to iterate over the dataset to perform inference/train on the whole (distributed) dataset. When leveraging a TPU this is a non-trivial task. An example of iterating over a distributed dataset is:\n\n```python\n for dist_batch in dist_ds:\n dist_step(dist_batch)\n```\n\nEvery step in the loop, which calls [**`strategy.run`**](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#run), will have a communication between the local VM (in our case, the Kaggle VM) and the remote TPU worker(s). \n\n**This is obviously not ideal.**\n\nHowever, you can iterate the distributed dataset inside a `tf.function` as shown by:\n\n``` python\n @tf.function\n def dist_run_on_dataset(dist_ds):\n \n for dist_batch in dist_ds:\n dist_step(dist_batch)\n \n dist_process_dataset(dist_ds)\n```\n\nThis way, all the operations conducted on the dataset are compiled into a graph which is sent to the remote TPU worker(s) for execution. This will vastly reduce the running time and limit the time TPUs will sit idle waiting for data from the local VM. See [**TPU: extreme optimizations**](https://www.kaggle.com/c/flower-classification-with-tpus/discussion/135443) for a good benchmark by [**Martin Görner**](https://www.kaggle.com/mgornergoogle).\n\nIn this notebook, we use a fixed number of training steps, so we can also use\n\n```python \n @tf.function\n def dist_process_dataset(dist_ds_iter):\n \n for _ in tf.range(n_stes):\n dist_step(next(dist_ds_iter))\n \n dist_ds_iter = iter(dist_ds)\n dist_process_dataset(dist_ds_iter)\n```\n\n---\n\n**With the above discussions, we are ready to define the routines used for training, validation and prediction. Let's get started!**\n\n
\n 📖 REFERENCE:
\n - Tutorial - Using Iterators
\n - Tutorial - Iterating Inside a
\n - Kaggle Discussion - TPU: Extreme Optimizations
\n - Kaggle Notebook - Custom Training Loop With 100+ Flowers on TPU
\n
\n","metadata":{"papermill":{"duration":0.13329,"end_time":"2021-06-15T18:03:20.470101","exception":false,"start_time":"2021-06-15T18:03:20.336811","status":"completed"},"tags":[],"id":"driven-marshall"}},{"cell_type":"markdown","source":"# 6. Model_training\n---\n\n\nIn this section we will define the training and validation routines as well as the final custom training loop that will execute everything we have worked on up until this point.","metadata":{"papermill":{"duration":0.13237,"end_time":"2021-06-15T18:03:20.734916","exception":false,"start_time":"2021-06-15T18:03:20.602546","status":"completed"},"tags":[],"id":"pleased-kansas"}},{"cell_type":"markdown","source":"## 6.1 INDIVIDUAL TRAIN STEP\n\n---\n\nINFORMATION","metadata":{"papermill":{"duration":0.130794,"end_time":"2021-06-15T18:03:20.998395","exception":false,"start_time":"2021-06-15T18:03:20.867601","status":"completed"},"tags":[],"id":"designed-remark"}},{"cell_type":"code","source":"def train_step(_image_batch, _inchi_batch):\n \"\"\" \n Forward pass of the training which calculates the gradients.\n \n Args:\n _image_batch: Distributed batches of image dataset\n _inchi_batch: Distributed batches of inchi dataset\n \n Returns:\n None; Calculates the gradients and loss metric\n \"\"\"\n _inchi_batch_input = _inchi_batch[:, :-1]\n _inchi_batch_target = _inchi_batch[:, 1:]\n combined_mask = create_mask(_inchi_batch_input, _inchi_batch_target)\n\n with tf.GradientTape() as tape:\n _image_embedding = encoder(_image_batch, training=True)\n prediction_batch, _ = transformer(_image_embedding, _inchi_batch_input, training=True, look_ahead_mask=combined_mask)\n \n # Update Loss Accumulator\n batch_loss = loss_fn(_inchi_batch_target, prediction_batch)/(MAX_LEN-1)\n\n # Update Accuracy Metric\n metrics[\"train_acc\"].update_state(_inchi_batch_target, prediction_batch, \n sample_weight=tf.where(tf.not_equal(_inchi_batch_target, PAD_TOKEN), 1.0, 0.0))\n\n\n # backpropagation using variables, gradients and loss\n # - split this into two seperate optimizers/lrs/etc in the future\n # - we use the batch loss accumulation to update gradients\n gradients = tape.gradient(batch_loss, encoder.trainable_variables + transformer.trainable_variables)\n gradients, _ = tf.clip_by_global_norm(gradients, 10.0)\n optimizer.apply_gradients(zip(gradients, encoder.trainable_variables+transformer.trainable_variables))\n \n metrics[\"batch_loss\"].update_state(batch_loss)\n metrics[\"train_loss\"].update_state(batch_loss)\n\n@tf.function\ndef dist_train_step(_image_batch, _inchi_batch):\n strategy.run(train_step, args=(_image_batch, _inchi_batch))","metadata":{"papermill":{"duration":0.146617,"end_time":"2021-06-15T18:03:21.278894","exception":false,"start_time":"2021-06-15T18:03:21.132277","status":"completed"},"tags":[],"id":"magnetic-domain","execution":{"iopub.status.busy":"2022-01-25T14:55:38.248926Z","iopub.execute_input":"2022-01-25T14:55:38.249996Z","iopub.status.idle":"2022-01-25T14:55:38.260292Z","shell.execute_reply.started":"2022-01-25T14:55:38.249953Z","shell.execute_reply":"2022-01-25T14:55:38.259336Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 6.2 INDIVIDUAL VAL STEP\n\n---\n\nINFORMATION","metadata":{"papermill":{"duration":0.130542,"end_time":"2021-06-15T18:03:21.537309","exception":false,"start_time":"2021-06-15T18:03:21.406767","status":"completed"},"tags":[],"id":"genuine-october"}},{"cell_type":"code","source":"def val_step(_image_batch, _inchi_batch):\n \"\"\" \n Forward pass of the validation step\n \n Args:\n image_batch: Distributed batches of image dataset\n inchi_batch: Distributed batches of inchi dataset\n \n Returns:\n predictions_seq_batch: Predictions on the validation batch\n \"\"\"\n \n # Initialize batch_loss\n batch_loss = tf.constant(0.0, TARGET_DTYPE) \n transformer_pred_batch = tf.ones((REPLICA_BATCH_SIZE, 1), dtype=tf.uint8)\n \n # Get image embedding (once)\n _image_embedding = encoder(_image_batch, training=False)\n \n # Teacher forcing - feeding the target as the next input\n for c_idx in range(1, MAX_LEN):\n gt_batch_id = _inchi_batch[:, c_idx]\n combined_mask = create_mask(_inchi_batch, transformer_pred_batch)\n \n # predictions.shape == (batch_size, seq_len, vocab_size)\n prediction_batch, attention_weights = transformer(_image_embedding, transformer_pred_batch, training=False, look_ahead_mask=combined_mask)\n predicted_batch_id = prediction_batch[:, -1:, :]\n \n # Update Loss Accumulator\n batch_loss += loss_fn(gt_batch_id, predicted_batch_id[:, -1])\n \n # Update Accuracy Metric\n metrics[\"val_acc\"].update_state(gt_batch_id, predicted_batch_id[:, -1],\n sample_weight=tf.where(tf.not_equal(gt_batch_id, PAD_TOKEN), 1.0, 0.0))\n\n # no teacher forcing, predicted char is next transformer input\n transformer_pred_batch = tf.concat([transformer_pred_batch, tf.cast(tf.argmax(predicted_batch_id, axis=-1), tf.uint8)], axis=-1)\n \n # Update Loss Metric\n metrics[\"val_loss\"].update_state(batch_loss)\n return transformer_pred_batch \n\n \n@tf.function\ndef dist_val_step(_val_dist_ds):\n _val_image_batch, _val_inchi_batch = next(_val_dist_ds)\n predictions_seq_batch_per_replica = strategy.run(val_step, args=(_val_image_batch, _val_inchi_batch))\n predictions_seq_batch_accum = strategy.gather(predictions_seq_batch_per_replica, axis=0)\n _val_inchi_batch_accum = strategy.gather(_val_inchi_batch, axis=0)\n return predictions_seq_batch_accum, _val_inchi_batch_accum","metadata":{"papermill":{"duration":0.14455,"end_time":"2021-06-15T18:03:21.810182","exception":false,"start_time":"2021-06-15T18:03:21.665632","status":"completed"},"tags":[],"id":"saving-insert","execution":{"iopub.status.busy":"2022-01-25T14:55:38.266613Z","iopub.execute_input":"2022-01-25T14:55:38.266953Z","iopub.status.idle":"2022-01-25T14:55:38.279541Z","shell.execute_reply.started":"2022-01-25T14:55:38.266921Z","shell.execute_reply":"2022-01-25T14:55:38.278562Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 6.3 INITIALIZE LOGGER\n\n---\n\nINFORMATION","metadata":{"papermill":{"duration":0.129364,"end_time":"2021-06-15T18:03:22.068649","exception":false,"start_time":"2021-06-15T18:03:21.939285","status":"completed"},"tags":[],"id":"temporal-warehouse"}},{"cell_type":"code","source":"class StatLogger():\n ''' This class initializes the Logger '''\n def __init__(self, verbose_frequency=100, print_style=\"tight\"):\n self.train_loss = []\n self.train_acc = []\n self.val_loss = []\n self.val_acc = []\n self.val_lsd = []\n self.step = []\n self.epoch = []\n self.lr = []\n \n self.current_step = 0\n self.epoch_start_time = 0\n self.batch_start_time = 0\n self.verbose_frequency = verbose_frequency\n self.print_style = print_style\n \n def print_last_val(self, current_time):\n if self.print_style==\"tight\":\n print(f\"| VAL DATA | STEP {VAL_STEPS:>4}/{VAL_STEPS} | \" \\\n f\"ACC: {str(self.val_acc[-1]*100)[:5]:<5} – \" \\\n f\"LOSS: {str(self.val_loss[-1])[:5]:<5} – \" \\\n f\"LSD: {str(self.val_lsd[-1]):<3} |\")\n else:\n print(f'\\n\\n{\"-\"*100}\\n{\"-\"*100}\\n' \\\n f'{\"-\"*25:<25}{\"VALIDATION ACCURACY : \"+str(self.val_acc[-1]*100): ^50}{\"-\"*25:>25}\\n' \\\n f'{\"-\"*100}\\n' \\\n f'{\"-\"*25:<25}{\"VALIDATION LOSS : \"+str(self.val_loss[-1]): ^50}{\"-\"*25:>25}\\n' \\\n f'{\"-\"*100}\\n' \\\n f'{\"-\"*25:<25}{\"VALIDATION LSD : \"+str(self.val_lsd[-1]): ^50}{\"-\"*25:>25}\\n' \\\n f'{\"-\"*100}\\n{\"-\"*100}\\n\\n')\n \n def print_current_train(self, step, train_acc, train_loss, batch_loss, current_time, current_lr):\n if self.print_style==\"tight\":\n print(f\"| TRAIN DATA | STEP {self.current_step:>4}/{TRAIN_STEPS} | \" \\\n f\"ACC: {str(train_acc*100)[:5]:<5} – \" \\\n f\"LOSS: {str(train_loss)[:5]:<5} – \" \\\n f\"LR: {current_lr:.2e} \" \\\n f\"| | TIME | EPOCH: {str(round((current_time-self.epoch_start_time)/3600,1))+'h':<5} – \" \\\n f\"SUBSET: {str(round((current_time-self.batch_start_time)*self.verbose_frequency,1))+'s':<6} – \" \\\n f\"BATCH: {str(round(current_time-self.batch_start_time,1))+'s':<5} |\")\n else:\n print(f'\\n\\n{\"-\"*100}\\n{\"-\"*100}\\n' \\\n f'{\"-\"*25:<25}{\"CURRENT STEP : \"+str(step)+\" OF \"+str(TRAIN_STEPS): ^50}{\"-\"*25:>25}\\n' \\\n f'{\"-\"*100}\\n' \\\n f'{\"-\"*25:<25}{\"CURRENT TRAIN ACCURACY : \"+str(train_acc*100): ^50}{\"-\"*25:>25}\\n' \\\n f'{\"-\"*100}\\n' \\\n f'{\"-\"*25:<25}{\"CURRENT TRAIN LOSS : \"+str(train_loss): ^50}{\"-\"*25:>25}\\n' \\\n f'{\"-\"*100}\\n' \\\n f'{\"-\"*25:<25}{\"LAST BATCH LOSS : \"+str(batch_loss): ^50}{\"-\"*25:>25}\\n' \\\n f'{\"-\"*100}\\n{\"-\"*100}\\n' \\\n f'{\"-\"*25:<25}{\"EPOCH ELAPSED TIME : \"+str(round(current_time-self.epoch_start_time,1))+\" SECONDS\": ^50}{\"-\"*25:>25}\\n' \\\n f'{\"-\"*100}\\n' \\\n f'{\"-\"*25:<25}{\"LAST SET OF BATCHES TOOK : ~\"+str(round((current_time-self.batch_start_time)*self.verbose_frequency,1))+\" SECONDS\": ^50}{\"-\"*25:>25}\\n' \\\n f'{\"-\"*100}\\n' \\\n f'{\"-\"*25:<25}{\"LAST SINGLE BATCH TOOK : \"+str(round(current_time-self.batch_start_time,1))+\" SECONDS\": ^50}{\"-\"*25:>25}\\n' \\\n f'{\"-\"*100}\\n{\"-\"*100}\\n\\n')","metadata":{"papermill":{"duration":0.149635,"end_time":"2021-06-15T18:03:22.35309","exception":false,"start_time":"2021-06-15T18:03:22.203455","status":"completed"},"tags":[],"id":"typical-philosophy","execution":{"iopub.status.busy":"2022-01-25T14:55:38.281122Z","iopub.execute_input":"2022-01-25T14:55:38.281605Z","iopub.status.idle":"2022-01-25T14:55:38.297887Z","shell.execute_reply.started":"2022-01-25T14:55:38.28157Z","shell.execute_reply":"2022-01-25T14:55:38.29712Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 6.4 CUSTOM TRAIN LOOP\n\n---\n\nINFORMATION","metadata":{"papermill":{"duration":0.128686,"end_time":"2021-06-15T18:03:22.612099","exception":false,"start_time":"2021-06-15T18:03:22.483413","status":"completed"},"tags":[],"id":"egyptian-winner"}},{"cell_type":"code","source":"# Instantiate our tool for logging\nstat_logger = StatLogger()\n \nfor epoch in range(1,EPOCHS+1):\n print(f'\\n\\n{\"=\"*100}\\n{\"=\"*25:<25}{\"EPOCH #\"+str(epoch): ^50}{\"=\"*25:>25}\\n{\"=\"*100}\\n')\n \n stat_logger.current_step=0\n stat_logger.epoch_start_time = time.time() # to compute epoch duration\n \n # create distributed versions of dataset to run on TPU with 8 computation units\n train_dist_ds = strategy.experimental_distribute_dataset(train_ds)\n val_dist_ds = iter(strategy.experimental_distribute_dataset(val_ds))\n \n for image_batch, inchi_batch in train_dist_ds:\n \n # Update current step\n stat_logger.batch_start_time = time.time()\n \n # Update the current step\n stat_logger.current_step += 1\n \n # Calculate training step\n dist_train_step(image_batch, inchi_batch)\n \n # end of epoch validation step\n if stat_logger.current_step == TRAIN_STEPS and epoch%2==0:\n print(\"\\n... VALIDATION DATASET STATISTICS ... \\n\")\n for _ in range(VAL_STEPS):\n preds, lbls = dist_val_step(val_dist_ds)\n metrics[\"val_lsd\"].update_state(get_levenshtein_distance(preds, lbls))\n \n # Record this epochs statistics\n stat_logger.train_loss.append(metrics[\"train_loss\"].result().numpy())\n stat_logger.train_acc.append(metrics[\"train_acc\"].result().numpy())\n stat_logger.val_loss.append(metrics[\"val_loss\"].result().numpy())\n stat_logger.val_acc.append(metrics[\"val_acc\"].result().numpy())\n stat_logger.val_lsd.append(metrics[\"val_lsd\"].result().numpy())\n stat_logger.step.append(stat_logger.current_step)\n stat_logger.epoch.append(epoch)\n stat_logger.lr.append(lr_scheduler(tf.cast(stat_logger.current_step+TRAIN_STEPS*(epoch-1), tf.float32)))\n \n # Reset the validation metrics as one epoch should not effect the next\n metrics[\"val_lsd\"].reset_states()\n metrics[\"val_acc\"].reset_states()\n metrics[\"val_loss\"].reset_states()\n metrics[\"train_acc\"].reset_states()\n metrics[\"train_loss\"].reset_states()\n metrics[\"batch_loss\"].reset_states()\n \n # Print validation scores\n stat_logger.print_last_val(current_time=time.time())\n \n # verbose logging step\n if stat_logger.current_step % stat_logger.verbose_frequency == 0: \n stat_logger.print_current_train(\n stat_logger.current_step,\n metrics[\"train_acc\"].result().numpy(), \n metrics[\"train_loss\"].result().numpy(), \n metrics[\"batch_loss\"].result().numpy(), \n current_time=time.time(),\n current_lr=lr_scheduler(tf.cast(stat_logger.current_step+TRAIN_STEPS*(epoch-1), tf.float32))\n )\n metrics[\"train_acc\"].reset_states()\n metrics[\"train_loss\"].reset_states()\n metrics[\"batch_loss\"].reset_states()\n\n # stop training when NaN loss is detected\n if stat_logger.current_step == TRAIN_STEPS:\n break\n \n # update learning rate\n # lr_scheduler.step(stat_logger.current_step+((epoch-1)*TRAIN_STEPS))\n \n # Save every other epoch (starting with first epoch)\n # Save after last epoch too...\n # if epoch%2==1 or epoch==EPOCHS:\n # save weights\n print(\"\\n...SAVING MODELS TO DISK ... \\n\")\n transformer.save_weights(f'./transformer_epoch_{epoch}.h5')\n encoder.save_weights(f'./encoder_epoch_{epoch}.h5')","metadata":{"papermill":{"duration":14151.927869,"end_time":"2021-06-15T21:59:14.668674","exception":false,"start_time":"2021-06-15T18:03:22.740805","status":"completed"},"tags":[],"id":"centered-waste","outputId":"d23e91dd-cffa-411d-b0bc-3ecf907cd353","execution":{"iopub.status.busy":"2022-01-25T14:55:38.300385Z","iopub.execute_input":"2022-01-25T14:55:38.300719Z","iopub.status.idle":"2022-01-25T14:57:26.563688Z","shell.execute_reply.started":"2022-01-25T14:55:38.300676Z","shell.execute_reply":"2022-01-25T14:57:26.562403Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 6.5 JUST-IN-CASE SAVE\n\n---\n\nINFORMATION","metadata":{"papermill":{"duration":0.35207,"end_time":"2021-06-15T21:59:15.370195","exception":false,"start_time":"2021-06-15T21:59:15.018125","status":"completed"},"tags":[],"id":"historic-pricing"}},{"cell_type":"code","source":"# My thing crashed so I loaded the weights from the last stable epoch to continue\n\n# transformer.save_weights(f'./transformer_epoch_safety_save.h5')\n# encoder.save_weights(f'./encoder_epoch_safety_save.h5')","metadata":{"papermill":{"duration":0.365427,"end_time":"2021-06-15T21:59:16.085603","exception":false,"start_time":"2021-06-15T21:59:15.720176","status":"completed"},"tags":[],"id":"minor-trouble","execution":{"iopub.status.busy":"2022-01-25T14:57:26.565144Z","iopub.status.idle":"2022-01-25T14:57:26.565508Z","shell.execute_reply.started":"2022-01-25T14:57:26.565319Z","shell.execute_reply":"2022-01-25T14:57:26.565342Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n# 7. INFER ON TEST DATA \n\nIn this section we will use our trained model to generate the predictions we will use to submit to the competition","metadata":{"papermill":{"duration":0.34958,"end_time":"2021-06-15T21:59:16.788003","exception":false,"start_time":"2021-06-15T21:59:16.438423","status":"completed"},"tags":[],"id":"stopped-benchmark"}},{"cell_type":"markdown","source":"## 7.1 INDIVIDUAL TEST STEP (AND DISTRIBUTED)\n\n---\n\nINFORMATION","metadata":{"papermill":{"duration":0.348916,"end_time":"2021-06-15T21:59:17.524214","exception":false,"start_time":"2021-06-15T21:59:17.175298","status":"completed"},"tags":[],"id":"spanish-lindsay"}},{"cell_type":"code","source":"def test_step(_image_batch):\n \"\"\" \n Forward pass of the testing step\n \n Args:\n image_batch: Distributed batches of image dataset\n inchi_batch: Distributed batches of inchi dataset\n \n Returns:\n predictions_seq_batch: Predictions on the test batch\n \"\"\"\n \n transformer_pred_batch = tf.ones((REPLICA_BATCH_SIZE, 1), dtype=tf.uint8)\n \n # Get image embedding (once)\n _image_embedding = encoder(_image_batch, training=False)\n \n # Teacher forcing - feeding the target as the next input\n for c_idx in range(1, MAX_LEN):\n \n combined_mask = create_mask(None, transformer_pred_batch)\n \n # predictions.shape == (batch_size, seq_len, vocab_size)\n prediction_batch, attention_weights = transformer(_image_embedding, transformer_pred_batch, training=False, look_ahead_mask=combined_mask)\n predicted_batch_id = prediction_batch[:, -1:, :]\n \n # no teacher forcing, predicted char is next transformer input\n transformer_pred_batch = tf.concat([transformer_pred_batch, tf.cast(tf.argmax(predicted_batch_id, axis=-1), tf.uint8)], axis=-1)\n \n return transformer_pred_batch \n\n \n@tf.function\ndef distributed_test_step(_img_batch, _img_ids):\n per_replica_seqs = strategy.run(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\n\n\n@tf.function\ndef distributed_test_step_v2(_test_dist_ds):\n _test_image_batch, _test_id_batch = next(_test_dist_ds)\n predictions_seq_batch_per_replica = strategy.run(test_step, args=(_test_image_batch,))\n predictions_seq_batch_accum = strategy.gather(predictions_seq_batch_per_replica, axis=0)\n _test_id_batch_accum = strategy.gather(_test_id_batch, axis=0)\n return predictions_seq_batch_accum, _test_id_batch_accum","metadata":{"papermill":{"duration":0.365232,"end_time":"2021-06-15T21:59:18.236679","exception":false,"start_time":"2021-06-15T21:59:17.871447","status":"completed"},"tags":[],"id":"cathedral-regression","execution":{"iopub.status.busy":"2022-01-25T14:57:26.56676Z","iopub.status.idle":"2022-01-25T14:57:26.567157Z","shell.execute_reply.started":"2022-01-25T14:57:26.56696Z","shell.execute_reply":"2022-01-25T14:57:26.566986Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 7.2 RAW INFERENCE LOOP\n\n---\n\nINFORMATION","metadata":{"papermill":{"duration":0.348529,"end_time":"2021-06-15T21:59:18.935783","exception":false,"start_time":"2021-06-15T21:59:18.587254","status":"completed"},"tags":[],"id":"selective-batch"}},{"cell_type":"code","source":"# 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\n# Create an iterator\ndist_test_ds = iter(strategy.experimental_distribute_dataset(test_ds))\nfor 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, preds], axis=0)\n all_pred_ids = tf.concat([all_pred_ids, tf.expand_dims(pred_ids, axis=-1)], axis=0)","metadata":{"papermill":{"duration":1557.692168,"end_time":"2021-06-15T22:50:18.876364","exception":false,"start_time":"2021-06-15T22:24:21.184196","status":"completed"},"tags":[],"id":"national-anthropology","outputId":"7709bd32-512a-40a2-f36e-09bd38e00d9e","execution":{"iopub.status.busy":"2022-01-25T14:57:26.570593Z","iopub.status.idle":"2022-01-25T14:57:26.570972Z","shell.execute_reply.started":"2022-01-25T14:57:26.570761Z","shell.execute_reply":"2022-01-25T14:57:26.570782Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 7.3 TEST PRED POST-PROCESSING\n\n---\n\nINFORMATION","metadata":{"papermill":{"duration":0.351029,"end_time":"2021-06-15T22:50:19.579439","exception":false,"start_time":"2021-06-15T22:50:19.22841","status":"completed"},"tags":[],"id":"compact-stupid"}},{"cell_type":"code","source":"def arr_2_inchi(arr):\n ''' Function to convert array to inchi '''\n inchi_str = ''\n for i in arr:\n c = int_2_tok.get(i)\n if c==\"\n - Tutorial - Using Iterators
\n - Tutorial - Iterating Inside a
tf.function\n - Kaggle Discussion - TPU: Extreme Optimizations
\n - Kaggle Notebook - Custom Training Loop With 100+ Flowers on TPU
\n