{"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":"# Installs","metadata":{}},{"cell_type":"code","source":"!pip install -q --upgrade pip\n!pip install -q efficientnet","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-05-22T04:11:24.305297Z","iopub.execute_input":"2021-05-22T04:11:24.306011Z","iopub.status.idle":"2021-05-22T04:11:53.955484Z","shell.execute_reply.started":"2021-05-22T04:11:24.305892Z","shell.execute_reply":"2021-05-22T04:11:53.954593Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"import unicodedata, re, os, io, time, pickle, math, random\nimport numpy as np\nimport 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","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:11:53.957Z","iopub.execute_input":"2021-05-22T04:11:53.957507Z","iopub.status.idle":"2021-05-22T04:11:59.376116Z","shell.execute_reply.started":"2021-05-22T04:11:53.957471Z","shell.execute_reply":"2021-05-22T04:11:59.375333Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Settings","metadata":{}},{"cell_type":"code","source":"SEED = 42\nos.environ['PYTHONHASHSEED'] = str(SEED)\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntf.random.set_seed(SEED)","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:12:03.219599Z","iopub.execute_input":"2021-05-22T04:12:03.220092Z","iopub.status.idle":"2021-05-22T04:12:03.226121Z","shell.execute_reply.started":"2021-05-22T04:12:03.220062Z","shell.execute_reply":"2021-05-22T04:12:03.224953Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"try:\n    TPU = tf.distribute.cluster_resolver.TPUClusterResolver()\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()\n\nREPLICAS = strategy.num_replicas_in_sync\nprint(f'REPLICAS: {REPLICAS}')","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:12:04.822269Z","iopub.execute_input":"2021-05-22T04:12:04.822796Z","iopub.status.idle":"2021-05-22T04:12:10.406636Z","shell.execute_reply.started":"2021-05-22T04:12:04.822765Z","shell.execute_reply":"2021-05-22T04:12:10.405721Z"},"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":{"execution":{"iopub.status.busy":"2021-05-22T04:12:10.408093Z","iopub.execute_input":"2021-05-22T04:12:10.408408Z","iopub.status.idle":"2021-05-22T04:12:10.741707Z","shell.execute_reply.started":"2021-05-22T04:12:10.408377Z","shell.execute_reply":"2021-05-22T04:12:10.740728Z"},"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":{"execution":{"iopub.status.busy":"2021-05-22T04:12:10.743188Z","iopub.execute_input":"2021-05-22T04:12:10.74347Z","iopub.status.idle":"2021-05-22T04:12:10.766826Z","shell.execute_reply.started":"2021-05-22T04:12:10.743444Z","shell.execute_reply":"2021-05-22T04:12:10.765776Z"},"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":{"execution":{"iopub.status.busy":"2021-05-22T04:12:11.573555Z","iopub.execute_input":"2021-05-22T04:12:11.573909Z","iopub.status.idle":"2021-05-22T04:12:11.579518Z","shell.execute_reply.started":"2021-05-22T04:12:11.573881Z","shell.execute_reply":"2021-05-22T04:12:11.578269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Make dataset","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\n\n# 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}')\n        \n# 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)\n\ndef 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","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:12:13.55664Z","iopub.execute_input":"2021-05-22T04:12:13.557008Z","iopub.status.idle":"2021-05-22T04:12:13.573781Z","shell.execute_reply.started":"2021-05-22T04:12:13.556974Z","shell.execute_reply":"2021-05-22T04:12:13.572921Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = get_train_dataset()\nbenchmark_dataset(train_dataset)","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:12:16.516164Z","iopub.execute_input":"2021-05-22T04:12:16.516536Z","iopub.status.idle":"2021-05-22T04:12:24.317283Z","shell.execute_reply.started":"2021-05-22T04:12:16.516504Z","shell.execute_reply":"2021-05-22T04:12:24.316245Z"},"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":{"execution":{"iopub.status.busy":"2021-05-22T04:12:24.31879Z","iopub.execute_input":"2021-05-22T04:12:24.319119Z","iopub.status.idle":"2021-05-22T04:12:25.206866Z","shell.execute_reply.started":"2021-05-22T04:12:24.31909Z","shell.execute_reply":"2021-05-22T04:12:25.206044Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# show first few train images\nshow_batch(train_dataset)","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:12:26.449425Z","iopub.execute_input":"2021-05-22T04:12:26.449798Z","iopub.status.idle":"2021-05-22T04:12:27.664848Z","shell.execute_reply.started":"2021-05-22T04:12:26.449766Z","shell.execute_reply":"2021-05-22T04:12:27.663834Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Validation","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()\nbenchmark_dataset(val_dataset)","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:12:28.242869Z","iopub.execute_input":"2021-05-22T04:12:28.243227Z","iopub.status.idle":"2021-05-22T04:12:35.516902Z","shell.execute_reply.started":"2021-05-22T04:12:28.243196Z","shell.execute_reply":"2021-05-22T04:12:35.5159Z"},"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":{"execution":{"iopub.status.busy":"2021-05-22T04:12:35.520081Z","iopub.execute_input":"2021-05-22T04:12:35.520374Z","iopub.status.idle":"2021-05-22T04:12:35.627609Z","shell.execute_reply.started":"2021-05-22T04:12:35.520345Z","shell.execute_reply":"2021-05-22T04:12:35.626431Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_batch(val_dataset)","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:12:37.649413Z","iopub.execute_input":"2021-05-22T04:12:37.649749Z","iopub.status.idle":"2021-05-22T04:12:38.666012Z","shell.execute_reply.started":"2021-05-22T04:12:37.649721Z","shell.execute_reply":"2021-05-22T04:12:38.665042Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Encoder","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.EfficientNetB2(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\n    \n# 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":{"execution":{"iopub.status.busy":"2021-05-22T04:12:39.877858Z","iopub.execute_input":"2021-05-22T04:12:39.878394Z","iopub.status.idle":"2021-05-22T04:12:43.340113Z","shell.execute_reply.started":"2021-05-22T04:12:39.878347Z","shell.execute_reply":"2021-05-22T04:12:43.339035Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Attension","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":{"execution":{"iopub.status.busy":"2021-05-22T04:12:48.223396Z","iopub.execute_input":"2021-05-22T04:12:48.224099Z","iopub.status.idle":"2021-05-22T04:12:48.244458Z","shell.execute_reply.started":"2021-05-22T04:12:48.224054Z","shell.execute_reply":"2021-05-22T04:12:48.243414Z"},"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":{"execution":{"iopub.status.busy":"2021-05-22T04:12:49.906872Z","iopub.execute_input":"2021-05-22T04:12:49.907253Z","iopub.status.idle":"2021-05-22T04:12:49.981357Z","shell.execute_reply.started":"2021-05-22T04:12:49.907221Z","shell.execute_reply":"2021-05-22T04:12:49.980296Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Decoder","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":{"execution":{"iopub.status.busy":"2021-05-22T04:12:52.291607Z","iopub.execute_input":"2021-05-22T04:12:52.292135Z","iopub.status.idle":"2021-05-22T04:12:52.306753Z","shell.execute_reply.started":"2021-05-22T04:12:52.292084Z","shell.execute_reply":"2021-05-22T04:12:52.306006Z"},"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":{"execution":{"iopub.status.busy":"2021-05-22T04:12:54.446111Z","iopub.execute_input":"2021-05-22T04:12:54.446677Z","iopub.status.idle":"2021-05-22T04:12:54.693983Z","shell.execute_reply.started":"2021-05-22T04:12:54.446626Z","shell.execute_reply":"2021-05-22T04:12:54.692736Z"},"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)\n\ntf.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":{"execution":{"iopub.status.busy":"2021-05-22T04:12:56.659373Z","iopub.execute_input":"2021-05-22T04:12:56.659939Z","iopub.status.idle":"2021-05-22T04:13:29.626473Z","shell.execute_reply.started":"2021-05-22T04:12:56.659893Z","shell.execute_reply":"2021-05-22T04:13:29.625478Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"encoder.summary()","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:13:29.62785Z","iopub.execute_input":"2021-05-22T04:13:29.628143Z","iopub.status.idle":"2021-05-22T04:13:29.65329Z","shell.execute_reply.started":"2021-05-22T04:13:29.628114Z","shell.execute_reply":"2021-05-22T04:13:29.652205Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"decoder.summary()","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:19:35.513991Z","iopub.execute_input":"2021-05-22T04:19:35.514437Z","iopub.status.idle":"2021-05-22T04:19:35.523163Z","shell.execute_reply.started":"2021-05-22T04:19:35.514401Z","shell.execute_reply":"2021-05-22T04:19:35.521444Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Learning Rate Scheduler","metadata":{}},{"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)\n\n# 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# Training configuration\nEPOCHS = 32\nWARMUP_STEPS = 500\nTRAIN_STEPS = 1000\nVERBOSE_FREQ = 100\nSTEPS_PER_EPOCH = TRAIN_STEPS // VERBOSE_FREQ\nTOTAL_STEPS = EPOCHS * TRAIN_STEPS\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":{"execution":{"iopub.status.busy":"2021-05-22T04:19:39.833173Z","iopub.execute_input":"2021-05-22T04:19:39.833568Z","iopub.status.idle":"2021-05-22T04:19:40.355451Z","shell.execute_reply.started":"2021-05-22T04:19:39.833533Z","shell.execute_reply":"2021-05-22T04:19:40.354677Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","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\n\n@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))\n        \ndef 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\n\n@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\n\ndef 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\n\ndef 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    )\n    \nclass 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        \n# 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","metadata":{"execution":{"iopub.status.busy":"2021-05-22T04:19:45.562709Z","iopub.execute_input":"2021-05-22T04:19:45.563251Z","iopub.status.idle":"2021-05-22T04:19:45.596315Z","shell.execute_reply.started":"2021-05-22T04:19:45.563204Z","shell.execute_reply":"2021-05-22T04:19:45.595237Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"STATS = Stats()\nLRREDUCE = LRReduce(optimizer, LR_SCHEDULE)\n\nstep_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        if epoch >= 16:\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":{"execution":{"iopub.status.busy":"2021-05-22T04:19:54.323882Z","iopub.execute_input":"2021-05-22T04:19:54.324238Z","iopub.status.idle":"2021-05-22T05:09:56.429389Z","shell.execute_reply.started":"2021-05-22T04:19:54.324206Z","shell.execute_reply":"2021-05-22T05:09:56.427318Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# History","metadata":{}},{"cell_type":"code","source":"STATS.plot_stat('train_loss')\nSTATS.plot_stat('train_acc')","metadata":{},"execution_count":null,"outputs":[]},{"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\n\ndef 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":{},"execution_count":null,"outputs":[]}]}