{"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":"This notebook is a baseline for an encoder and decoder model written in Tensorflow and running on a TPU. Several notebooks, examples and documentation were used as a source of inspiration","metadata":{}},{"cell_type":"code","source":"!pip install -q --upgrade pip\n!pip install -q efficientnet","metadata":{"execution":{"iopub.status.busy":"2021-06-10T12:13:45.904824Z","iopub.execute_input":"2021-06-10T12:13:45.905188Z","iopub.status.idle":"2021-06-10T12:14:15.412031Z","shell.execute_reply.started":"2021-06-10T12:13:45.905156Z","shell.execute_reply":"2021-06-10T12:14:15.410813Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport tensorflow as tf\nfrom tensorflow import keras\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":{"execution":{"iopub.status.busy":"2021-06-10T12:14:15.414665Z","iopub.execute_input":"2021-06-10T12:14:15.414965Z","iopub.status.idle":"2021-06-10T12:14:16.223741Z","shell.execute_reply.started":"2021-06-10T12:14:15.414934Z","shell.execute_reply":"2021-06-10T12:14:16.222786Z"},"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","metadata":{"execution":{"iopub.status.busy":"2021-06-10T12:15:01.9284Z","iopub.execute_input":"2021-06-10T12:15:01.928754Z","iopub.status.idle":"2021-06-10T12:15:01.935045Z","shell.execute_reply.started":"2021-06-10T12:15:01.928725Z","shell.execute_reply":"2021-06-10T12:15:01.934003Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if 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}')\n\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":{"execution":{"iopub.status.busy":"2021-06-10T12:15:10.118382Z","iopub.execute_input":"2021-06-10T12:15:10.118769Z","iopub.status.idle":"2021-06-10T12:15:15.860846Z","shell.execute_reply.started":"2021-06-10T12:15:10.118736Z","shell.execute_reply":"2021-06-10T12:15:15.85998Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEBUG = False\n\nIMG_HEIGHT = 256\nIMG_WIDTH = 448\nN_CHANNELS = 3\n\nMAX_INCHI_LEN = 200\n\nBATCH_SIZE_BASE = 6 if DEBUG else (64 if TPU else 12)\nBATCH_SIZE = BATCH_SIZE_BASE*REPLICAS\nBATCH_SIZE_DEBUG = 2\n\nN_TEST_IMGS = 1616107\nN_TEST_STEPS = N_TEST_IMGS // BATCH_SIZE + 1\n\nTARGET_DTYPE = tf.bfloat16 if TPU else tf.float32\nLABEL_DTYPE = tf.uint8\n\nVAL_SIZE = int(1e3) if DEBUG else int(100e3)\nVAL_STEPS = VAL_SIZE // BATCH_SIZE\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\nif TPU:\n    GCS_DS_PATH = KaggleDatasets().get_gcs_path('molecular-translation-images-cleaned-tfrecords')\n    \nAUTO = tf.data.experimental.AUTOTUNE","metadata":{"execution":{"iopub.status.busy":"2021-06-10T12:17:13.294733Z","iopub.execute_input":"2021-06-10T12:17:13.295085Z","iopub.status.idle":"2021-06-10T12:17:13.592152Z","shell.execute_reply.started":"2021-06-10T12:17:13.295057Z","shell.execute_reply":"2021-06-10T12:17:13.591204Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open('../input/molecular-translation-images-cleaned-tfrecords/vocabulary_to_int.pkl', 'rb') as handle:\n    vocabulary_to_int = pickle.load( handle)\n    \nwith open('../input/molecular-translation-images-cleaned-tfrecords/int_to_vocabulary.pkl', 'rb') as handle:\n    int_to_vocabulary = pickle.load( handle)\n    \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":{"execution":{"iopub.status.busy":"2021-06-10T12:17:43.573156Z","iopub.execute_input":"2021-06-10T12:17:43.573521Z","iopub.status.idle":"2021-06-10T12:17:43.605326Z","shell.execute_reply.started":"2021-06-10T12:17:43.573493Z","shell.execute_reply":"2021-06-10T12:17:43.604221Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"VOCAB_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-06-10T12:17:51.42741Z","iopub.execute_input":"2021-06-10T12:17:51.427765Z","iopub.status.idle":"2021-06-10T12:17:51.433187Z","shell.execute_reply.started":"2021-06-10T12:17:51.427735Z","shell.execute_reply":"2021-06-10T12:17:51.432122Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@tf.function\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":{"execution":{"iopub.status.busy":"2021-06-10T12:18:01.058459Z","iopub.execute_input":"2021-06-10T12:18:01.059256Z","iopub.status.idle":"2021-06-10T12:18:01.068744Z","shell.execute_reply.started":"2021-06-10T12:18:01.05918Z","shell.execute_reply":"2021-06-10T12:18:01.06761Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_dataset(bs=BATCH_SIZE, val=False):\n    ignore_order = tf.data.Options()\n    ignore_order.experimental_deterministic = False\n    \n    if val:\n        FNAMES_TRAIN_TFRECORDS = tf.io.gfile.glob(f'{GCS_DS_PATH}/val/*.tfrecords')\n    \n    FNAMES_TRAIN_TFRECORDS = tf.io.gfile.glob(f'{GCS_DS_PATH}/train/*.tfrecords')\n    dataset = tf.data.TFRecordDataset(FNAMES_TRAIN_TFRECORDS, num_parallel_reads=AUTO)\n    dataset = dataset.with_options(ignore_order)\n    dataset = dataset.prefetch(AUTO) \n    dataset = dataset.repeat()\n    dataset = dataset.map(decode_tfrecord, num_parallel_calls=AUTO)\n    dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)\n    dataset = dataset.prefetch(1) \n    \n    return dataset\n\ntrain_dataset = get_dataset()","metadata":{"execution":{"iopub.status.busy":"2021-06-10T12:18:20.766692Z","iopub.execute_input":"2021-06-10T12:18:20.76725Z","iopub.status.idle":"2021-06-10T12:18:21.130028Z","shell.execute_reply.started":"2021-06-10T12:18:20.7672Z","shell.execute_reply":"2021-06-10T12:18:21.128951Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"val_dataset = get_dataset(val=True)","metadata":{"execution":{"iopub.status.busy":"2021-06-10T12:18:58.68834Z","iopub.execute_input":"2021-06-10T12:18:58.688925Z","iopub.status.idle":"2021-06-10T12:18:58.939534Z","shell.execute_reply.started":"2021-06-10T12:18:58.688875Z","shell.execute_reply":"2021-06-10T12:18:58.938591Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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='noisy-student')\n        \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, 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":{"execution":{"iopub.status.busy":"2021-06-10T12:19:17.492714Z","iopub.execute_input":"2021-06-10T12:19:17.493086Z","iopub.status.idle":"2021-06-10T12:19:17.501085Z","shell.execute_reply.started":"2021-06-10T12:19:17.493058Z","shell.execute_reply":"2021-06-10T12:19:17.499943Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"imgs, 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(), img0.dtype)\nprint('train img0 mean: %.3f, std: %.3f, min: %.3f, max: %.3f, %s'%train_batch_info)","metadata":{"execution":{"iopub.status.busy":"2021-06-10T12:19:29.252769Z","iopub.execute_input":"2021-06-10T12:19:29.253167Z","iopub.status.idle":"2021-06-10T12:19:35.950839Z","shell.execute_reply.started":"2021-06-10T12:19:29.253134Z","shell.execute_reply":"2021-06-10T12:19:35.949771Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with tf.device('/CPU:0'):\n    encoder = Encoder()\n    encoder_res = encoder(imgs[:BATCH_SIZE_DEBUG], debug = True)\n    \nprint('Encode output shape: (batch_size, seq_len, units) {}'.format(encoder_res.shape))","metadata":{"execution":{"iopub.status.busy":"2021-06-10T12:19:47.581623Z","iopub.execute_input":"2021-06-10T12:19:47.582156Z","iopub.status.idle":"2021-06-10T12:20:22.996891Z","shell.execute_reply.started":"2021-06-10T12:19:47.582109Z","shell.execute_reply":"2021-06-10T12:20:22.995461Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Decoder(keras.Model):\n    def __init__(self, vocab_size, encoder_dim, decoder_dim, char_embedding_dim):\n        super(Decoder, self).__init__()\n        self.init_h = keras.layers.Dense(units=decoder_dim, input_shape=[encoder_dim], name='encoder_res_to_hiddent_init')\n        self.init_c = keras.layers.Dense(units=decoder_dim, input_shape=[encoder_dim], name='encoder_res_to_inp_act_init')\n        self.lstm_cell = keras.layers.LSTMCell(decoder_dim, name='lstm_char_predictor')\n        self.do = keras.layers.Dropout(0.3, name='prediction_dropout')\n        self.fcn = keras.layers.Dense(units=vocab_size, input_shape=[decoder_dim], dtype=tf.float32, name='lstm_output_to_char_probs')\n        self.embedding = keras.layers.Embedding(vocab_size, char_embedding_dim, name='char_embedding')\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        char = self.embedding(char, training=training)\n        char = tf.squeeze(char, axis=1)\n\n        lstm_input = char\n        \n        if debug:\n            print(f'lstm_input shape: {lstm_input.shape}')\n        _, (h_new, c_new) = self.lstm_cell(lstm_input, (h, c), training=training)\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)\n        c = self.init_c(mean_encoder_out, training=training)\n        \n        return h, c","metadata":{"execution":{"iopub.status.busy":"2021-06-10T12:20:22.998654Z","iopub.execute_input":"2021-06-10T12:20:22.998995Z","iopub.status.idle":"2021-06-10T12:20:23.014161Z","shell.execute_reply.started":"2021-06-10T12:20:22.998964Z","shell.execute_reply":"2021-06-10T12:20:23.013326Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with tf.device('/CPU:0'):\n    decoder = Decoder(VOCAB_SIZE, 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    print('Decoder output shape: (batch_size, vocab_size {}'.format(preds.shape))","metadata":{"execution":{"iopub.status.busy":"2021-06-10T12:20:23.015473Z","iopub.execute_input":"2021-06-10T12:20:23.015891Z","iopub.status.idle":"2021-06-10T12:20:23.296381Z","shell.execute_reply.started":"2021-06-10T12:20:23.015862Z","shell.execute_reply":"2021-06-10T12:20:23.295208Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"START_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":{"execution":{"iopub.status.busy":"2021-06-10T12:20:55.277752Z","iopub.execute_input":"2021-06-10T12:20:55.27813Z","iopub.status.idle":"2021-06-10T12:20:55.28513Z","shell.execute_reply.started":"2021-06-10T12:20:55.278099Z","shell.execute_reply":"2021-06-10T12:20:55.283642Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tf.keras.backend.clear_session()\n\nwith strategy.scope():\n    mixed_precision.set_policy('mixed_bfloat16' if TPU else 'float32')\n    \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    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, 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-06-10T12:21:05.109565Z","iopub.execute_input":"2021-06-10T12:21:05.109953Z","iopub.status.idle":"2021-06-10T12:21:54.532376Z","shell.execute_reply.started":"2021-06-10T12:21:05.10992Z","shell.execute_reply":"2021-06-10T12:21:54.531384Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"EPOCHS = 1\nWARMUP_STEPS = 500\nTRAIN_STEPS = 1000\nVERBOSE_FREQ = 100\nSTEPS_PER_EPOCH = TRAIN_STEPS // VERBOSE_FREQ\nTOTAL_STEPS = EPOCHS * TRAIN_STEPS","metadata":{"execution":{"iopub.status.busy":"2021-06-10T12:21:54.533552Z","iopub.execute_input":"2021-06-10T12:21:54.533957Z","iopub.status.idle":"2021-06-10T12:21:54.537492Z","shell.execute_reply.started":"2021-06-10T12:21:54.533929Z","shell.execute_reply":"2021-06-10T12:21:54.536754Z"},"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":{"execution":{"iopub.status.busy":"2021-06-10T12:21:54.538925Z","iopub.execute_input":"2021-06-10T12:21:54.539436Z","iopub.status.idle":"2021-06-10T12:21:54.557951Z","shell.execute_reply.started":"2021-06-10T12:21:54.539393Z","shell.execute_reply":"2021-06-10T12:21:54.557009Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def 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":{"execution":{"iopub.status.busy":"2021-06-10T12:21:54.559179Z","iopub.execute_input":"2021-06-10T12:21:54.559518Z","iopub.status.idle":"2021-06-10T12:21:54.580608Z","shell.execute_reply.started":"2021-06-10T12:21:54.559488Z","shell.execute_reply":"2021-06-10T12:21:54.579626Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@tf.function()\ndef distributed_train_step(dataset):\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            for idx in range(1, SEQ_LEN_OUT):\n                t = targ[:, idx]\n                t = tf.reshape(t, [BATCH_SIZE_BASE])\n                predictions, h, c = decoder(dec_input, h, c, enc_output, training=True)\n                total_loss += loss_function(t, predictions)\n                train_accuracy.update_state(t, predictions)\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    train_loss.reset_states()\n    train_accuracy.reset_states()\n    \n    for _ in tf.range(tf.convert_to_tensor(VERBOSE_FREQ)):\n        strategy.run(train_step, args=next(dataset))","metadata":{"execution":{"iopub.status.busy":"2021-06-10T12:21:54.582037Z","iopub.execute_input":"2021-06-10T12:21:54.582324Z","iopub.status.idle":"2021-06-10T12:21:54.597548Z","shell.execute_reply.started":"2021-06-10T12:21:54.582297Z","shell.execute_reply":"2021-06-10T12:21:54.596342Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2021-06-10T12:21:58.715228Z","iopub.execute_input":"2021-06-10T12:21:58.71562Z","iopub.status.idle":"2021-06-10T12:21:58.724682Z","shell.execute_reply.started":"2021-06-10T12:21:58.715588Z","shell.execute_reply":"2021-06-10T12:21:58.723541Z"},"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":{"execution":{"iopub.status.busy":"2021-06-10T12:22:05.996677Z","iopub.execute_input":"2021-06-10T12:22:05.997155Z","iopub.status.idle":"2021-06-10T12:22:06.00263Z","shell.execute_reply.started":"2021-06-10T12:22:05.997125Z","shell.execute_reply":"2021-06-10T12:22:06.001458Z"},"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":{"execution":{"iopub.status.busy":"2021-06-10T12:22:14.342433Z","iopub.execute_input":"2021-06-10T12:22:14.342822Z","iopub.status.idle":"2021-06-10T12:22:14.349401Z","shell.execute_reply.started":"2021-06-10T12:22:14.34279Z","shell.execute_reply":"2021-06-10T12:22:14.348523Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2021-06-10T12:22:21.818724Z","iopub.execute_input":"2021-06-10T12:22:21.819592Z","iopub.status.idle":"2021-06-10T12:22:21.827187Z","shell.execute_reply.started":"2021-06-10T12:22:21.819549Z","shell.execute_reply":"2021-06-10T12:22:21.82583Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2021-06-10T12:22:32.345092Z","iopub.execute_input":"2021-06-10T12:22:32.34548Z","iopub.status.idle":"2021-06-10T12:22:32.355004Z","shell.execute_reply.started":"2021-06-10T12:22:32.345447Z","shell.execute_reply":"2021-06-10T12:22:32.353495Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"LR_SCHEDULE = [lrfn(step, 1e-8, 2e-3, 1e-4 ,EPOCHS) for step in range(TOTAL_STEPS)]\n\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":{"execution":{"iopub.status.busy":"2021-06-10T12:22:48.574397Z","iopub.execute_input":"2021-06-10T12:22:48.57475Z","iopub.status.idle":"2021-06-10T12:22:48.615191Z","shell.execute_reply.started":"2021-06-10T12:22:48.574718Z","shell.execute_reply":"2021-06-10T12:22:48.614244Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"step_total = 0\nfor epoch in range(EPOCHS):\n    print(f'*****EPOCH: {epoch+1}*****')\n    t_start = time.time()\n    t_start_batch = time.time()\n    total_loss = 0\n    \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        distributed_train_step(train_dist_dataset)\n        STATS.update_stats()\n        encoder.save_weights(f'./encoder_epoch_{epoch+1}.h5')\n        decoder.save_weights(f'./decoder_epoch_{epoch+1}.h5')\n        \n        if step == STEPS_PER_EPOCH:\n            val_ls_distance = get_val_metrics(val_dist_dataset)\n            log(step, t_start_batch, val_ls_distance)\n        else:\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        LRREDUCE.step(epoch * TRAIN_STEPS + step * VERBOSE_FREQ - 1)\n        \n        if np.isnan(total_loss):\n            break\n            \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-06-10T12:22:59.923488Z","iopub.execute_input":"2021-06-10T12:22:59.923881Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"END_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":"@tf.function\ndef decode_tfrecord_test(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":"def 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}/test/*.tfrecords')\n    else:\n        FNAMES_TRAIN_TFRECORDS = tf.io.gfile.glob('/kaggle/input/molecular-translation-images-cleaned-tfrecords/test/*.tfrecords')\n        \n    test_dataset = tf.data.TFRecordDataset(FNAMES_TRAIN_TFRECORDS, num_parallel_reads=AUTO if TPU else cpu_count())\n    test_dataset = test_dataset.with_options(ignore_order)\n    test_dataset = test_dataset.prefetch(AUTO)\n    test_dataset = test_dataset.map(decode_tfrecord_test, num_parallel_calls=AUTO if TPU else cpu_count())\n    test_dataset = test_dataset.batch(BATCH_SIZE)\n    test_dataset = test_dataset.prefetch(1)\n    \n    return test_dataset\n\ntest_dataset = get_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":"# 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('./encoder_epoch_1.h5')\n    encoder.trainable = False\n    encoder.compile()\n\n    decoder = Decoder(VOCAB_SIZE, ENCODER_DIM, DECODER_DIM, CHAR_EMBEDDING_DIM)\n    h, c = decoder.init_hidden_state(encoder_res, training=False)\n    preds, h, c = decoder(tf.ones([BATCH_SIZE, 1]), h, c, encoder_res)\n    decoder.load_weights('./decoder_epoch_1.h5')\n    decoder.trainable = False\n    decoder.compile()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def 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, training=False)\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":"@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":"@tf.function\ndef test_step_last_batch(imgs):\n    return prediction_step(imgs)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"predictions_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":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission = pd.DataFrame({ 'image_id': predictions_img_ids, 'InChI': predictions_inchi }, dtype='string')\nsubmission.head()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission.to_csv('submission.csv', index=False)","metadata":{},"execution_count":null,"outputs":[]}]}