{"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\nThis is the prediction part of the TPU training notebook found [here](https://www.kaggle.com/markwijkhuizen/tensorflow-tpu-training-baseline-lb-16-92)\nAt first I was not planning to make the prediction notebook public. However, many high scoring prediction notebooks have been publishes lately, publishing mine won't flood the leaderboard anymore.\n\nThis is the prediction notebook using the trained model from [this](https://www.kaggle.com/markwijkhuizen/tensorflow-tpu-training-baseline-lb-16-92) notebook.\n\nTo save on TPU quota this notebook will use a GPU for the predictions, which will take about 3 hours.\n\nIf you have any questions, leave a comment :D\n\n*V2 UPDATES*\n\n* Forgot to turn off internet, only non-internet notebooks are allowed to make a submission.\n\n*V3 UPDATES*\n\n* Forgot to add BMS competition data, submit button was still not visible...\n\n*V4 UPDATES*\n\n* Prediction loop now runs on a TPU","metadata":{}},{"cell_type":"code","source":"# install Keras EfficientNet models with noisy-student weights\n!pip install -q '/kaggle/input/birdcall-identification-submission-custom/Keras_Applications-1.0.8-py3-none-any.whl'\n!pip install -q '/kaggle/input/birdcall-identification-submission-custom/efficientnet-1.1.0-py3-none-any.whl'","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\nfrom multiprocessing import cpu_count\n\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":"# Detect hardware, return appropriate distribution strategy\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}')\n\n# # set half precision policy\nmixed_precision.set_policy('mixed_bfloat16' if TPU else 'float32')\n\nprint(f'Compute dtype: {mixed_precision.global_policy().compute_dtype}')\nprint(f'Variable dtype: {mixed_precision.global_policy().variable_dtype}')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"IMG_HEIGHT = 256\nIMG_WIDTH = 448\nN_CHANNELS = 3\nMAX_INCHI_LEN = 200\n\nBATCH_SIZE_BASE = 128 if TPU else 64\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')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# dictionary to integer encode the vocabulary\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 convert the integer encoding to vocabulary\nwith open('/kaggle/input/molecular-translation-images-cleaned-tfrecords/int_to_vocabulary.pkl', 'rb') as handle:\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":"markdown","source":"# Dataset","metadata":{}},{"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","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        \n    train_dataset = tf.data.TFRecordDataset(FNAMES_TRAIN_TFRECORDS, num_parallel_reads=AUTO if TPU else 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 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()","metadata":{"trusted":true},"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","metadata":{}},{"cell_type":"code","source":"class Encoder(tf.keras.Model):\n    def __init__(self):\n        super(Encoder, self).__init__()\n        \n        self.feature_maps = efn.EfficientNetB0(include_top=False, weights=None)\n        global ENCODER_DIM\n        ENCODER_DIM = self.feature_maps.layers[-1].output_shape[-1]\n        \n        self.reshape = tf.keras.layers.Reshape([-1, ENCODER_DIM], name='reshape_featuere_maps')\n\n    def call(self, x, training):\n        x = self.feature_maps(x, training=False)\n        x = self.reshape(x, training=False)\n        return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Attention","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):\n        # dense hidden state to attention units size and expand dimension\n        h_expand = tf.expand_dims(h, axis=1) # expand dimension\n            \n        h_dense = self.H(h_expand, training=False)\n        \n        # dense features to units size\n        encoder_res_dense = self.E(encoder_res, training=False) # dense to attention\n\n        # add vectors\n        score = tf.nn.relu(h_dense + encoder_res_dense)\n        score = self.V(score, training=False)\n        \n        # create alpha vector size (bs, layers)        \n        attention_weights = tf.nn.softmax(score, axis=1)\n        \n        # create attention weights (bs, layers)\n        context_vector = encoder_res * attention_weights\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":"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        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":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"START_TOKEN = tf.constant(vocabulary_to_int.get('<start>'), dtype=tf.int32)\nEND_TOKEN = tf.constant(vocabulary_to_int.get('<end>'), dtype=tf.int32)\nPAD_TOKEN = tf.constant(vocabulary_to_int.get('<pad>'), dtype=tf.int32)","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('/kaggle/input/tensorflow-tpu-training-baseline/encoder_epoch_10.h5')\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('/kaggle/input/tensorflow-tpu-training-baseline/decoder_epoch_10.h5')\n    decoder.trainable = False\n    decoder.compile()","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":"# Prediction Step","metadata":{}},{"cell_type":"markdown","source":"Improved performance of `int2char` function by using sets(constant lookup time) and returning on `<end>` token","metadata":{}},{"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","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":"markdown","source":"# Predictions","metadata":{}},{"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](https://www.kaggle.com/dragonzhang). Predictions on TPU take less than 20 minutes, about 10 times faster than on a GPU :D","metadata":{}},{"cell_type":"code","source":"# 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()]","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()","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":[]}]}