{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"**Hello fellow Kagglers,**\n\n\nThis notebook is a baseline for an encoder/decoder model with attention written in Tensorflow and running on a TPU. Several notebooks, examples and documentation were used as a source of inspiration, especially the two Kaggle notebooks, a big thanks for sharing that work:\n\n**Kaggle Notebook**\n\n[Pytorch training by Eric Pasewark](https://www.kaggle.com/yasufuminakama/inchi-resnet-lstm-with-attention-starter)\n\n[Pytorch training by Y.Nakama](https://www.kaggle.com/yasufuminakama/inchi-resnet-lstm-with-attention-starter)\n\n**Tensorflow Code Examples/Documentation**\n\n[Tensorflow encoder/decoder attention baseline](https://www.tensorflow.org/tutorials/text/nmt_with_attention)\n\n[Custom Tensorflow model](https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit)\n\n[TPU training in Tensorflow](https://www.tensorflow.org/tutorials/distribute/custom_training)\n\n**My own preprocessing notebook**\n\n[Advanced Image Cleaning and TFRecord Generation](https://www.kaggle.com/markwijkhuizen/advanced-image-cleaning-and-tfrecord-generation)\n\n**Prediction Notebook (available several hours after V3 completes running)**\n\n[BMS - Tensorflow TPU Predictions](https://www.kaggle.com/markwijkhuizen/bms-tensorflow-tpu-predictions)\n\nI will not disclose the prediction notebook to prevent people from simply copying and submitting this notebook and thereby flooding the leaderboard with equal scores.\n\nIf you have any questions or remarks, feel free to leave a comment :D\n\nWhen publishing a notebook based on this notebook, please don't forget to reference this notebook.\n\nA small disclaimer, this is the first time I am playing around with sequence predictions and encoder/decoder models. Keep this in mind when reading the notebook, many improvements will be possible.\n\n**VERSION 2 UPDATES**\n\n* Dataset converted to iterator. Without iterator the dataset starts at the beginning each epoch, thereby using only the first part of the train dataset. Credits go to [Darien Shettler](https://www.kaggle.com/dschettler8845) for pointing this out in the comments.\n\n* Dynamically assign encoder dimensions. This idea is based on [Andy Penrose's](https://www.kaggle.com/andypenrose) comment\n\n* Optimized training loop, this idea is based on [this](https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu) training notebook made by [Martin Görner](https://www.kaggle.com/mgornergoogle). An example of this in the Tensorflow documentation can be found [here](https://www.tensorflow.org/guide/tpu#improving_performance_by_multiple_steps_within_tffunction). Multiple training steps are performed in one run on the TPU, 100 to be precise. Also, the batch of images and labels are retrieved directly on the TPU, rather than on the CPU to be then send to the TPU. This reduces the training step duration from 45 second to 38 seconds, a reduction of 16\\% :D.\n\n**VERSION 3 UPDATES**\n\n* Updates the attention mechanism based on [this](https://www.kaggle.com/konradb/model-train-efficientnet) notebook. This improves both the score and efficiency, and epoch now takes only 27 seconds, TPU's are awesome ;)\n\n* Modified learning rate scheduler, using lower learning rates.\n\n* Reduced the character embedding dimension.\n\n* Made [prediction notebook](https://www.kaggle.com/markwijkhuizen/bms-tensorflow-tpu-predictions) public, will be finished after V3 has finished.","metadata":{}},{"cell_type":"code","source":"# install tensorflow implementations of EfficientNet with noisy-student weights\n!pip install -q --upgrade pip\n!pip install -q efficientnet","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport tensorflow as tf\nimport matplotlib.pyplot as plt\nimport efficientnet.tfkeras as efn\n\nfrom tensorflow.keras.mixed_precision import experimental as mixed_precision\nfrom kaggle_datasets import KaggleDatasets\nfrom tqdm.notebook import tqdm\n\nimport unicodedata\nimport re\nimport numpy as np\nimport os\nimport io\nimport time\nimport pickle\nimport math\nimport random","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# seed everything\nSEED = 42\nos.environ['PYTHONHASHSEED'] = str(SEED)\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntf.random.set_seed(SEED)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Detect hardware, set appropriate distribution strategy (GPU/TPU)\ntry:\n    TPU = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection. No parameters necessary if TPU_NAME environment variable is set. On Kaggle this is always the case.\n    print('Running on TPU ', TPU.master())\nexcept ValueError:\n    print('Running on GPU')\n    TPU = None\n\nif TPU:\n    tf.config.experimental_connect_to_cluster(TPU)\n    tf.tpu.experimental.initialize_tpu_system(TPU)\n    strategy = tf.distribute.experimental.TPUStrategy(TPU)\nelse:\n    strategy = tf.distribute.get_strategy() # default distribution strategy in Tensorflow. Works on CPU and single GPU.\n\nREPLICAS = strategy.num_replicas_in_sync\nprint(f'REPLICAS: {REPLICAS}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEBUG = False\n\n# image resolution\nIMG_HEIGHT = 256\nIMG_WIDTH = 448\nN_CHANNELS = 3\n# maximum InChI length is 200 to prevent too much padding\nMAX_INCHI_LEN = 200\n\n# batch sizes\nBATCH_SIZE_BASE = 6 if DEBUG else (64 if TPU else 12)\nBATCH_SIZE = BATCH_SIZE_BASE * REPLICAS\nBATCH_SIZE_DEBUG = 2\n\n# target data type, bfloat16 when using TPU to improve throughput\nTARGET_DTYPE = tf.bfloat16 if TPU else tf.float32\n # minimal memory usage of labels\nLABEL_DTYPE= tf.uint8\n\n# 100K validation images are used\nVAL_SIZE = int(1e3) if DEBUG else int(100e3)\nVAL_STEPS = VAL_SIZE // BATCH_SIZE\n\n# ImageNet mean and std to normalize training images accordingly\nIMAGENET_MEAN = tf.constant([0.485, 0.456, 0.406], dtype=tf.float32)\nIMAGENET_STD = tf.constant([0.229, 0.224, 0.225], dtype=tf.float32)\n\n# Google Cloud Dataset path to training and validation images\nGCS_DS_PATH = KaggleDatasets().get_gcs_path('molecular-translation-images-cleaned-tfrecords')\n\n# Tensorflow AUTO flag, used in datasets\nAUTO = tf.data.experimental.AUTOTUNE","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# dictionary to translate a character to the integer encoding\nwith open('/kaggle/input/molecular-translation-images-cleaned-tfrecords/vocabulary_to_int.pkl', 'rb') as handle:\n    vocabulary_to_int = pickle.load( handle)\n\n# dictionary to decode an integer encoded character back to the character\nwith open('/kaggle/input/molecular-translation-images-cleaned-tfrecords/int_to_vocabulary.pkl', 'rb') as handle:\n    int_to_vocabulary = pickle.load( handle)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# configure model\nVOCAB_SIZE = len(vocabulary_to_int.values())\nSEQ_LEN_OUT = MAX_INCHI_LEN\nDECODER_DIM = 512\nCHAR_EMBEDDING_DIM = 256\nATTENTION_UNITS = 256\n\nprint(f'VOCAB_SIZE: {VOCAB_SIZE}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Datasets\n\nSince there are 2.4M training images, reading those images efficiently is of great importance to not get IO bottlenecked, which is a key problem with TPU's. With the TFRecord format the images are read in batches and not one by one. As will be shown, ~6000 images per second can be loaded. To further improve training speed images are converted to the bfloat16 data type. This is a 16 bits float with the range of a 32 bit float, but a lower precision than a 16 bits float. When training a neural network the 10th decimal is not interesting, however the full range of a 32 bits float is needed, therefore this data format fits the needs perfectly with just half memory usage of a conventional 32 bits float. More info on the awesome bfloat16 data type can be found [here](https://cloud.google.com/tpu/docs/bfloat16)","metadata":{}},{"cell_type":"code","source":"# decodes TFRecord\ndef decode_tfrecord(record_bytes):\n    features = tf.io.parse_single_example(record_bytes, {\n        'image': tf.io.FixedLenFeature([], tf.string),\n        'InChI': tf.io.FixedLenFeature([MAX_INCHI_LEN], tf.int64),\n    })\n\n    # decode the PNG and explicitly reshape to image size (required on TPU)\n    image = tf.io.decode_png(features['image'])    \n    image = tf.reshape(image, [IMG_HEIGHT, IMG_WIDTH, 1])\n    # normalize according to ImageNet mean and std\n    image = tf.cast(image, tf.float32)  / 255.0\n    image = (image - IMAGENET_MEAN) / IMAGENET_STD\n    \n    if TPU: # if running on TPU image needs to be cast to bfloat16\n        image = tf.cast(image, TARGET_DTYPE)\n    \n    InChI = tf.reshape(features['InChI'], [MAX_INCHI_LEN])\n    InChI = tf.cast(InChI, LABEL_DTYPE)\n    \n    return image, InChI","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Benchmark function to test the dataset throughput performance\ndef benchmark_dataset(dataset, num_epochs=3, n_steps_per_epoch=25, bs=BATCH_SIZE):\n    start_time = time.perf_counter()\n    for epoch_num in range(num_epochs):\n        epoch_start = time.perf_counter()\n        for idx, (images, labels) in enumerate(dataset.take(n_steps_per_epoch)):\n            if idx is 1 and epoch_num is 0:\n                print(f'image shape: {images.shape}, image dtype: {images.dtype}')\n                print(f'labels shape: {labels.shape}, label dtype: {labels.dtype}')\n            pass\n        epoch_t = time.perf_counter() - epoch_start\n        mean_step_t = round(epoch_t / n_steps_per_epoch * 1000, 1)\n        n_imgs_per_s = int(1 / (mean_step_t / 1000) * bs)\n        print(f'epoch {epoch_num} took: {round(epoch_t, 2)} sec, mean step duration: {mean_step_t}ms, images/s: {n_imgs_per_s}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# plots the first images of the dataset\ndef show_batch(dataset, rows=3, cols=2):\n    imgs, lbls = next(iter(dataset))\n    fig, axes = plt.subplots(nrows=rows, ncols=cols, figsize=(cols*7, rows*4))\n    for r in range(rows):\n        for c in range(cols):\n            img = imgs[r*cols+c].numpy().astype(np.float32)\n            img += abs(img.min())\n            img /= img.max()\n            axes[r, c].imshow(img)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Dataset","metadata":{}},{"cell_type":"code","source":"def get_train_dataset(bs=BATCH_SIZE):\n    ignore_order = tf.data.Options()\n    ignore_order.experimental_deterministic = False\n    \n    FNAMES_TRAIN_TFRECORDS = tf.io.gfile.glob(f'{GCS_DS_PATH}/train/*.tfrecords')\n    train_dataset = tf.data.TFRecordDataset(FNAMES_TRAIN_TFRECORDS, num_parallel_reads=AUTO)\n    train_dataset = train_dataset.with_options(ignore_order)\n    train_dataset = train_dataset.prefetch(AUTO) # optimize automatically\n    train_dataset = train_dataset.repeat()\n    train_dataset = train_dataset.map(decode_tfrecord, num_parallel_calls=AUTO)  # optimize automatically\n    train_dataset = train_dataset.batch(BATCH_SIZE, drop_remainder=True)\n    train_dataset = train_dataset.prefetch(1) # just 1 prefetched batch is needed\n    \n    return train_dataset\n\ntrain_dataset = get_train_dataset()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"benchmark_dataset(train_dataset)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# display statistics about the first image to check if the images are decoded correctly\nimgs, lbls = next(iter(train_dataset))\nprint(f'imgs.shape: {imgs.shape}, lbls.shape: {lbls.shape}')\nimg0 = imgs[0].numpy().astype(np.float32)\ntrain_batch_info = (img0.mean(), img0.std(), img0.min(), img0.max(), imgs.dtype)\nprint('train img 0 mean: %.3f, 0 std: %.3f, min: %.3f, max: %.3f, %s' % train_batch_info)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# show first few train images\nshow_batch(train_dataset)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Validation Dataset","metadata":{}},{"cell_type":"code","source":"def get_val_dataset(bs=BATCH_SIZE):\n    FNAMES_TRAIN_TFRECORDS = tf.io.gfile.glob(f'{GCS_DS_PATH}/val/*.tfrecords')\n    val_dataset = tf.data.TFRecordDataset(FNAMES_TRAIN_TFRECORDS, num_parallel_reads=AUTO)\n    val_dataset = val_dataset.prefetch(AUTO)\n    val_dataset = val_dataset.repeat()\n    val_dataset = val_dataset.map(decode_tfrecord, num_parallel_calls=AUTO)\n    val_dataset = val_dataset.batch(BATCH_SIZE, drop_remainder=True)\n    val_dataset = val_dataset.prefetch(1)\n    \n    return val_dataset\n\nval_dataset = get_val_dataset()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"benchmark_dataset(val_dataset)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"val_imgs, val_lbls = next(iter(val_dataset))\nprint(f'val_imgs.shape: {val_imgs.shape}, val_lbls.shape: {val_lbls.shape}')\nval_img0 = val_imgs[0].numpy().astype(np.float32)\nval_batch_info = (val_img0.mean(), val_img0.std(), val_img0.min(), val_img0.max(), val_imgs.dtype)\nprint('val img 0 mean: %.3f, 0 std: %.3f, min: %.3f, max: %.3f, %s' % train_batch_info)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_batch(val_dataset)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Encoder\nAn encoder/decoder model with attention is used, which is based on [this](https://www.tensorflow.org/tutorials/text/nmt_with_attention) Tensorflow example.\n\nThe encoder creates the feature maps of the images, which are then used in the encoder. EfficientNetB0 with pretrained noisy-student weights creates 1280 feature maps with dimensions of $14\\cdot8$ pixels. These feature maps are flattened by a reshape: $14\\cdot8\\cdot1280 \\Rightarrow 112\\cdot1280$.","metadata":{}},{"cell_type":"code","source":"class Encoder(tf.keras.Model):\n    def __init__(self):\n        super(Encoder, self).__init__()\n        \n        # output: (bs, 1280, 14, 8)\n        self.feature_maps = efn.EfficientNetB0(include_top=False, weights='noisy-student')\n        # set global encoder dimension variable\n        global ENCODER_DIM\n        ENCODER_DIM = self.feature_maps.layers[-1].output_shape[-1]\n        \n        # output: (bs, 1280, 112)\n        self.reshape = tf.keras.layers.Reshape([-1, ENCODER_DIM], name='reshape_featuere_maps')\n\n    def call(self, x, training, debug=False):\n        x = self.feature_maps(x, training=training)\n        if debug:\n            print(f'feature maps shape: {x.shape}')\n            \n        x = self.reshape(x, training=training)\n        if debug:\n            print(f'feature maps reshaped shape: {x.shape}')\n        \n        return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Example enoder output\nwith tf.device('/CPU:0'):\n    encoder = Encoder()\n    encoder_res = encoder(imgs[:BATCH_SIZE_DEBUG], debug=True)\n\nprint ('Encoder output shape: (batch size, sequence length, units) {}'.format(encoder_res.shape))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Attention\nDuring the decoding phase the important features from the encoder will differ for each character predicted. The attention mechanism takes as input the hidden state from the LSTM, which is the LSTM state after the last predicted character, and encoder features. The hidden LSTM state will differ each prediction iteration, but the encoder result remains the same. Using this hidden LSTM state the attention mechanism learns which parts of the feature maps are important. The feature maps have a dimension of 8*14 pixels whicha re flattened to a vector of size 112. The attention mechanism creates a importancy score for each pixel, which is a probability distribution summing to 1, over the 112 pixels and multiplies it with the feature map vectors to create a single value for each feature map.\n\nTo make it a bit less abstract, take the next InChI as an example\n\n```C13H5F5N2/c14-7-3-6(5-19)1-2-10(7)20-13-11(17)8(15)4-9(16)12(13)18/h1-4,20H```\n\nAfter predicting C13H5 the attention mechanism should focus on features containing F atoms and leave any feature maps on C or H atoms aside. The LSTM hidden state should tell the attention mechanism it has predicted C13H5 so far and the attention mechanism will learn it has to focus on F atoms after C and H atoms are predicted.","metadata":{}},{"cell_type":"code","source":"class BahdanauAttention(tf.keras.layers.Layer):\n    def __init__(self, units):\n        super(BahdanauAttention, self).__init__()\n        self.H = tf.keras.layers.Dense(units, name='hidden_to_attention_units')\n        self.E = tf.keras.layers.Dense(units, name='encoder_res_to_attention_units')\n        self.V = tf.keras.layers.Dense(1, name='score_to_alpha')\n\n    def call(self, h, encoder_res, training, debug=False):\n        # dense hidden state to attention units size and expand dimension\n        h_expand = tf.expand_dims(h, axis=1) # expand dimension\n        if debug:\n            print(f'h shape: {h.shape}, encoder_res shape: {encoder_res.shape}')\n            print(f'h_expand shape: {h_expand.shape}')\n            \n        h_dense = self.H(h_expand, training=training)\n        \n        # dense features to units size\n        encoder_res_dense = self.E(encoder_res, training=training) # dense to attention\n\n        # add vectors\n        score = tf.nn.relu(h_dense + encoder_res_dense)\n        if debug:\n            print(f'h_dense shape: {h_dense.shape}')\n            print(f'encoder_res_dense shape: {encoder_res_dense.shape}')\n            print(f'score tanh shape: {score.shape}')\n        score = self.V(score, training=training)\n        \n        # create alpha vector size (bs, layers)        \n        attention_weights = tf.nn.softmax(score, axis=1)\n        if debug:\n            score_np = score.numpy().astype(np.float32)\n            print(f'score V shape: {score.shape}, score min: %.3f score max: %.3f' % (score_np.min(), score_np.max()))\n            print(f'attention_weights shape: {attention_weights.shape}')\n            aw = attention_weights.numpy().astype(np.float32)\n            aw_print_data = (aw.min(), aw.max(), aw.mean(), aw.sum())\n            print(f'aw shape: {aw.shape} aw min: %.3f, aw max: %.3f, aw mean: %.3f,aw sum: %.3f' % aw_print_data)\n        \n        # create attention weights (bs, layers)\n        context_vector = encoder_res * attention_weights\n        if debug:\n            print(f'first attention weights: {attention_weights.numpy().astype(np.float32)[0,0]}')\n            print(f'first encoder_res: {encoder_res.numpy().astype(np.float32)[0,0,0]}')\n            print(f'first context_vector: {context_vector.numpy().astype(np.float32)[0,0,0]}')\n            \n            print(f'42th attention weights: {attention_weights.numpy().astype(np.float32)[0,42]}')\n            print(f'42th encoder_res: {encoder_res.numpy().astype(np.float32)[0,42,42]}')\n            print(f'42th context_vector: {context_vector.numpy().astype(np.float32)[0,42,42]}')\n            \n            print(f'encoder_res abs sum: {abs(encoder_res.numpy().astype(np.float32)).sum()}')\n            print(f'context_vector abs sum: {abs(context_vector.numpy().astype(np.float32)).sum()}')\n            \n            print(f'encoder_res shape: {encoder_res.shape}, attention_weights shape: {attention_weights.shape}')\n            print(f'context_vector shape: {context_vector.shape}')\n        \n        # reduce to ENCODER_DIM features\n        context_vector = tf.reduce_sum(context_vector, axis=1)\n        \n        return context_vector","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with tf.device('/CPU:0'):\n    attention_layer = BahdanauAttention(ATTENTION_UNITS)\n    context_vector, attention_weights = attention_layer(tf.zeros([BATCH_SIZE_DEBUG, DECODER_DIM]), encoder_res, debug=True)\n\nprint('context_vector shape: (batch size, units) {}'.format(context_vector.shape))\nprint('attention_weights shape: (batch_size, sequence_length, 1) {}'.format(attention_weights.shape))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Decoder\nThe decoder takes the encoder features and predicts one character at a time using an LSTMCell. The LSTMCell takes a concatinated context from the attention mechanism and an embedded character as input. The LSTMCell hidden and carry states are initialized with the encoder features. A 30\\% dropout is used on the LSTMCell output before making the final prediction.","metadata":{}},{"cell_type":"code","source":"class Decoder(tf.keras.Model):\n    def __init__(self, vocab_size, attention_units, encoder_dim, decoder_dim, char_embedding_dim):\n        super(Decoder, self).__init__()\n        \n        # LSTM hidden and carry state initialization\n        self.init_h = tf.keras.layers.Dense(units=decoder_dim, input_shape=[encoder_dim], name='encoder_res_to_hidden_init')\n        self.init_c = tf.keras.layers.Dense(units=decoder_dim, input_shape=[encoder_dim], name='encoder_res_to_inp_act_init')\n        # The LSTM cell\n        self.lstm_cell = tf.keras.layers.LSTMCell(decoder_dim, name='lstm_char_predictor')\n        # dropout before prediction\n        self.do = tf.keras.layers.Dropout(0.30, name='prediction_dropout')\n        # fully connected prediction layer\n        self.fcn = tf.keras.layers.Dense(units=vocab_size, input_shape=[decoder_dim], dtype=tf.float32, name='lstm_output_to_char_probs')\n        # character embedding layer\n        self.embedding = tf.keras.layers.Embedding(vocab_size, char_embedding_dim, name='character_embedding')\n\n        # used for attention\n        self.attention = BahdanauAttention(attention_units)\n\n    def call(self, char, h, c, enc_output, training, debug=False):\n        if debug:\n            print(f'char shape: {char.shape}, h shape: {h.shape}, c shape: {c.shape}, enc_output shape: {enc_output.shape}')\n        # embed previous character\n        char = self.embedding(char, training=training)\n        char = tf.squeeze(char, axis=1)\n        if debug:\n            print(f'char embedded and squeezed shape: {char.shape}')\n        # get attention alpha and context vector\n        context = self.attention(h, enc_output, training=training)\n\n        # concat context and char to create lstm input\n        lstm_input = tf.concat((context, char), axis=-1)\n        if debug:\n            print(f'lstm_input shape: {lstm_input.shape}')\n        \n        # LSTM call, get new h, c\n        _, (h_new, c_new) = self.lstm_cell(lstm_input, (h, c), training=training)\n        \n        # compute predictions with dropout\n        output = self.do(h_new, training=training)\n        output = self.fcn(output, training=training)\n\n        return output, h_new, c_new\n    \n    def init_hidden_state(self, encoder_out, training):\n        mean_encoder_out = tf.math.reduce_mean(encoder_out, axis=1)\n        h = self.init_h(mean_encoder_out, training=training)  # (batch_size, decoder_dim)\n        c = self.init_c(mean_encoder_out, training=training)\n        return h, c","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with tf.device('/CPU:0'):\n    decoder = Decoder(VOCAB_SIZE, ATTENTION_UNITS, ENCODER_DIM, DECODER_DIM, CHAR_EMBEDDING_DIM)\n    h, c = decoder.init_hidden_state(encoder_res[:BATCH_SIZE_DEBUG], training=False)\n    preds, h, c = decoder(lbls[:BATCH_SIZE_DEBUG, :1], h, c, encoder_res, debug=True)\n\nprint ('Decoder output shape: (batch_size, vocab size) {}'.format(preds.shape))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"# The start/end/pad tokens will be removed from the string when computing the Levenshtein distance\nSTART_TOKEN = tf.constant(vocabulary_to_int.get('<start>'), dtype=tf.int64)\nEND_TOKEN = tf.constant(vocabulary_to_int.get('<end>'), dtype=tf.int64)\nPAD_TOKEN = tf.constant(vocabulary_to_int.get('<pad>'), dtype=tf.int64)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tf.keras.backend.clear_session()\n\n# initialize the model, a dummy call to the encoder and deocder is made to allow the summaries to be printed\nwith strategy.scope():\n    # # set half precision policy\n    mixed_precision.set_policy('mixed_bfloat16' if TPU else 'float32')\n\n    # enable XLA optmizations\n    tf.config.optimizer.set_jit(True)\n\n    print(f'Compute dtype: {mixed_precision.global_policy().compute_dtype}')\n    print(f'Variable dtype: {mixed_precision.global_policy().variable_dtype}')\n    \n    # Sparse categorical cross entropy loss is used\n    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)\n\n    def loss_function(real, pred):\n        per_example_loss = loss_object(real, pred)\n\n        return tf.nn.compute_average_loss(per_example_loss, global_batch_size=BATCH_SIZE)\n    \n    # Metrics\n    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()\n    train_loss = tf.keras.metrics.Sum()\n    val_loss = tf.keras.metrics.Sum()\n\n\n    # Encoder\n    encoder = Encoder()\n    encoder.build(input_shape=[BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, N_CHANNELS])\n    encoder_res = encoder(imgs[:2], training=False)\n    \n    # Decoder\n    decoder = Decoder(VOCAB_SIZE, ATTENTION_UNITS, ENCODER_DIM, DECODER_DIM, CHAR_EMBEDDING_DIM)\n    h, c = decoder.init_hidden_state(encoder_res, training=False)\n    preds, h, c = decoder(lbls[:2, :1], h, c, encoder_res, training=False)\n    \n    # Adam Optimizer\n    optimizer = tf.keras.optimizers.Adam()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"encoder.summary()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"decoder.summary()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Learning Rate Scheduler\nMany different learning rate schedulers have been tried, and the current configuration worked best. An exponential warmup with a staircase decay is used.","metadata":{}},{"cell_type":"code","source":"# Training configuration\nEPOCHS = 10\nWARMUP_STEPS = 500\nTRAIN_STEPS = 1000\nVERBOSE_FREQ = 100\nSTEPS_PER_EPOCH = TRAIN_STEPS // VERBOSE_FREQ\nTOTAL_STEPS = EPOCHS * TRAIN_STEPS","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def lrfn(step, WARMUP_LR_START, LR_START, LR_FINAL, DECAYS):\n    # exponential warmup\n    if step < WARMUP_STEPS:\n        warmup_factor = (step / WARMUP_STEPS) ** 2\n        lr = WARMUP_LR_START + (LR_START - WARMUP_LR_START) * warmup_factor\n    # staircase decay\n    else:\n        power = (step - WARMUP_STEPS) // ((TOTAL_STEPS - WARMUP_STEPS) / (DECAYS + 1))\n        decay_factor =  ((LR_START / LR_FINAL) ** (1 / DECAYS)) ** power\n        lr = LR_START / decay_factor\n\n    return round(lr, 8)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# plot the learning rate schedule\ndef plot_lr_schedule(lr_schedule, name):\n    plt.figure(figsize=(15,8))\n    plt.plot(lr_schedule)\n    schedule_info = f'start: {lr_schedule[0]}, max: {max(lr_schedule)}, final: {lr_schedule[-1]}'\n    plt.title(f'Step Learning Rate Schedule {name}, {schedule_info}', size=16)\n    plt.grid()\n    plt.show()\n\n# Learning rate for encoder\nLR_SCHEDULE = [lrfn(step, 1e-8, 2e-3, 1e-4 ,EPOCHS) for step in range(TOTAL_STEPS)]\nplot_lr_schedule(LR_SCHEDULE, 'Ecnoder')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Levenshtein distance","metadata":{}},{"cell_type":"code","source":"# converts a dense to a sparse tensor\n# sparse tensors are required to compute the Levenshtein distance\ndef dense_to_sparse(dense):\n    ones = tf.ones(dense.shape)\n    indices = tf.where(ones)\n    values = tf.gather_nd(dense, indices)\n    sparse = tf.SparseTensor(indices, values, dense.shape)\n    \n    return sparse\n\n# computes the levenshtein distance between the predictions and labels\ndef get_levenshtein_distance(preds, lbls):\n    preds = tf.cast(preds, tf.int64)\n\n    preds = tf.where(tf.not_equal(preds, START_TOKEN) & tf.not_equal(preds, END_TOKEN) & tf.not_equal(preds, PAD_TOKEN), preds, y=0)\n    \n    lbls = strategy.gather(lbls, axis=0)\n    lbls = tf.cast(lbls, tf.int64)\n    lbls = tf.where(tf.not_equal(lbls, START_TOKEN) & tf.not_equal(lbls, END_TOKEN) & tf.not_equal(lbls, PAD_TOKEN), lbls, y=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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training Step\nTeacher forcing when predicting the characters, thus each iteration the correct character is fed to the decoder and not the predicted one.","metadata":{}},{"cell_type":"code","source":"@tf.function()\ndef distributed_train_step(dataset):\n    # Step function\n    def train_step(inp, targ):\n        total_loss = 0.0\n\n        with tf.GradientTape() as tape:\n            enc_output = encoder(inp, training=True)\n            h, c = decoder.init_hidden_state(enc_output, training=True)\n            dec_input = tf.expand_dims(targ[:, 0], 1)\n\n            # Teacher forcing - feeding the target as the next input\n            for idx in range(1, SEQ_LEN_OUT):\n                t = targ[:, idx]\n                t = tf.reshape(t, [BATCH_SIZE_BASE])\n                # passing enc_output to the decoder\n                predictions, h, c = decoder(dec_input, h, c, enc_output, training=True)\n\n                # update loss and train metrics\n                total_loss += loss_function(t, predictions)\n                train_accuracy.update_state(t, predictions)\n                # using teacher forcing\n                dec_input = tf.expand_dims(t, 1)\n\n        variables = encoder.trainable_variables + decoder.trainable_variables\n        gradients = tape.gradient(total_loss, variables)\n        gradients, _ = tf.clip_by_global_norm(gradients, 10.0)\n        optimizer.apply_gradients(zip(gradients, variables))\n\n        batch_loss = total_loss / (SEQ_LEN_OUT - 1)\n        train_loss.update_state(batch_loss)\n    \n    # reset metrics\n    train_loss.reset_states()\n    train_accuracy.reset_states()\n    # perform VERBOSE_FREQ train steps\n    for _ in tf.range(tf.convert_to_tensor(VERBOSE_FREQ)):\n        strategy.run(train_step, args=next(dataset))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Validation step\nWhen predicting teacher forcing is not applied as this will also not be the case when performing the test predictions. The predicted character is thus fed to the decoder to predict the next character","metadata":{}},{"cell_type":"code","source":"def validation_step(inp, targ):\n    total_loss = 0.0\n    enc_output = encoder(inp, training=False)\n    h, c = decoder.init_hidden_state(enc_output, training=False)\n    dec_input = tf.expand_dims(targ[:, 0], 1)\n\n    predictions_seq = tf.expand_dims(targ[:, 0], 1)\n\n    # Teacher forcing - feeding the target as the next input\n    for t in range(1, SEQ_LEN_OUT):\n        # passing enc_output to the decoder\n        predictions, h, c = decoder(dec_input, h, c, enc_output, training=False)\n\n        # add loss \n        # update loss and train metrics\n        total_loss += loss_function(targ[:, t], predictions)\n        \n        # add predictions to pred_seq\n        dec_input = tf.math.argmax(predictions, axis=1, output_type=tf.int32)\n        dec_input = tf.expand_dims(dec_input, axis=1)\n        dec_input = tf.cast(dec_input, LABEL_DTYPE)\n        predictions_seq = tf.concat([predictions_seq, dec_input], axis=1)\n        \n    batch_loss = total_loss / (SEQ_LEN_OUT - 1)\n    val_loss.update_state(batch_loss)\n    \n    return predictions_seq","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@tf.function\ndef distributed_val_step(dataset):\n    inp_val, targ_val = next(dataset)\n    per_replica_predictions_seq = strategy.run(validation_step, args=(inp_val, targ_val))\n    predictions_seq = strategy.gather(per_replica_predictions_seq, axis=0)\n    \n    return predictions_seq, targ_val","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_val_metrics(val_dist_dataset):\n    # reset metrics\n    val_loss.reset_states()\n    total_ls_distance = 0.0\n    \n    for step in range(VAL_STEPS):\n        predictions_seq, targ = distributed_val_step(val_dist_dataset)\n        levenshtein_distance = get_levenshtein_distance(predictions_seq, targ)\n        total_ls_distance += levenshtein_distance\n    \n    return total_ls_distance / VAL_STEPS","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Log\nA custom training loop requires a custom log system.","metadata":{}},{"cell_type":"code","source":"def log(batch, t_start_batch, val_ls_distance=False):\n    print(\n        f'Step %s|' % f'{batch * VERBOSE_FREQ}/{TRAIN_STEPS}'.ljust(10, ' '),\n        f'loss: %.3f,' % (train_loss.result() / VERBOSE_FREQ),\n        f'acc: %.3f, ' % train_accuracy.result(),\n    end='')\n    \n    if val_ls_distance:\n        print(\n            f'val_loss: %.3f, ' % (val_loss.result() / VERBOSE_FREQ),\n            f'val lsd: %s,' % ('%.1f' % val_ls_distance).ljust(5, ' '),\n        end='')\n    # always end with batch duration and line break\n    print(\n        f'lr: %s,' % ('%.1E' % LRREDUCE.get_lr()).ljust(7),\n        f't: %s sec' % int(time.time() - t_start_batch),\n    )","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training History\nclass to keep track of training metrics and plot them after training","metadata":{}},{"cell_type":"code","source":"class Stats():\n    def __init__(self):\n        self.stats = {\n            'train_loss': [],\n            'train_acc': [],\n        }\n        \n    def update_stats(self):\n        self.stats['train_loss'].append(train_loss.result() / VERBOSE_FREQ)\n        self.stats['train_acc'].append(train_accuracy.result())\n        \n    def get_stats(self, metric):\n        return self.stats[metric]\n        \n    def plot_stat(self, metric):\n        plt.figure(figsize=(15,8))\n        plt.xticks(fontsize=16)\n        plt.yticks(fontsize=16)\n        plt.plot(self.stats[metric])\n        plt.grid()\n        plt.title(f'{metric} stats', size=24)\n        plt.show()\n        \nSTATS = Stats()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# custom learning rate scheduler\nclass LRReduce():\n    def __init__(self, optimizer, lr_schedule):\n        self.opt = optimizer\n        self.lr_schedule = lr_schedule\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        \nLRREDUCE = LRReduce(optimizer, LR_SCHEDULE)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Custom Training Loop","metadata":{}},{"cell_type":"code","source":"step_total = 0\nfor epoch in range(EPOCHS):\n    print(f'***** EPOCH {epoch + 1} *****')\n    \n    t_start = time.time()\n    t_start_batch = time.time()\n    total_loss = 0\n    \n    # create distributed versions of dataset\n    train_dist_dataset = iter(strategy.experimental_distribute_dataset(train_dataset))\n    val_dist_dataset = iter(strategy.experimental_distribute_dataset(val_dataset))\n\n    for step in range(1, STEPS_PER_EPOCH + 1):\n        # train step\n        distributed_train_step(train_dist_dataset)\n        STATS.update_stats()\n        # save epoch weights\n        encoder.save_weights(f'./encoder_epoch_{epoch+1}.h5')\n        decoder.save_weights(f'./decoder_epoch_{epoch+1}.h5')\n            \n        # end of epoch validation\n        if step == STEPS_PER_EPOCH:\n            val_ls_distance = get_val_metrics(val_dist_dataset)\n            # log with validation\n            log(step, t_start_batch, val_ls_distance)\n        else:\n            # normal log\n            log(step, t_start_batch)\n            # reset start time batch\n            t_start_batch = time.time()\n            \n        total_loss += train_loss.result()\n        \n        # learning rate step\n        LRREDUCE.step(epoch * TRAIN_STEPS + step * VERBOSE_FREQ - 1)\n        \n        # stop training when NaN loss is detected, this can be caused by exploding gradients\n        if np.isnan(total_loss):\n            break\n            \n    # stop training when NaN loss is detected\n    if np.isnan(total_loss):\n        break\n\n    print(f'Epoch {epoch} Loss {round(total_loss.numpy() / TRAIN_STEPS, 3)}, time: {int(time.time() - t_start)} sec\\n')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training History","metadata":{}},{"cell_type":"code","source":"STATS.plot_stat('train_loss')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"STATS.plot_stat('train_acc')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prediction Examples\nShow the prediction for some validation images.","metadata":{}},{"cell_type":"code","source":"# convert the integer encoded predictions to a string\ndef int2char(i_str):\n    res = ''\n    for i in i_str.numpy():\n        c = int_to_vocabulary.get(i)\n        if c not in ['<start>', '<end>', '<pad>']:\n            res += c\n    return res","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def evaluate(img, actual=None):\n    # get encoder output and initiate LSTM hidden and carry state\n    enc_out = encoder(tf.expand_dims(img, axis=0), training=False)\n    h, c = decoder.init_hidden_state(enc_out, training=False)\n    \n    # the \"<start>\" token is used as first character when predicting\n    dec_input = tf.expand_dims([vocabulary_to_int.get('<start>')], 0)\n    result = ''\n    \n    for t in tqdm(range(SEQ_LEN_OUT)):\n        predictions, h, c = decoder(dec_input, h, c, enc_out, training=False)\n        predicted_id = tf.argmax(predictions[0]).numpy()\n        predicted_char = int_to_vocabulary.get(predicted_id)\n\n        # stop predicting when \"<end>\" token is predicted\n        if predicted_char == '<end>':\n            break\n        \n        # add every character except \"<start>\"\n        if result != '<start>':\n            result += predicted_char\n\n        # predicted charachter is used as input to the decoder to predict the next character\n        dec_input = tf.expand_dims([predicted_id], 0)\n    \n    # plot the molecule image\n    plt.figure(figsize=(7, 4))\n    plt.imshow(img.numpy().astype(np.float32))\n    plt.show()\n    print(f'predicted: \\t{result}')\n    print(f'actual: \\t{int2char(actual)}')\n\nfor n in range(3):\n    evaluate(val_imgs[n], actual=val_lbls[n])","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}