2.6 INITIAL DATAFRAME INSTANTIATION
\n\n---\n","metadata":{"id":"8ChLuaTnUekk"}},{"cell_type":"code","source":"!gsutil ls gs://kds-ebd140840cd463059a4465c0dddd014bba523491bf03cfa668287a31/","metadata":{"id":"GIB3TnkdeHzB","executionInfo":{"status":"ok","timestamp":1621257408279,"user_tz":-480,"elapsed":1812,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"78427c38-4d00-43fd-9486-6ec6729bc355"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#copy gcs csv files locally\n!gsutil cp gs://kds-ebd140840cd463059a4465c0dddd014bba523491bf03cfa668287a31/train_labels_w_extra.csv .\n!gsutil cp gs://kds-ebd140840cd463059a4465c0dddd014bba523491bf03cfa668287a31/sample_submission_w_extra.csv .\n","metadata":{"id":"z-0Sq5EjdKJl","executionInfo":{"status":"ok","timestamp":1621257454213,"user_tz":-480,"elapsed":19256,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"95ffb372-9eb5-4870-d2b2-632deec3f8e7"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... INITIAL DATAFRAME INSTANTIATION STARTING ...\\n\")\n\n# Load the train and submission dataframes\ntrain_df = pd.read_csv(\"train_labels_w_extra.csv\")\nss_df = pd.read_csv(\"sample_submission_w_extra.csv\")\n\n# --- Distribution Information ---\nN_EX = len(train_df)\nN_TEST = len(ss_df)\nN_VAL = 80_000 # Fixed from dataset creation information\nN_TRAIN = N_EX-N_VAL\n\n# --- Batching Information ---\nBATCH_SIZE_DEBUG = 2\nREPLICA_BATCH_SIZE = 64 # Could probably be 128\n\nif DEBUG:\n REPLICA_BATCH_SIZE = BATCH_SIZE_DEBUG\nOVERALL_BATCH_SIZE = REPLICA_BATCH_SIZE*N_REPLICAS\n\n\n# --- Input Image Information ---\nIMG_SHAPE = (192,384,3)\n\n# --- Autocalculate Training/Validation/Testing Information ---\nTRAIN_STEPS = N_TRAIN // OVERALL_BATCH_SIZE\nVAL_STEPS = N_VAL // OVERALL_BATCH_SIZE\nTEST_STEPS = int(np.ceil(N_TEST/OVERALL_BATCH_SIZE))\n\n# This is for padding our test dataset so we only have whole batches\nREQUIRED_DATASET_PAD = OVERALL_BATCH_SIZE-N_TEST%OVERALL_BATCH_SIZE\n\n# --- Modelling Information ---\nATTN_EMB_DIM = 192\nN_RNN_UNITS = 512\n\nprint(f\"\\n... # OF TRAIN+VAL EXAMPLES : {N_EX:<7} ...\")\nprint(f\"... # OF TRAIN EXAMPLES : {N_TRAIN:<7} ...\")\nprint(f\"... # OF VALIDATION EXAMPLES : {N_VAL:<7} ...\")\nprint(f\"... # OF TEST EXAMPLES : {N_TEST:<7} ...\\n\")\n\nprint(f\"\\n... REPLICA BATCH SIZE : {REPLICA_BATCH_SIZE} ...\")\nprint(f\"... OVERALL BATCH SIZE : {OVERALL_BATCH_SIZE} ...\\n\")\n\nprint(f\"\\n... IMAGE SHAPE : {IMG_SHAPE} ...\\n\")\n\nprint(f\"\\n... TRAIN STEPS PER EPOCH : {TRAIN_STEPS:<5} ...\")\nprint(f\"... VAL STEPS PER EPOCH : {VAL_STEPS:<5} ...\")\nprint(f\"... TEST STEPS PER EPOCH : {TEST_STEPS:<5} ...\\n\")\n\nprint(f\"\\n... ATTENTION EMBEDDING DIMENSION : {ATTN_EMB_DIM:<5} ...\")\nprint(f\"... NUMBER OF UNITS IN LSTM : {N_RNN_UNITS:<5} ...\\n\")\n\nprint(\"\\n... TRAIN DATAFRAME ...\\n\")\ndisplay(train_df.head(3))\n\nprint(\"\\n... SUBMISSION DATAFRAME ...\\n\")\ndisplay(ss_df.head(3))\n\nprint(\"\\n... INITIAL DATAFRAME INSTANTIATION COMPLETED...\\n\")","metadata":{"id":"rgT8z9wEUekl","executionInfo":{"status":"ok","timestamp":1621257535594,"user_tz":-480,"elapsed":21298,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"cfbfdb6d-c289-4b6e-cbe2-01653191c475","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"2.7 USER INPUT VARIABLES
\n\n---\n","metadata":{"id":"uc5UzxzBUekl"}},{"cell_type":"code","source":"print(\"\\n... SPECIAL VARIABLE SETUP STARTING ...\\n\")\n\n\n# Whether to start training using previously checkpointed model\nLOAD_MODEL = False\nENCODER_CKPT_PATH = \"\"\nDECODER_CKPT_PATH = \"\"\n\nif LOAD_MODEL:\n print(f\"\\n... ENCODER MODEL TRAINING WILL RESUME FROM PREVIOUS CHECKPOINT:\\n\\t-->{ENCODER_CKPT_PATH}\\n\")\n print(f\"... DECODER MODEL TRAINING WILL RESUME FROM PREVIOUS CHECKPOINT:\\n\\t-->{DECODER_CKPT_PATH}\\n\")\nelse:\n print(f\"\\n... MODEL TRAINING WILL START FROM SCRATCH ...\\n\")\n\n \nprint(\"\\n... SPECIAL VARIABLE SETUP COMPLETED ...\\n\")","metadata":{"id":"MXRcm7hfUekm","executionInfo":{"status":"ok","timestamp":1621257535596,"user_tz":-480,"elapsed":21290,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"3c4661cf-15da-4a15-fb22-c35ec37f44cb","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n\n\n\n\n\n
\n 3 HELPER FUNCTION & CLASSESS ⤒\n
","metadata":{"id":"91fOX6bpUekm"}},{"cell_type":"markdown","source":"3.1 GENERAL HELPER FUNCTIONS
\n\n---","metadata":{"id":"ipszRYU3Uekn"}},{"cell_type":"code","source":"def flatten_l_o_l(nested_list):\n \"\"\" Flatten a list of lists \"\"\"\n return [item for sublist in nested_list for item in sublist]\n\n\ndef tf_load_image(path, img_size=(192,384,3), invert=False):\n \"\"\" Load an image with the correct size and shape \n \n Args:\n path (tf.string): Path to the image to be loaded\n img_size (tuple, optional): Size to reshape image to (required for TPU)\n tile_to_3_channel (bool, optional): Whether to tile the single channel\n image to 3 channels which will be required for most off-the-shelf models\n invert (bool, optional): Whether or not to invert the background/foreground\n \n Returns:\n 3 channel tf.Constant image ready for training/inference\n \n \"\"\"\n img = decode_img(tf.io.read_file(path), img_size, n_channels=3, invert=invert) \n return img\n \n \ndef decode_image(image_data, resize_to=(192,384,3)):\n \"\"\" Function to decode the tf.string containing image information \n \n \n Args:\n image_data (tf.string): String containing encoded image data from tf.Example\n resize_to (tuple, optional): Size that we will reshape the tensor to (required for TPU)\n \n Returns:\n Tensor containing the resized single-channel image in the appropriate dtype\n \"\"\"\n image = tf.image.decode_png(image_data, channels=3)\n image = tf.reshape(image, resize_to)\n return tf.cast(image, TARGET_DTYPE)\n \n \n# sparse tensors are required to compute the Levenshtein distance\ndef dense_to_sparse(dense):\n \"\"\" Convert a dense tensor to a sparse tensor \n \n Args:\n dense (Tensor): TBD\n \n Returns:\n A sparse tensor \n \"\"\"\n indices = tf.where(tf.ones_like(dense))\n values = tf.reshape(dense, (MAX_LEN*OVERALL_BATCH_SIZE,))\n sparse = tf.SparseTensor(indices, values, dense.shape)\n return sparse\n\ndef get_levenshtein_distance(preds, lbls):\n \"\"\" Computes the Levenshtein distance between the predictions and labels \n \n Args:\n preds (tensor): Batch of predictions\n lbls (tensor): Batch of labels\n \n Returns:\n The mean Levenshtein distance calculated across the batch\n \"\"\"\n preds = tf.where(tf.not_equal(lbls, END_TOKEN) & tf.not_equal(lbls, PAD_TOKEN), preds, 0)\n lbls = tf.where(tf.not_equal(lbls, END_TOKEN), lbls, 0)\n\n preds_sparse = dense_to_sparse(preds)\n lbls_sparse = dense_to_sparse(lbls)\n\n batch_distance = tf.edit_distance(preds_sparse, lbls_sparse, normalize=False)\n mean_distance = tf.math.reduce_mean(batch_distance)\n \n return mean_distance","metadata":{"id":"tRHd39nmUekn","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n\n\n\n\n\n
4 PREPARE THE DATASET ⤒
\n\nIn this section we prepare the **`tf.data.Datasets`** we will use for training and validation","metadata":{"id":"KFbLZ1F9Ueko"}},{"cell_type":"code","source":"print(\"\\n\\n... STARTING PREPARING VARIABLES FOR DATASET ...\\n\")\n\ntok_2_int = {c.strip(\"\\\\\"):i for i,c in enumerate(TOKEN_LIST)}\nint_2_tok = {v:k for k,v in tok_2_int.items()}\n\n# Max Length Was Determined Previously Using... \n# >>> MAX_LEN = train_df.InChI.progress_apply(lambda x: len(re.findall(\"|\".join(TOKEN_LIST), x))).max()+1\nMAX_LEN = ((train_df.inchi_token_len.max()+1)//2) # //2 yields 138... which is half of max length (speeds up training)\nVOCAB_LEN = len(int_2_tok)\n\nprint(f\"\\t--> TOKEN TO INTEGER MAP : {tok_2_int}\")\nprint(f\"\\t--> INTEGER TO TOKEN MAP : {int_2_tok}\")\nprint(f\"\\t--> MAX # OF TOKENS IN INCHI : {MAX_LEN}\")\nprint(f\"\\t--> LENGTH OF VOCAB : {VOCAB_LEN}\")\n\nprint(f\"\\n\\n\\t--> CONVERTED INCHI STRINGS :\")\nfor i, row in train_df.iloc[:N_VAL].sample(3).iterrows():\n print(f\"\\n\\t\\t--> EXAMPLE #{i} FROM THE VALIDATION DATASET\")\n print(\"\\t\\t\\t--> RAW INCHI : \", row[\"InChI\"])\n\nprint(\"\\n\\n... PREPARING VARIABLES FOR DATASET COMPLETED ...\\n\")","metadata":{"id":"W2tNE5sZUeko","executionInfo":{"status":"ok","timestamp":1621257535599,"user_tz":-480,"elapsed":21279,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"aacb19b7-8d09-4bfb-f240-c88358954b1c","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"4.1 READ TFRECORD FILES - CREATE THE RAW DATASET(S)
\n\n---\n\nHere we will leverage **`tf.data.TFRecordDataset`** to read the TFRecord files.\n* The simplest way is to specify a list of filenames (paths) of TFRecord files.\n* It is a subclass of **`tf.data.Dataset`**.\n\nThis newly created raw dataset contains **`tf.train.Example`** messages, and when iterated over it, we get scalar string tensors.","metadata":{"id":"LxjV_QmNUekp"}},{"cell_type":"code","source":"print(\"\\n... CREATE TFRECORD RAW DATASETS STARTING ...\\n\")\n\n# Create tf.data.Dataset from filepaths for conversion later\nraw_train_ds = tf.data.TFRecordDataset(TRAIN_TFREC_PATHS, num_parallel_reads=None)\nraw_val_ds = tf.data.TFRecordDataset(VAL_TFREC_PATHS, num_parallel_reads=None)\nraw_test_ds = tf.data.TFRecordDataset(TEST_TFREC_PATHS, num_parallel_reads=None)\n\n# raw_test_ds = tf.data.TFRecordDataset(TEST_TFREC_PATHS, num_parallel_reads=None)\n\nprint(f\"\\n... THE RAW TF.DATA.TFRECORDDATASET OBJECT:\\n\\t--> {raw_train_ds}\\n\")\n\nprint(\"\\n... CREATE TFRECORD RAW DATASETS COMPLETED ...\\n\")","metadata":{"id":"bZgcKtYCUekp","executionInfo":{"status":"ok","timestamp":1621257535600,"user_tz":-480,"elapsed":21270,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"229bcfb7-c45c-4963-dc36-866da9f522fb","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"4.2 WHAT TO DO IF YOU DON'T KNOW THE FEATURE DESCRIPTIONS OF THE DATASET?
\n\n---\n\nIf you are the author who created the TFRecord files, you definitely know how to define the feature description to parse the raw dataset.\n\nOtherwise, you can use like\n\n```python\nexample = tf.train.Example()\nexample.ParseFromString(serialized_example.numpy())\n```\n\nto check the information. You will get something like\n\n```python\nfeatures {\n feature {\n key: \"class\"\n value {\n int64_list {\n value: 57\n }\n }\n }\n feature {\n key: \"id\"\n value {\n bytes_list {\n value: \"338ab7bac\"\n }\n }\n }\n feature {\n key: \"image\"\n value {\n bytes_list {\n value: ...\n }\n }\n }\n ...\n}\n```\n\nThis should give you enough information to define the feature description.","metadata":{"id":"9XEcihC_Uekp"}},{"cell_type":"code","source":"print(\"\\n... RAW TFRECORD INVESTIGATION TO DETERMINE FEATURE DESCRIPTIONS STARTED ...\\n\")\n\nprint(\"\\n... EXAMPLE OF TRUNCATED RAW TFRECORD/TFEXAMPLE FROM TRAINING DATASET TO SHOW HOW TO FIND FEATURE DESCRIPTIONS:\\n\")\n# See an example\nfor raw in raw_train_ds.take(1):\n example = tf.train.Example()\n example.ParseFromString(raw.numpy())\n for i, (k,v) in enumerate(example.features.feature.items()):\n print(f\"\\tFEATURE #{i+1}\")\n print(f\"\\t\\t--> KEY = {k}\")\n if k!=\"image\":\n try:\n print(f\"\\t\\t\\t--> TRUNCATED-VALUE = {v.int64_list.value[:15]} ...\\n\")\n except:\n print(f\"\\t\\t\\t--> TRUNCATED-VALUE = {v.bytes_list.value[0][:25]} ...\\n\")\n else:\n print(f\"\\t\\t\\t--> TRUNCATED-VALUE = {str(v.bytes_list.value[0][:25])} ...\\n\") \n\nprint(\"\\n... RAW TFRECORD INVESTIGATION TO DETERMINE FEATURE DESCRIPTIONS COMPLETED ...\\n\")","metadata":{"id":"W4vqT1TRUekq","executionInfo":{"status":"ok","timestamp":1621257538624,"user_tz":-480,"elapsed":24290,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"be2a673f-c4c0-4eab-a401-f4f529a26a69","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"4.3 PARSE THE RAW DATASET(S)
\n\n---\n\n\nThe general recipe to parse the string tensors in the raw dataset looks something like this:\n\n\n\n**STEP 1.** Create a description of the features. For example:\n\n```python\nfeature_description = { \n 'feature0': tf.io.FixedLenFeature([], tf.int64),\n 'feature1': tf.io.FixedLenFeature([], tf.string),\n 'feature2': tf.io.FixedLenFeature([], tf.float32),\n ...\n}\n```\n\n
\n\n**STEP 2.** Define a parsing function by using `tf.io.parse_single_example` and the defined feature description.\n```python\ndef _parse_function(example):\n \"\"\"\n Args:\n example: A string tensor representing a `tf.train.Example`.\n \"\"\"\n\n # Parse `example`.\n parsed_example = tf.io.parse_single_example(example, feature_description)\n\n return parsed_example\n```\n\n
\n\n**STEP 3.** Map the raw dataset by `_parse_function`.\n```python\ndataset = raw_dataset.map(_parse_function)\n```\n\n
\n\n---\n\n
\n\n**In the following cell, we apply the above recipe to our BMS tfrecord dataset.**\n\n\n
\n π NOTE:
- The parsed images are
- The InChI strings and Image IDs will just be left as byte string tensors.\n
\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 π REFERENCE:
\n - Tutorial - TFRecord and tf.Example
\n - TFRecordDataset Documentation
\n - Decoding PNGs Documentation
\n
\n","metadata":{"id":"gMasPEdIUekq"}},{"cell_type":"code","source":"def decode(serialized_example, is_test=False, tokenized_inchi=True):\n \"\"\" Parses a set of features and label from the given `serialized_example`.\n \n It is used as a map function for `dataset.map`\n\n Args:\n serialized_example (tf.Example): A serialized example containing the\n following features:\n β 'image'\n βΒ 'image_id'\n β 'inchi'\n is_test (bool, optional): Whether to allow for the InChI feature\n drop_id (bool, optional): Whether or not to drop the ID feature\n \n Returns:\n A decoded tf.data.Dataset object representing the tfrecord dataset\n \"\"\"\n feature_dict = {\n 'image': tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value=''),\n }\n \n if not is_test:\n if tokenized_inchi:\n feature_dict[\"inchi\"] = tf.io.FixedLenFeature(shape=[MAX_LEN], dtype=tf.int64, default_value=[0]*MAX_LEN)\n else:\n feature_dict[\"inchi\"] = tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value='')\n else:\n feature_dict[\"image_id\"] = tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value='')\n \n # Define a parser\n features = tf.io.parse_single_example(serialized_example, features=feature_dict)\n \n # Decode the tf.string\n image = decode_image(features['image'], resize_to=IMG_SHAPE)\n \n # Figure out the correct information to return\n if is_test:\n image_id = features[\"image_id\"] \n return image, image_id\n else:\n if tokenized_inchi:\n target = tf.cast(features[\"inchi\"], tf.uint8)\n else:\n target = features[\"inchi\"]\n return image, target","metadata":{"id":"kPtkLlkEUekr","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... DECODING RAW TFRECORD DATASETS STARTING ...\\n\")\n\n# Decode the tfrecords completely ββ decode is our `_parse_function` (from recipe above)\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":{"id":"E26ZA6f2Uekr","executionInfo":{"status":"ok","timestamp":1621257539992,"user_tz":-480,"elapsed":25649,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"c50f2685-1d28-434c-ba09-3be90ba62dee","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n - Tutorial - TFRecord and tf.Example
\n - TFRecordDataset Documentation
\n - Decoding PNGs Documentation
\n
4.4 WORKING WITH `TF.DATA.DATASET` OBJECTS
\n\n---\n\nWith the above parsing methods defined, we can define how to load the dataset with more options and further apply shuffling, bacthing, etc. In particular the following methods and attributes are of special interest to us:\n* Use **`num_parallel_reads`** in **`tf.data.TFRecordDataset`** to read files in parallel.\n* Set **`tf.data.Options.experimental_deterministic=False`** and use it to get a new dataset that ignores the order of elements.\n* Use **`num_parallel_calls`** in **`tf.data.Dataset.map()`** method to have parallel processing.\n* Use **`tf.data.Dataset.prefetch()`** to allow later batches to be prepared while the current batch is being processed.\n* Use **`tf.data.AUTOTUNE`** to automatically determine parallelization argument values\n\nThe parallel processing and prefetching are particular important when working with TPU:\n* This is because a TPU can process batches very quickly\n* The dataset pipeline should be able to provide data for TPU efficiently, otherwise the TPU will be idle.\n\n**In the cell below we will create the functions and configuration template which will later be used to create our respective datasets**\n\n\n π REFERENCE:
\n - Guide - tf.data: Build TensorFlow Input Pipelines
\n - Guide - Better Performance With the tf.data API
\n - tf.data.Dataset Documentation
\n
","metadata":{"id":"iEGRNrXoUeks"}},{"cell_type":"code","source":"# def load_dataset(filenames, is_test=False, ordered=False, tokenized_inchi=True):\n# \"\"\"Read from TFRecords.\n \n# For optimal performance, reading from multiple files at once and disregarding data order (if `ordered=False`).\n# - If pulling InChI from TFRecords than order does not matter since we will \n# be shuffling the data anyway (for training dataset).\n \n# Args:\n# filenames (list of strings): List of paths to that point to the respective TFRecord files\n# is_test (bool, optional): Whether or not to include the image ID or label in the returned dataset\n# ordered (bool, optional): Whether to ensured ordered results or maximize parallelization\n# tokenized_inchi (bool, optional): Whether our dataset includes the tokenized inchi or we will be \n# creating it from the caption numpy array\n \n# Returns:\n# Decoded tf.data.Dataset object\n# \"\"\"\n\n# options = tf.data.Options()\n# if not ordered:\n# # disable order, increase speed\n# options.experimental_deterministic = False\n# N_PARALLEL=tf.data.AUTOTUNE\n# else:\n# N_PARALLEL=None\n \n# # If not-ordered, this will read in by automatically interleaving multiple tfrecord files.\n# dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=N_PARALLEL)\n \n# # If not-ordered, this will ensure that we use data as soon as it \n# # streams in, rather than in its original order.\n# dataset = dataset.with_options(options) \n \n# # parse and return a dataset w/ the appropriate configuration\n# dataset = dataset.map(\n# lambda x: decode(x, is_test, tokenized_inchi),\n# num_parallel_calls=N_PARALLEL,\n# )\n \n# return dataset\n\n# def get_dataset(filenames, batch_size, \n# is_test=False, \n# shuffle_buffer_size=1, \n# repeat_dataset=True, \n# preserve_file_order=False, \n# drop_remainder=True,\n# tokenized_inchi=True,\n# external_inchi_dataset=None):\n# \"\"\" Get a tf.data.Dataset w/ the appropriate configuration\n \n# Args:\n# TBD\n \n# Returns:\n# TBD\n \n# \"\"\"\n# # Load the dataset\n# dataset = load_dataset(filenames, is_test, preserve_file_order, tokenized_inchi)\n \n# # If 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\n\n\ndef load_dataset(filenames, is_test=False, ordered=False, tokenized_inchi=True):\n \"\"\"Read from TFRecords.\n \n For optimal performance, reading from multiple files at once and disregarding data order (if `ordered=False`).\n - If pulling InChI from TFRecords than order does not matter since we will \n be shuffling the data anyway (for training dataset).\n \n Args:\n filenames (list of strings): List of paths to that point to the respective TFRecord files\n is_test (bool, optional): Whether or not to include the image ID or label in the returned dataset\n ordered (bool, optional): Whether to ensured ordered results or maximize parallelization\n tokenized_inchi (bool, optional): Whether our dataset includes the tokenized inchi or we will be \n creating it from the caption numpy array\n \n Returns:\n Decoded tf.data.Dataset object\n \"\"\"\n\n options = tf.data.Options()\n if not ordered:\n # disable order, increase speed\n options.experimental_deterministic = False\n N_PARALLEL=tf.data.AUTOTUNE\n else:\n N_PARALLEL=None\n \n # If not-ordered, this will read in by automatically interleaving multiple tfrecord files.\n dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=N_PARALLEL)\n \n # If not-ordered, this will ensure that we use data as soon as it \n # streams in, rather than in its original order.\n dataset = dataset.with_options(options) \n \n # parse and return a dataset w/ the appropriate configuration\n dataset = dataset.map(\n lambda x: decode(x, is_test, tokenized_inchi),\n num_parallel_calls=N_PARALLEL,\n )\n \n return dataset\n\ndef get_dataset(filenames, batch_size, \n is_test=False, \n shuffle_buffer_size=1, \n repeat_dataset=True, \n preserve_file_order=False, \n drop_remainder=True,\n tokenized_inchi=True,\n external_inchi_dataset=None,\n test_padding=0):\n \"\"\" Get a tf.data.Dataset w/ the appropriate configuration\n \n Args:\n TBD\n test_padding (int, optional): Amount required to pad dataset to have only full batches\n \n Returns:\n TBD\n \n \"\"\"\n # Load the dataset\n dataset = load_dataset(filenames, is_test, preserve_file_order, tokenized_inchi)\n \n if test_padding!=0:\n pad_dataset = tf.data.Dataset.from_tensor_slices((\n tf.zeros((test_padding, *IMG_SHAPE), dtype=TARGET_DTYPE), # Fake Images\n tf.constant([\"000000000000\",]*test_padding, dtype=tf.string)) # Fake IDs\n )\n dataset = dataset.concatenate(pad_dataset)\n \n # If we are training than we will want to repeat the dataset. \n # We will determine the number of steps (or updates) later for 1 training epoch.\n if repeat_dataset:\n dataset = dataset.repeat()\n \n # If we need to add on manually the inchi\n if external_inchi_dataset is not None:\n # Zip the datasets and tile the 1 channel image to 3 channels & drop the old inchi value\n dataset = tf.data.Dataset.zip((dataset, external_inchi_dataset))\n dataset = dataset.map(lambda x,y: (tf.tile(tf.expand_dims(x[0], -1), tf.constant([1,1,3], tf.int32)), y))\n \n # Shuffling\n if shuffle_buffer_size!=1:\n dataset = dataset.shuffle(shuffle_buffer_size)\n \n # Batching\n dataset = dataset.batch(batch_size, drop_remainder=drop_remainder)\n \n # prefetch next batch while training (autotune prefetch buffer size)\n dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)\n \n return dataset","metadata":{"id":"oQ4noX9_Ueks","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"####### ####### ####### ####### ####### ####### ####### #######\n\n# Template Configuration\nDS_TEMPLATE_CONFIG = dict(\n filenames=[],\n batch_size=1,\n is_test=False, \n shuffle_buffer_size=1, \n repeat_dataset=True, \n preserve_file_order=False, \n drop_remainder=True,\n tokenized_inchi=True,\n external_inchi_dataset=None,\n test_padding=0\n)\n\n####### ####### ####### ####### ####### ####### ####### #######\n\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=False,\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, VAL_DS_CONFIG, TEST_DS_CONFIG]): # , 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":{"id":"_x_EBOMyUekt","executionInfo":{"status":"ok","timestamp":1621257554969,"user_tz":-480,"elapsed":40617,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"3219b420-0e17-464f-ba8f-a2f51d229046","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\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\n\n\n\n
5 MODEL PREPERATION ⤒
\n\nIn this section we prepare the models for training. We will be using a model architecture very similar to that found within the [**Show, Attend, and Tell Research Paper**](https://arxiv.org/pdf/1502.03044.pdf).\n\n\n\n
![](\"https://kelvinxu.github.io/projects/diags/model_diag.png\")
","metadata":{"id":"TTbfDAVyUeku"}},{"cell_type":"markdown","source":"
5.1 UNDERSTANDING THE MODELS - ENCODER
\n\n---\n\nWe will be leveraging an [**EfficientNetV2**](https://arxiv.org/pdf/2104.00298) model to act as the Encoder CNN in our network. \n* On **TPU/GPU/CPU** we will use an **EfficientNetV2-B2** model\n\n\n\n***Basic View of EfficientNetB0 Architecture w/ 380x380x3 Input ... this is very similar to EfficientNetV2***\n
\n\nOur encoder will create feature maps for each image which will in turn be passed to the decoder side of the network. \n\n
\n\n**In the following cell, we will create a function to generate our encoder model.**\n\n\n
\n π NOTE:
- For different encoder architectures we will have a different number of feature maps. i.e. If we utilized EfficientNetV2B7 we would have 2560 feature maps instead of the 1280 feature maps that EfficientNetV2B0 produces. \n
\n\n- For different encoder architectures we will have a different number of feature maps. i.e. If we utilized EfficientNetV2B7 we would have 2560 feature maps instead of the 1280 feature maps that EfficientNetV2B0 produces. \n
\n π REFERENCE:
\n - TF Tutorial β Transformer Model for Language Understanding
\n - TF Tutorial β Image Captioning
\n - TF Tutorial β Neural Machine Translation w/ Attention
\n
","metadata":{"id":"60DcSXjkUeku"}},{"cell_type":"code","source":"# Will yield, at the lowest feature level, (6,12,208) or (72,208)\ndef get_efficientnetv2_backbone(model_name, include_top=False, input_shape=(192,384,3), pooling=None, weights=None):\n # Catch unsupported arguments\n if pooling or weights or include_top:\n raise NotImplementedError(\"\\n...At this time we only want to use the raw \" \\\n \"(no pretraining), headless, features with no pooling ...\\n\")\n backbone = effnetv2_model.EffNetV2Model(model_name=model_name) \n backbone(tf.ones((BATCH_SIZE_DEBUG,*input_shape)), training=False, features_only=True)\n return backbone","metadata":{"id":"GKQ5l39sUekv","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... ENCODER MODEL CREATION STARTING ...\\n\")\n\n# SAMPLE IMAGES\nSAMPLE_IMGS, SAMPLE_LBLS = next(iter(val_ds.unbatch().batch(BATCH_SIZE_DEBUG)))\n\n# ENCODER_CONFIG\nPREPROCESSING_FN = tf.keras.applications.efficientnet.preprocess_input\nBB_FN = get_efficientnetv2_backbone\n\n# This will be the dimension the network outputs flattened\nIMG_EMB_DIM = (6,12,208)\nprint(f\"\\n... RAW BB OUTPUT: {IMG_EMB_DIM} ...\\n\")\nIMG_EMB_DIM = (IMG_EMB_DIM[0]*IMG_EMB_DIM[1], IMG_EMB_DIM[2])\n\nclass Encoder(tf.keras.Model):\n def __init__(self, image_embedding_dim, preprocessing_fn, backbone_fn, image_shape, do_permute=False, include_top=False, pretrained_weights=None, scale_factor=0):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n super(Encoder, self).__init__()\n \n self.image_embedding_dim = image_embedding_dim\n self.preprocessing_fn = preprocessing_fn\n self.encoder_backbone = backbone_fn(model_name=EV2_NAME, include_top=include_top, weights=pretrained_weights, input_shape=image_shape) \n self.reshape = tf.keras.layers.Reshape(self.image_embedding_dim, name='image_embedding')\n self.permute = tf.keras.layers.Permute([2, 1], name='permute_features')\n self.do_permute = do_permute\n self.include_top = include_top\n self.scale_factor = scale_factor\n \n def call(self, x, training):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n x = self.preprocessing_fn(x)\n x = self.encoder_backbone(x, training=training, features_only=not self.include_top)[self.scale_factor]\n x = self.reshape(x, training=training)\n if self.do_permute:\n x = self.permute(x, training=training)\n return x\n \n# Example enoder output\nwith tf.device('/CPU:0'):\n encoder = Encoder(IMG_EMB_DIM, PREPROCESSING_FN, BB_FN, IMG_SHAPE, do_permute=IMG_EMB_DIM[1]\n - TF Tutorial β Transformer Model for Language Understanding
\n - TF Tutorial β Image Captioning
\n - TF Tutorial β Neural Machine Translation w/ Attention
\n
\n\n* **Recurrent Mechanism β Embedding & LSTM/GRU [(Partly Quoting From Mark Wijkhuizen's Notebook)](https://www.kaggle.com/markwijkhuizen/tensorflow-tpu-training-baseline-lb-16-92)**\n * The Embedding layer will learn a representation to convert the chemical token/characters into a representative vector\n * The LSTMCell takes a concatenated context vector from the attention mechanism and an embedded token value as input. \n * The LSTMCell hidden and carry states are initialized using the encoder features.\n\n
\n\n***Basic Overview of the Image Captioning Process With Some Rudimentary Math***\n
![\"attn_gif\"](\"https://i.ibb.co/H2fkVpL/attn.gif\")
\n\n**In the following cell, we will create a function to generate our decoder model.**\n\n
\n π NOTE:
- We decided to utilize an LSTM cell for the Recurrent portion of our decoder network. This decision is based on the concept that LSTMs can generally handle longer sequences better than GRUs. That being said, other options do exist. \n
\n\n- We decided to utilize an LSTM cell for the Recurrent portion of our decoder network. This decision is based on the concept that LSTMs can generally handle longer sequences better than GRUs. That being said, other options do exist. \n
\n π REFERENCE:
\n - TF Tutorial β Transformer Model for Language Understanding
\n - TF Tutorial β Image Captioning
\n - TF Tutorial β Neural Machine Translation w/ Attention
\n
","metadata":{"id":"7NrMUV-QUekv"}},{"cell_type":"code","source":"print(\"\\n... ATTENTION MECHANISM LAYER CREATION STARTING ...\\n\")\n# TO BE REPLACED W/ --> tfa.seq2seq.BahdanauAttention\nclass BahdanauAttention(tf.keras.layers.Layer):\n def __init__(self, attn_emb_depth):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n super(BahdanauAttention, self).__init__()\n self.W1 = tf.keras.layers.Dense(attn_emb_depth)\n self.W2 = tf.keras.layers.Dense(attn_emb_depth)\n self.V = tf.keras.layers.Dense(1)\n\n def call(self, hidden, features, training):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n # hidden shape == (batch_size, hidden size)\n # hidden_with_time_axis shape == (batch_size, 1, hidden size)\n # features shape == (batch_size, img_emb_size, hidden size)\n # we are doing this to broadcast addition along the time axis to calculate the score\n hidden_with_time_axis = tf.expand_dims(hidden, 1)\n\n # score shape == (batch_size, img_emb_size, 1)\n # we get 1 at the last axis because we are applying score to self.V\n # img_emb_size refers to the product of the image embedding width and height (144 for EB)\n # the shape of the tensor before applying self.V is (batch_size, img_emb_size, attention_emb_depth)\n score = self.V(\n tf.nn.tanh(self.W1(hidden_with_time_axis, training=training) + \\\n self.W2(features, training=training)), \n training=training\n )\n\n # attention_weights shape (before & after) == (batch_size, img_emb_size, 1)\n attention_weights = tf.nn.softmax(score, axis=1)\n\n # context_vector shape after multiplication is (batch_size, img_emb_size, 1)\n context_vector = attention_weights * features\n \n # context_vector shape after sum == (batch_size, hidden_size)\n context_vector = tf.reduce_sum(context_vector, axis=1) # this is axis 2 in some implementations?\n \n return context_vector, attention_weights\n \n \n# ATTN_EMB_DIM = 512\n# N_RNN_UNITS = 1024\nwith tf.device('/CPU:0'):\n attn_layer = BahdanauAttention(ATTN_EMB_DIM)\n context_vector, attn_weights = attn_layer(tf.zeros([BATCH_SIZE_DEBUG, IMG_EMB_DIM[0]]), img_embedding_batch)\n\nprint(f'\\n... Context Vector Shape : (batch_size, img_emb_dim_depth) : {context_vector.shape} ...')\nprint(f'... Attention Weights Shape : (batch_size, img_emb_dim_shape, 1) : {attn_weights.shape} ...\\n')\n\nprint(\"\\n... ATTENTION MECHANISM LAYER CREATION FINISHED ...\\n\")","metadata":{"id":"EDcKwj7LUekw","executionInfo":{"status":"ok","timestamp":1621257568663,"user_tz":-480,"elapsed":54294,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"c63ae0ee-9f22-4ac2-e03d-aa44949f1be0","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... DECODER MODEL CREATION STARTING ...\\n\")\n\nclass Decoder(tf.keras.Model):\n def __init__(self, vocab_len, attn_emb_depth, img_emb_dim, n_rnn_units, dropout_rate=0.25):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n super(Decoder, self).__init__()\n \n # Basic Parameters\n self.vocab_len = vocab_len\n self.attn_emb_depth = attn_emb_depth\n self.img_emb_dim = img_emb_dim\n self.n_rnn_units = n_rnn_units\n self.dropout_rate = dropout_rate\n \n # Attention Mechanism\n self.attention_layer = BahdanauAttention(self.attn_emb_depth)\n \n # LSTM hidden and carry state initialization\n self.init_h_layer = tf.keras.layers.Dense(\n units=n_rnn_units, input_shape=[img_emb_dim], name='img_embedding__hidden_init_layer'\n )\n self.init_c_layer = tf.keras.layers.Dense(\n units=n_rnn_units, input_shape=[img_emb_dim], name='img_embedding__input_act_init_layer'\n )\n \n # Character Embedding Layer In The Decoder\n # - TODO ... `mask_zero=True` ???\n self.embedding_layer = tf.keras.layers.Embedding(vocab_len, n_rnn_units, )\n \n # The LSTM cell\n self.lstm_cell_layer = tf.keras.layers.LSTMCell(n_rnn_units, name='lstm_cell_layer')\n \n # Dropout Layer to Prevent Overfitting\n self.dropout_layer = tf.keras.layers.Dropout(self.dropout_rate, name='dropout_layer')\n \n # Fully Connected Prediction Layer \n self.fcn_layer = tf.keras.layers.Dense(units=vocab_len, input_shape=[n_rnn_units], dtype=tf.float32, name='fc_prediction_layer')\n\n def call(self, token, hidden, memory, img_emb, training):\n \"\"\" TODO\n \n Args:\n TODO \n \n Returns:\n TODO\n \"\"\"\n # img_emb shape == (batch_size, img_emb_shape, hidden_size)\n # --> We ignore the attention weights for now (_)\n context_vector, _ = self.attention_layer(hidden, img_emb, training=training)\n\n # shape after passing token through embedding == (batch_size, 1, n_rnn_units)\n # shape after passing through squeeze == (batch_size, n_rnn_units)\n tok_emb = tf.squeeze(self.embedding_layer(token, training=training), axis=1)\n \n # shape after concatenation == (batch_size, n_rnn_units + hidden_size)\n x = tf.concat((context_vector, tok_emb), axis=-1)\n\n # passing the concatenated vector to the LSTM cell\n # - also getting the new hidden (h) and memory (c) vectors\n _, (hidden_new, memory_new) = self.lstm_cell_layer(x, (hidden, memory), training=training)\n \n #### ##### ##### ####\n # MORE LSTM LAYERS? #\n #### ##### ##### ####\n \n # compute prediction logits and leverage dropout\n output = self.dropout_layer(hidden_new, training=training)\n output = self.fcn_layer(output, training=training)\n\n return output, hidden_new, memory_new\n \n def init_hidden_state(self, img_emb, training):\n mean_encoder_out = tf.math.reduce_mean(img_emb, axis=1)\n hidden = self.init_h_layer(mean_encoder_out, training=training) # (batch_size, n_rnn_units)\n memory = self.init_c_layer(mean_encoder_out, training=training)\n return hidden, memory\n\n\nwith tf.device('/CPU:0'):\n decoder = Decoder(VOCAB_LEN, ATTN_EMB_DIM, IMG_EMB_DIM[0], N_RNN_UNITS)\n h, c = decoder.init_hidden_state(img_embedding_batch[:BATCH_SIZE_DEBUG], training=False)\n pred_output, h, c = decoder(tf.random.uniform((BATCH_SIZE_DEBUG, 1)), h, c, img_embedding_batch)\n\nprint(f'\\n... Decoder Output Shape : (batch_size, vocab_len) : {pred_output.shape} ...\\n')\n\nprint(\"\\n... DECODER MODEL CREATION FINISHED ...\\n\")","metadata":{"id":"BZyzEM0FUekx","executionInfo":{"status":"ok","timestamp":1621257569926,"user_tz":-480,"elapsed":55552,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"910c4187-440e-4c8f-fd4e-ef022c2944ab","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n - TF Tutorial β Transformer Model for Language Understanding
\n - TF Tutorial β Image Captioning
\n - TF Tutorial β Neural Machine Translation w/ Attention
\n
5.3 CREATE A LEARNING RATE SCHEDULER
\n\n---\n\nAdapting the learning rate for your stochastic gradient descent optimization procedure can increase performance and reduce training time. Sometimes this is called learning rate annealing or adaptive learning rates. Here we will call this approach a learning rate schedule. See [**this article**](https://machinelearningmastery.com/using-learning-rate-schedules-deep-learning-models-python-keras/) for a basic tutorial on learning rade schedules.\n\n\nWe will utilize a basic step function following a warmup phase. Warmup is commonly used in learning rate schedule where we start training a model with a much smaller learning rate and increase it during the first few epochs/steps until the initial learning rate is used.\n\nIntuitively, this method will allow a model to adjust itself less before it becomes more familiar with the dataset. This usually prevents breaking pretrained weights. For adaptive optimisers like Adam, warmup also allows the optimizers to compute bettere statistics of the gradients.","metadata":{"id":"1RpQcq0oUeky"}},{"cell_type":"code","source":"print(\"\\n... LEARNING RATE SCHEDULE CREATION STARTING ...\\n\")\n\n# Part of the Training Configuration\nEPOCHS = 36\nTOTAL_STEPS = TRAIN_STEPS*EPOCHS\n\n# Learning Rate Scheduler Configuration\nWARM_STEPS = 100\nWARM_START_LR = 1e-5\nPEAK_START_LR = 2e-3\nFINAL_LR = 1e-4\n\ndef 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\ndef 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 \nclass 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\nlr_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]\nplot_lr_schedule(lr_schedule)\n\nprint(\"\\n... LEARNING RATE SCHEDULE CREATION FINISHED ...\\n\")","metadata":{"id":"FaGpq5KLUeky","executionInfo":{"status":"ok","timestamp":1621257760600,"user_tz":-480,"elapsed":2721,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"30d75d9e-fdd3-491a-8860-09a9e7ee5dfa","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":{"id":"DYWxCU9-Uekz"}},{"cell_type":"code","source":"class Config():\n def __init__(self,):\n self.encoder_config = {}\n self.decoder_config = {}\n self.lr_config = {}\n def initialize_encoder_config(self, image_embedding_dim, preprocessing_fn, backbone_fn, image_shape, do_permute=False, pretrained_weights=None):\n self.encoder_config = dict(\n image_embedding_dim=image_embedding_dim, \n preprocessing_fn=preprocessing_fn, \n backbone_fn=backbone_fn, \n image_shape=image_shape, \n do_permute=do_permute, \n pretrained_weights=pretrained_weights,\n )\n def initialize_decoder_config(self, vocab_len, attn_emb_depth, img_emb_dim, n_rnn_units, dropout_rate=0.05):\n self.decoder_config = dict(\n vocab_len=vocab_len, \n attn_emb_depth=attn_emb_depth, \n img_emb_dim=img_emb_dim, \n n_rnn_units=n_rnn_units, \n dropout_rate=dropout_rate,\n )\n def initialize_lr_config(self, total_steps, warm_lr_start, warm_steps, peak_lr_start, lr_final, n_epochs):\n self.lr_config = dict(\n total_steps=total_steps, \n warm_lr_start=warm_lr_start, \n warm_steps=warm_steps, \n peak_lr_start=peak_lr_start, \n lr_final=lr_final, \n n_epochs=n_epochs,\n )\n \ntraining_config = Config()\ntraining_config.initialize_encoder_config(image_embedding_dim=IMG_EMB_DIM, \n preprocessing_fn=PREPROCESSING_FN, \n backbone_fn=BB_FN, \n image_shape=IMG_SHAPE, \n do_permute=IMG_EMB_DIM[1]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":{"id":"TGbuWYKNUekz"}},{"cell_type":"code","source":"print(\"\\n... TRAINING PREPERATION STARTING ...\\n\")\n\ndef prepare_for_training(lr_config, encoder_config, decoder_config, encoder_wts=None, decoder_wts=None, verbose=0):\n \"\"\" Declare required objects under TPU session scope and return ready for training\n \n Args:\n 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 decoder_config (dict): Keyword arguments mapped to desired values for decoder model instantiation \n encoder_wts (str, optional): Path to pretrained model weights for encoder\n decoder_wts (str, optional): Path to pretrained model weights for decoder\n verbose (bool, optional): Whether or not to print model information and plot lr schedule\n \n Returns:\n loss_fn - TBD\n metrics - TBD\n optimizer - TBD\n lr_scheduler - TBD\n encoder - TBD\n decoder - TBD\n \n \"\"\"\n \n\n # Everything must be declared within the scope when leveraging the TPU strategy\n # - This will still function properly if scope is set to another type of accelerator\n with strategy.scope():\n \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 print(\"\\t--> CREATING OPTIMIZER ...\")\n # Instiate an optimizer\n optimizer = tf.keras.optimizers.Adam()\n \n print(\"\\t--> CREATING LEARNING RATE SCHEDULER ...\")\n # Declare the learning rate schedule (try this as actual lr schedule and list...)\n lr_schedule = [\n lr_schedule_fn(step=step, **lr_config) \\\n for step in range(lr_config[\"total_steps\"])\n ]\n lr_scheduler = LRS(optimizer, lr_schedule)\n \n # Instantiate the encoder model \n print(\"\\t--> CREATING ENCODER MODEL ARCHITECTURE ...\")\n encoder = Encoder(**encoder_config)\n initialization_batch = encoder(\n tf.ones(((REPLICA_BATCH_SIZE,)+encoder_config[\"image_shape\"]), dtype=TARGET_DTYPE), \n training=False,\n )\n \n # Instantiate the decoder model\n print(\"\\t--> CREATING DECODER MODEL ARCHITECTURE...\")\n decoder = Decoder(**decoder_config)\n init_h, init_c = decoder.init_hidden_state(initialization_batch[:REPLICA_BATCH_SIZE], training=False)\n pred_output, h, c = decoder(\n tf.random.uniform((REPLICA_BATCH_SIZE, 1)), \n init_h, init_c, initialization_batch, \n training=False,\n )\n \n # Load weights after variable initialization\n if encoder_wts is not None:\n print(\"\\t--> LOADING ENCODER MODEL WEIGHTS ...\")\n encoder.load_weights(encoder_wts)\n if decoder_wts is not None:\n print(\"\\t--> LOADING DECODER MODEL WEIGHTS ...\")\n decoder.load_weights(decoder_wts)\n \n # Show the model architectures and plot the learning rate\n if verbose:\n print(\"\\n\\n... ENCODER MODEL SUMMARY...\\n\")\n print(encoder.summary())\n\n print(\"\\n\\n... DECODER MODEL SUMMARY...\\n\")\n print(decoder.summary())\n\n print(\"\\n\\n... LR SCHEDULE PLOT...\\n\")\n plot_lr_schedule(lr_schedule)\n \n return loss_fn, metrics, optimizer, lr_scheduler, encoder, decoder\n \n \nprint(\"\\n... GENERATING THE FOLLOWING:\")\n# Instantiate our required training components in the correct scope\nloss_fn, metrics, optimizer, lr_scheduler, encoder, decoder = \\\n prepare_for_training(lr_config=training_config.lr_config,\n encoder_config=training_config.encoder_config,\n decoder_config=training_config.decoder_config,\n encoder_wts=(ENCODER_CKPT_PATH if ENCODER_CKPT_PATH!=\"\" else None),\n decoder_wts=(DECODER_CKPT_PATH if DECODER_CKPT_PATH!=\"\" else None),\n verbose=0,)\n\nprint(\"\\n... TRAINING PREPERATION FINISHED ...\\n\")","metadata":{"id":"bbTV0zMVUek0","executionInfo":{"status":"ok","timestamp":1621257791359,"user_tz":-480,"elapsed":33461,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"ee4aba64-3dd4-4076-f832-b33f71994b16","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":{"id":"g8D_SWZ2Uek1"}},{"cell_type":"code","source":"print(\"-\"*100) # Just for printing\n\naccumulated_loss = np.zeros((BATCH_SIZE_DEBUG,))\nfor i in range(0, MAX_LEN, 1):\n loss = loss_fn(SAMPLE_LBLS[:, i], pred_output).numpy()\n accumulated_loss+=loss\n \n if i%50==0:\n print(f\"\\n\\nLABEL FOR STEP {i+1} OF {MAX_LEN} FOR THE {BATCH_SIZE_DEBUG} EXAMPLES IN OUR DEMO BATCH:\\n\\t--> {SAMPLE_LBLS[:, i]}\")\n print(f\"\\nOUR PREDICTIONS:\\n\\tSHAPE: {pred_output.shape}\\n\\tARGMAX: {tf.argmax(pred_output, axis=1)}\")\n print(f\"\\nTHE CALCULATED INDIVIDUAL STEP LOSS:\\n\\t--> {loss_fn(SAMPLE_LBLS[:, i], pred_output).numpy()}\\n\\n\")\n print(\"-\"*100) # Just for printing\n \nprint(f\"\\n\\n... ACCUMULATED LOSS:\\n\\t--> {accumulated_loss}\")\nprint(f\"\\n\\n... AVERAGE LOSS BY EXAMPLE:\\n\\t--> {accumulated_loss/MAX_LEN}\")\nprint(f\"\\n\\n... AVERAGE LOSS ACROSS BATCH:\\n\\t--> {np.sum(accumulated_loss)/(MAX_LEN*BATCH_SIZE_DEBUG)}\\n\\n\")\n\nprint(\"-\"*100) # Just for printing","metadata":{"id":"i6JTTuBRUek1","executionInfo":{"status":"ok","timestamp":1621257792281,"user_tz":-480,"elapsed":34375,"user":{"displayName":"dragon zhang","photoUrl":"","userId":"17320165301386381342"}},"outputId":"1194db30-adb6-48bd-8f20-16fea5c095f4","trusted":true},"execution_count":null,"outputs":[]},{"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":{"id":"bNjsPcluUek2"}},{"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":{"id":"8oc8c0jcUek2"}},{"cell_type":"markdown","source":"\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
\n\n\n\n\n\n
6 MODEL TRAINING ⤒
\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":{"id":"8d0xEFA4Uek3"}},{"cell_type":"markdown","source":"6.1 INDIVIDUAL TRAIN STEP
\n\n---\n\nINFORMATION","metadata":{"id":"QhW9xK9nUek3"}},{"cell_type":"code","source":"def train_step(_image_batch, _inchi_batch):\n \"\"\" Forward pass (calculate gradients)\n \n Args:\n _image_batch (): TBD\n _inchi_batch (): TBD\n \n Returns:\n tbd\n \"\"\"\n \n batch_loss = tf.constant(0.0, tf.float32) \n with tf.GradientTape() as tape:\n # image_batch_embedding has shape --> (REPLICA_BATCH_SIZE, IMG_EMB_DIM)\n image_batch_embedding = encoder(_image_batch, training=True)\n \n # hidden and memory both have the shape --> (REPLICA_BATCH_SIZE, N_RNN_UNITS)\n hidden_batch, memory_batch = decoder.init_hidden_state(image_batch_embedding, training=True)\n \n # decoder_input has a shape --> (REPLICA_BATCH_SIZE, 1)\n decoder_input_batch = tf.ones((REPLICA_BATCH_SIZE, 1), dtype=tf.uint8)\n \n # Teacher forcing - feeding the target as the next input\n for c_idx in range(1, MAX_LEN):\n gt_batch = _inchi_batch[:, c_idx]\n \n # passing enc_output to the decoder\n prediction_batch, hidden_batch, memory_batch = \\\n decoder(decoder_input_batch, hidden_batch, memory_batch, image_batch_embedding, training=True)\n \n # Update Loss Accumulator\n batch_loss += loss_fn(gt_batch, prediction_batch)\n \n # Update Accuracy Metric\n metrics[\"train_acc\"].update_state(gt_batch, prediction_batch, \n sample_weight=tf.where(tf.not_equal(gt_batch, PAD_TOKEN), 1.0, 0.0))\n \n # teacher forcing, use correct character as next input to LSTMCell\n decoder_input_batch = tf.expand_dims(gt_batch, 1)\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 + decoder.trainable_variables)\n gradients, _ = tf.clip_by_global_norm(gradients, 10.0)\n \n # Normalize loss across all characters \n batch_loss = batch_loss/(MAX_LEN-1)\n \n metrics[\"batch_loss\"].update_state(batch_loss)\n metrics[\"train_loss\"].update_state(batch_loss)\n \n optimizer.apply_gradients(zip(gradients, encoder.trainable_variables+decoder.trainable_variables))\n\n@tf.function\ndef dist_train_step(_image_batch, _inchi_batch):\n strategy.run(train_step, args=(_image_batch, _inchi_batch))","metadata":{"id":"NnAf2CAjUek4","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"6.2 INDIVIDUAL VAL STEP
\n\n---\n\nINFORMATION","metadata":{"id":"_KSt9qngUek4"}},{"cell_type":"code","source":"def val_step(_image_batch, _inchi_batch):\n \"\"\" Forward pass (calculate gradients)\n \n Args:\n image_batch (): TBD\n inchi_batch (): TBD\n \n Returns:\n tbd\n \"\"\"\n \n # Initialize batch_loss\n batch_loss = tf.constant(0.0, tf.float32) \n \n # image_batch_embedding has shape --> (REPLICA_BATCH_SIZE, IMG_EMB_DIM)\n image_batch_embedding = encoder(_image_batch, training=False)\n\n # hidden and memory both have the shape --> (REPLICA_BATCH_SIZE, N_RNN_UNITS)\n hidden_batch, memory_batch = decoder.init_hidden_state(image_batch_embedding, training=False)\n\n # decoder_input and predictions_seq share the shape --> (REPLICA_BATCH_SIZE, 1)\n decoder_input_batch = tf.ones((REPLICA_BATCH_SIZE, 1), dtype=tf.uint8)\n predictions_seq_batch = tf.ones((REPLICA_BATCH_SIZE, 1), dtype=tf.uint8)\n\n # Teacher forcing - feeding the target as the next input\n for c_idx in range(1, MAX_LEN):\n gt_batch = _inchi_batch[:, c_idx]\n \n # passing enc_output to the decoder\n prediction_batch, hidden_batch, memory_batch = \\\n decoder(decoder_input_batch, hidden_batch, memory_batch, image_batch_embedding, training=False)\n\n # Update Loss Accumulator\n batch_loss += loss_fn(gt_batch, prediction_batch)\n \n # Update Accuracy Metric\n metrics[\"val_acc\"].update_state(gt_batch, prediction_batch,\n sample_weight=tf.where(tf.not_equal(gt_batch, PAD_TOKEN), 1.0, 0.0))\n\n # no teacher forcing, predicted char is next LSTMCell input\n decoder_input_batch = tf.expand_dims(tf.cast(tf.math.argmax(prediction_batch, axis=1, output_type=tf.int32), tf.uint8), axis=1)\n predictions_seq_batch = tf.concat([predictions_seq_batch, decoder_input_batch], axis=1)\n \n # Normalize loss across all characters \n batch_loss = batch_loss/(MAX_LEN-1)\n\n # Update Levenshtein Distance Metric & Loss Metric\n metrics[\"val_loss\"].update_state(batch_loss)\n \n return predictions_seq_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":{"id":"CMeY-zesUek5","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"6.3 INITIALIZE LOGGER
\n\n---\n\nINFORMATION","metadata":{"id":"h-A7Zew0Uek7"}},{"cell_type":"code","source":"class StatLogger():\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":{"id":"olB5r_AVUek7","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"6.4 CUSTOM TRAIN LOOP
\n\n---\n\nINFORMATION","metadata":{"id":"RG-ODl4DUek8"}},{"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:\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.lr)\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.lr\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 encoder.save_weights(f'./encoder_epoch_{epoch}.h5')\n decoder.save_weights(f'./decoder_epoch_{epoch}.h5')","metadata":{"id":"c18utTJUUek8","outputId":"0e3c9209-7c0b-4ad9-cb9f-bdf6efb61184","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"6.5 JUST-IN-CASE SAVE
\n\n---\n\nINFORMATION","metadata":{"id":"7PYZD70bUek9"}},{"cell_type":"code","source":"encoder.save_weights(f'./encoder_epoch_safety_save.h5')\ndecoder.save_weights(f'./decoder_epoch_safety_save.h5')","metadata":{"id":"TSDFNnjCUek9","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"6.6 VIEW PREDICTIONS & DISTRIBUTION OF LEVENSHTEIN DISTANCE FOR VAL DATASET
\n\n---\n\nINFORMATION","metadata":{"id":"szrOMXESUek-"}},{"cell_type":"code","source":"pred_s = []\nlbls_s = []\nLDS = []\n\nfor _ in tqdm(range(VAL_STEPS), total=VAL_STEPS):\n preds, lbls = dist_val_step(val_dist_ds)\n for i, (p, l) in enumerate(zip(preds.numpy(), lbls.numpy())):\n try:\n p_s = \"\".join([int_2_tok[x] for x in p[:np.where(p==2)[0][0]]])\n except:\n p_s = \"\".join([int_2_tok[x] for x in p])\n l_s = \"\".join([int_2_tok[x] for x in l[:np.where(l==2)[0][0]]])\n LD = Levenshtein.distance(p_s,l_s)\n \n if i % 1000==0:\n print(f\"Levenshtein Distance: {LD}\")\n print(f\"\\t--> {p_s}\")\n print(f\"\\t--> {l_s}\\n\")\n\n pred_s.append(p_s)\n lbls_s.append(l_s)\n LDS.append(LD)\n \nld_df = pd.DataFrame({\"lbl_inchi\":lbls_s, \"pred_inchi\":pred_s, \"lev_dist\":LDS})\nld_df.to_csv(\"./levenshtein_on_val.csv\")\n\nprint(ld_df.describe())\ndisplay(ld_df)\n\npx.histogram(ld_df.lev_dist, log_y=True)","metadata":{"id":"K63PrhHPUek-","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n\n\n\n\n\n