{"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":"#IMG_HEIGHT = 256\n#IMG_WIDTH = 448\n\nIMG_HEIGHT = 256\nIMG_WIDTH = 512\n\nN_CHANNELS = 3\nMAX_INCHI_LEN = 200\n\n#this works for B4 on TPU. let try GPU\n\n#BATCH_SIZE_BASE = 128 if TPU else 64\n\n#above too big for model b4 on TPU\n\nBATCH_SIZE_BASE = 32 if TPU else 16 \n\nBATCH_SIZE = BATCH_SIZE_BASE * REPLICAS\n\nN_TEST_IMGS = 1616107\nN_TEST_STEPS = N_TEST_IMGS // BATCH_SIZE + 1\n\nTARGET_DTYPE = tf.bfloat16 if TPU else tf.float32\n\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\nAUTO = tf.data.experimental.AUTOTUNE\n\nif TPU: # get Google Cloud path to dataset for TPU\n    # Given Data Train/Val/Test\n    #GCS_DS_PATH_IMGS = KaggleDatasets().get_gcs_path('molecular-translation-images-cleaned-tfrecords')\n    \n    #try new tfrecords with 512x256\n    GCS_DS_PATH_IMGS = KaggleDatasets().get_gcs_path('effnb1-tf-data')\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#with open('/kaggle/input/molecular-translation-images-cleaned-tfrecords/vocabulary_to_int.pkl', 'rb') as handle:\n\nwith open('../input/effnb1-tf-data/vocabulary_to_int.pkl', 'rb') as handle:\n    vocabulary_to_int   = pickle.load( handle)\n\n# dictionary to convert the integer encoding to vocabulary\n#with open('/kaggle/input/molecular-translation-images-cleaned-tfrecords/int_to_vocabulary.pkl', 'rb') as handle:\nwith open('../input/effnb1-tf-data/int_to_vocabulary.pkl', 'rb') as handle:\n\n    int_to_vocabulary  = pickle.load( handle)\n    \nprint(f'vocabulary_to_int head: {list(vocabulary_to_int.items())[:5]}')\nprint(f'int_to_vocabulary head: {list(int_to_vocabulary.items())[:5]}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# configure problem\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'vocabulary size: {VOCAB_SIZE}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Decodes the TFRecords to a tuple yielding the image and image_id\n@tf.function\ndef decode_tfrecord_train(record_bytes):\n    features = tf.io.parse_single_example(record_bytes, {\n        'image': tf.io.FixedLenFeature([], tf.string),\n        'image_id': tf.io.FixedLenFeature([], tf.string),\n    })\n\n    image = tf.io.decode_png(features['image'])    \n    image = tf.reshape(image, [IMG_HEIGHT, IMG_WIDTH, 1])\n    image = tf.cast(image, tf.float32)  / 255.0\n    image = (image - IMAGENET_MEAN) / IMAGENET_STD\n    image = tf.cast(image, TARGET_DTYPE)\n    \n    image_id = features['image_id']\n    \n    return image, image_id\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Benchmark function to finetune the dataset\ndef benchmark_dataset(dataset, num_epochs=3, bs=BATCH_SIZE, N_IMGS_PER_EPOCH=int(100e3) if TPU else 5000):\n    n_steps_per_epoch = N_IMGS_PER_EPOCH // (num_epochs * bs)\n    start_time = time.perf_counter()\n    for epoch_num in range(num_epochs):\n        epoch_start = time.perf_counter()\n        for idx, (images, image_id) 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            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 few images\ndef show_batch(dataset, rows=3, cols=2):\n    fig, axes = plt.subplots(nrows=rows, ncols=cols, figsize=(cols*7, rows*4))\n    imgs, img_ids = next(iter(dataset.unbatch().batch(rows*cols)))\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            axes[r, c].set_title(img_ids[r*cols+c].numpy().decode(), size=16)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#  dataset for the test images\ndef get_test_dataset(bs=BATCH_SIZE):\n    ignore_order = tf.data.Options()\n    ignore_order.experimental_deterministic = False\n    \n    if TPU:\n        FNAMES_TRAIN_TFRECORDS = tf.io.gfile.glob(f'{GCS_DS_PATH_IMGS}/test/*.tfrecords')\n    else:\n        #FNAMES_TRAIN_TFRECORDS = tf.io.gfile.glob('/kaggle/input/molecular-translation-images-cleaned-tfrecords/test/*.tfrecords')\n        FNAMES_TRAIN_TFRECORDS = tf.io.gfile.glob('../input/effnb1-tf-data/test/*.tfrecords')\n        \n    train_dataset = tf.data.TFRecordDataset(FNAMES_TRAIN_TFRECORDS, num_parallel_reads=AUTO if TPU else os.cpu_count())\n    train_dataset = train_dataset.with_options(ignore_order)\n    train_dataset = train_dataset.prefetch(AUTO)\n    train_dataset = train_dataset.map(decode_tfrecord_train, num_parallel_calls=AUTO if TPU else os.cpu_count())\n    train_dataset = train_dataset.batch(BATCH_SIZE)\n    train_dataset = train_dataset.prefetch(1)\n    \n    return train_dataset\n\ntest_dataset = get_test_dataset()\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# benchmark dataset, should bve roughly 300 images a second\nbenchmark_dataset(test_dataset)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"imgs, img_ids = next(iter(test_dataset))\nprint(f'imgs.shape: {imgs.shape}, img_ids.shape: {img_ids.shape}')\nprint(f'imgs dtype: {imgs.dtype}, img_ids dtype: {img_ids.dtype}')\nimg0 = imgs[0].numpy().astype(np.float32)\ntrain_batch_info = (img0.mean(), img0.std(), img0.min(), img0.max())\nprint('train img 0 mean: %.3f, 0 std: %.3f, min: %.3f, max: %.3f' % train_batch_info)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_batch(test_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        #self.feature_maps = efn.EfficientNetB2(include_top=False, weights='noisy-student')\n        #self.feature_maps = efn.EfficientNetB4(include_top=False, weights='noisy-student')\n        self.feature_maps = efn.EfficientNetB7(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        \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])\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_layer(tf.zeros([BATCH_SIZE, DECODER_DIM]), encoder_res, debug=True)\n\nprint('context_vector shape: (batch size, units) {}'.format(context_vector.shape))\n ","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        self.vocab_size = vocab_size\n        self.attention_units = attention_units\n        self.encoder_dim = encoder_dim\n        self.decoder_dim = decoder_dim\n        \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        self.lstm_cell = tf.keras.layers.LSTMCell(decoder_dim, name='lstm_char_predictor')\n        self.fcn = tf.keras.layers.Dense(units=vocab_size, input_shape=[decoder_dim], dtype=tf.float32, name='lstm_output_to_char_probs')\n        self.do = tf.keras.layers.Dropout(0.30, name='prediction_dropout')\n        \n        self.embedding = tf.keras.layers.Embedding(vocab_size, char_embedding_dim)\n\n        # used for attention\n        self.attention = BahdanauAttention(self.attention_units)\n\n    def call(self, char, h, c, enc_output):\n        # embed previous character\n        char = self.embedding(char, training=False)\n        char = tf.squeeze(char, axis=1)\n        # get attention alpha and context vector\n        context = self.attention(h, enc_output, training=False)\n\n        # concat context and char to create lstm input\n        lstm_input = tf.concat((context, char), axis=-1)\n        \n        # LSTM call, get new h, c\n        _, (h_new, c_new) = self.lstm_cell(lstm_input, (h, c), training=False)\n        \n        # compute predictions with dropout\n        output = self.do(h_new, training=False)\n        output = self.fcn(output, training=False)\n\n        return output, h_new, c_new\n    \n    def init_hidden_state(self, encoder_out):\n        mean_encoder_out = tf.math.reduce_mean(encoder_out, axis=1)\n        h = self.init_h(mean_encoder_out, training=False)  # (batch_size, decoder_dim)\n        c = self.init_c(mean_encoder_out, training=False)\n        return h, c","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"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":"# Models\ntf.keras.backend.clear_session()\n\n# enable XLA optmizations\ntf.config.optimizer.set_jit(True)\n\nwith strategy.scope():\n    encoder = Encoder()\n    encoder.build(input_shape=[BATCH_SIZE, IMG_HEIGHT, IMG_WIDTH, N_CHANNELS])\n    encoder_res = encoder(imgs[:BATCH_SIZE])\n    #encoder.load_weights('../input/b2-encdec/encoder_epoch_12.h5')\n    #encoder.load_weights('../input/b4-acc996-encdec/encoder_epoch_12.h5')\n    \n    #encoder.load_weights('../input/512x256-acc998-encdec/encoder_epoch_20.h5')\n    \n    #encoder.load_weights('../input/512x256-valloss-1-encdec/encoder_epoch_20.h5')\n    \n    #encoder.load_weights('../input/b7-acc997-51x256-encdec/encoder_epoch_30.h5')\n    \n    encoder.load_weights('../input/b7-lsd2d9-512x256-encdec/encoder_epoch_12.h5')\n    \n    \n    encoder.trainable = False\n    encoder.compile()\n\n    decoder = Decoder(VOCAB_SIZE, ATTENTION_UNITS, ENCODER_DIM, DECODER_DIM, CHAR_EMBEDDING_DIM)\n    h, c = decoder.init_hidden_state(encoder_res)\n    preds, h, c = decoder(tf.ones([BATCH_SIZE, 1]), h, c, encoder_res)\n    #decoder.load_weights('../input/b2-encdec/decoder_epoch_12.h5')\n    #decoder.load_weights('../input/b4-acc996-encdec/decoder_epoch_12.h5')\n    \n    #decoder.load_weights('../input/512x256-acc998-encdec/decoder_epoch_20.h5')\n    \n    #decoder.load_weights('../input/512x256-valloss-1-encdec/decoder_epoch_20.h5')\n    \n    #decoder.load_weights('../input/b7-acc997-51x256-encdec/decoder_epoch_30.h5')\n    decoder.load_weights('../input/b7-lsd2d9-512x256-encdec/decoder_epoch_12.h5')\n    \n    decoder.trainable = False\n    decoder.compile()\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"encoder.summary()\n\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"decoder.summary()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Predictions","metadata":{}},{"cell_type":"markdown","source":"Predictions\nDue to popular demand the prediction loop has been modified to run on a TPU. The last batch will be run on a single TPU core, not distributed over all 8. This adaption is needed due to the different batch size which cannot be ditributed evenly over all TPU cores and will therefore throw an error as pointed out by dragon zhang. Predictions on TPU take less than 20 minutes, about 10 times faster than on a GPU :D\n","metadata":{}},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# converts and integer encoded InChI prediction to a correct InChI string\n# Note the \"InChI=1S/\" part is prepended and all <start>/<end>/<pad> tokens are ignored\n\nEND_TOKEN = vocabulary_to_int.get('<end>')\nSTART_TOKEN = vocabulary_to_int.get('<start>')\nPAD_TOKEN =  vocabulary_to_int.get('<pad>')\n\ndef int2char(i_str):\n    res = 'InChI=1S/'\n    for i in i_str:\n        if i == END_TOKEN:\n            return res\n        elif i != START_TOKEN and i != PAD_TOKEN:\n            res += int_to_vocabulary.get(i)\n    return res","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Makes the InChI prediction for a given image\ndef prediction_step(imgs):\n    # get the feature maps from the encoder\n    encoder_res = encoder(imgs)\n    # initialize the hidden LSTM states given the feature maps\n    h, c = decoder.init_hidden_state(encoder_res)\n    \n    # initialize the prediction results with the <start> token\n    predictions_seq = tf.fill([len(imgs), 1], value=vocabulary_to_int.get('<start>'))\n    predictions_seq = tf.cast(predictions_seq, tf.int32)\n    # first encoder input is always the <start> token\n    dec_input = tf.expand_dims([vocabulary_to_int.get('<start>')] * len(imgs), 1)\n\n    # Teacher forcing - feeding the target as the next input\n    for t in range(1, SEQ_LEN_OUT):\n        # make character prediction and receive new LSTM states\n        predictions, h, c = decoder(dec_input, h, c, encoder_res)\n        \n        # softmax prediction to get prediction classes\n        dec_input = tf.math.argmax(predictions, axis=1, output_type=tf.int32)\n               \n        # expand dimension of prediction to make valid encoder input\n        dec_input = tf.expand_dims(dec_input, axis=1)\n        \n        # add character to predictions\n        predictions_seq = tf.concat([predictions_seq, dec_input], axis=1)\n            \n    return predictions_seq\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# distributed test step, will also run on TPU :D\n@tf.function\ndef distributed_test_step(imgs):\n    per_replica_predictions = strategy.run(prediction_step, args=[imgs])\n    predictions = strategy.gather(per_replica_predictions, axis=0)\n    \n    return predictions","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# perform a test step on a single device, used for last batch with random size\n@tf.function\ndef test_step_last_batch(imgs):\n    return prediction_step(imgs)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Due to popular demand the prediction loop has been modified to run on a TPU. The last batch will be run on a single TPU core, not distributed over all 8. This adaption is needed due to the different batch size which cannot be ditributed evenly over all TPU cores and will therefore throw an error as pointed out by dragon zhang. Predictions on TPU take less than 20 minutes, about 10 times faster than on a GPU :D","metadata":{}},{"cell_type":"code","source":"\n# list with predicted InChI's\npredictions_inchi = []\n# List with image id's\npredictions_img_ids = []\n# Distributed test set, needed for TPU\ntest_dist_dataset = strategy.experimental_distribute_dataset(test_dataset)\n\n# Prediction Loop\nfor step, (per_replica_imgs, per_repliac_img_ids) in tqdm(enumerate(test_dist_dataset), total=N_TEST_STEPS):\n    # special step for last batch which has a different size\n    # this step will take about half a minute because the function needs to be compiled\n    if TPU and step == N_TEST_STEPS - 1:\n        imgs_single_device = strategy.gather(per_replica_imgs, axis=0)\n        preds = test_step_last_batch(imgs_single_device)\n    else:\n        # make test step and get predictions\n        preds = distributed_test_step(per_replica_imgs)\n    \n    # get image ids\n    img_ids = strategy.gather(per_repliac_img_ids, axis=0)\n    \n    # decode integer encoded predictions to characters and add to InChI's prediction list\n    predictions_inchi += [int2char(p) for p in preds.numpy()]\n    # add image id's to list\n    predictions_img_ids += [e.decode() for e in img_ids.numpy()]\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# create DataFrame with image ids and predicted InChI's\nsubmission = pd.DataFrame({ 'image_id': predictions_img_ids, 'InChI': predictions_inchi }, dtype='string')\n# save as CSV file so we can submit it :D\nsubmission.to_csv('submission.csv', index=False)\n# show head of submission, sanity check\npd.options.display.max_colwidth = 200\nsubmission.head()\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# submission csv info, important, it should contain 1616107 rows!!!\nsubmission.info()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def normalize_inchi(inchi):\n    try:\n        mol = Chem.MolFromInchi(inchi)\n        if mol is None:\n            return inchi\n        else: \n            try:\n                return Chem.MolToInchi(mol)\n            except:\n                return inchi\n\n    except: return inchi","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"#submission['InChI']=submission.apply(lambda x: normalize_inchi(x['InChI']), axis=1)\n\n\n \n","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}