{"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":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\n#for dirname, _, filenames in os.walk('/kaggle/input'):\n    #for filename in filenames:\n        #print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"0f1dbb20-88e0-4349-9146-93f6c2552c2d","_cell_guid":"fb7823a2-0a0f-41fc-8636-c4336690dc2b","collapsed":false,"papermill":{"duration":0.034077,"end_time":"2021-03-22T14:43:11.721067","exception":false,"start_time":"2021-03-22T14:43:11.68699","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from IPython.core.magic import register_cell_magic\n@register_cell_magic\ndef skip(line, cell=None):\n    '''Skips execution of the current line/cell if line evaluates to True.'''\n    if eval(line):\n        return\n        \n    get_ipython().run_cell(cell)","metadata":{"_uuid":"e96d48c5-c2ec-46a9-9813-1e3a92f696ae","_cell_guid":"1300590a-d283-4192-b4fd-169b1d2e8bf4","collapsed":false,"_kg_hide-input":true,"_kg_hide-output":true,"papermill":{"duration":0.034089,"end_time":"2021-03-22T14:43:11.782703","exception":false,"start_time":"2021-03-22T14:43:11.748614","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport tensorflow as tf\nfrom tensorflow import keras \nfrom tensorflow.keras import backend as K \nfrom tensorflow.keras.preprocessing.text import Tokenizer\nfrom tensorflow.keras.preprocessing.sequence import pad_sequences\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport os, cv2\nimport random\nimport dill\nimport gc\nimport time\nfrom kaggle_datasets import KaggleDatasets\nimport tensorflow_datasets.public_api as tfds\nimport math\nfrom tqdm.notebook import tqdm\nimport tensorflow_addons as tfa\nfrom tensorflow.python.distribute import values as value_lib\nfrom mt_utils import CstTokenizer","metadata":{"_uuid":"98bc8d20-9708-4be9-aa18-9dbf2a79cdee","_cell_guid":"fc5e8cb7-7686-49c6-9d02-32a5efe43de8","collapsed":false,"papermill":{"duration":7.328647,"end_time":"2021-03-22T14:43:19.13692","exception":false,"start_time":"2021-03-22T14:43:11.808273","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed=123456789\nrandom.seed(seed)\nnp.random.seed(seed)\ntf.random.set_seed(seed)\nos.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'","metadata":{"_uuid":"45fcbc4f-ff27-4a94-a52e-5f91ab056f1e","_cell_guid":"56e7c069-b99e-4f10-bec5-2c946324cc5e","collapsed":false,"_kg_hide-output":true,"papermill":{"duration":0.034434,"end_time":"2021-03-22T14:43:19.198157","exception":false,"start_time":"2021-03-22T14:43:19.163723","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# NEW on TPU in TensorFlow 24: shorter cross-compatible TPU/GPU/multi-GPU/cluster-GPU detection code\ntpu = None\ntry: # detect TPUs\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect() # TPU detection\n    strategy = tf.distribute.TPUStrategy(tpu)\nexcept ValueError: # detect GPUs\n    #strategy = tf.distribute.MirroredStrategy() # for GPU or multi-GPU machines\n    strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n    #strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy() # for clusters of multi-GPU machines\n\n    \n#strategy,tpu = tf.distribute.MirroredStrategy(devices=[\"TPU:0\", \"TPU:1\",\"TPU:2\"]),True\n\nprint(\"Number of accelerators: \", strategy.num_replicas_in_sync)\n\n\nAUTO = tf.data.experimental.AUTOTUNE\nREPLICAS = strategy.num_replicas_in_sync\nprint(f'REPLICAS: {REPLICAS}')","metadata":{"_uuid":"dfacfaff-d743-43f0-aa0c-db4224435e0c","_cell_guid":"faa70578-8f49-42b2-9acf-12a3f2401cd8","collapsed":false,"_kg_hide-input":true,"papermill":{"duration":6.120063,"end_time":"2021-03-22T14:43:25.344807","exception":false,"start_time":"2021-03-22T14:43:19.224744","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# References:\nhttps://www.tensorflow.org/tutorials/distribute/custom_training\n\nhttps://www.tensorflow.org/guide/tpu\n\nhttps://www.tensorflow.org/tutorials/text/image_captioning\n\nhttps://www.tensorflow.org/addons/tutorials/networks_seq2seq_nmt\n\n#### Dataset links : https://www.kaggle.com/tchaye59/mt-tfrecord-custom-vocab & https://www.kaggle.com/tchaye59/mtcustomvocabimg\n#### Pretraining : https://www.kaggle.com/tchaye59/mt-pretraining","metadata":{"_uuid":"1f7d0c2e-8527-4e99-a8c6-a51844c13d11","_cell_guid":"7bf2345e-b420-4d15-b038-dcf689770fb6","trusted":true}},{"cell_type":"code","source":"GCS_PATH = KaggleDatasets().get_gcs_path('mtcustomvocabimg')\nSUBMIT = True","metadata":{"_uuid":"32947c42-3ffb-4d7d-9872-508182d8db90","_cell_guid":"c7c0cdb8-c285-4c79-bd33-3d3dd3eb043f","collapsed":false,"papermill":{"duration":0.389381,"end_time":"2021-03-22T14:43:25.820254","exception":false,"start_time":"2021-03-22T14:43:25.430873","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"start = '<start>'\nend = '<end>'\nmax_seq = 393","metadata":{"_uuid":"934ef807-0cd6-458c-b8f0-1938831d3438","_cell_guid":"edd6dd19-7d60-4bd6-8171-501ed7724b12","collapsed":false,"papermill":{"duration":0.034152,"end_time":"2021-03-22T14:43:25.880882","exception":false,"start_time":"2021-03-22T14:43:25.84673","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# tokenizer\ntokenizer = CstTokenizer()\nstart_index = tokenizer.word_index[start]\nend_index = tokenizer.word_index[end]\ntokenizer.word_index","metadata":{"_uuid":"9d23088a-f930-4090-91c5-8740de9c1723","_cell_guid":"5c774bea-a00c-469e-bac2-f8af12a6305d","collapsed":false,"papermill":{"duration":0.044694,"end_time":"2021-03-22T14:43:25.951677","exception":false,"start_time":"2021-03-22T14:43:25.906983","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{"_uuid":"77d2bef1-edb6-4782-901a-6d92945850da","_cell_guid":"38b9d459-99a3-4c8e-baae-2212fdd00918","papermill":{"duration":0.027908,"end_time":"2021-03-22T14:43:26.070758","exception":false,"start_time":"2021-03-22T14:43:26.04285","status":"completed"},"tags":[],"trusted":true}},{"cell_type":"code","source":"class TrainDataset(tfds.core.GeneratorBasedBuilder):\n    VERSION = tfds.core.Version('0.1.0')\n    \n    def _split_generators(self, dl_manager):\n        return [\n            tfds.core.SplitGenerator(\n                    name=f'train',\n                    gen_kwargs={\n                    },\n            )\n        ]\n    \n    def _info(self):\n        return tfds.core.DatasetInfo(\n            builder=self,\n            description=(\"\"),\n            features=tfds.features.FeaturesDict({\n                \"image\": tfds.features.Image(shape=(None,None,1)),\n                \"target\": tfds.features.Tensor(shape=(max_seq,),dtype=tf.int8),\n            }),\n        )\n    \n    def _generate_examples(self,**args):\n        pass\n    \n\nclass TestDataset(tfds.core.GeneratorBasedBuilder):\n    VERSION = tfds.core.Version('0.1.0')\n    \n    def _split_generators(self, dl_manager):\n        return [\n            tfds.core.SplitGenerator(\n                    name=f'test',\n                    gen_kwargs={\n                    },\n            )\n        ]\n    \n    def _info(self):\n        return tfds.core.DatasetInfo(\n            builder=self,\n            description=(\"\"),\n            features=tfds.features.FeaturesDict({\n                \"image\": tfds.features.Image(shape=(None,None,1),),\n                \"image_id\": tfds.features.Text(),\n            }),\n        )\n    \n    def _generate_examples(self,**args):\n        pass","metadata":{"_uuid":"c94e276f-a942-47ad-b3c8-362a735c5bb5","_cell_guid":"3f9bb40b-e930-450d-83ff-57abc0a29b61","collapsed":false,"_kg_hide-input":true,"_kg_hide-output":true,"papermill":{"duration":0.038883,"end_time":"2021-03-22T14:43:26.136727","exception":false,"start_time":"2021-03-22T14:43:26.097844","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE_PER_REPLICA = 512\nGLOBAL_BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync\ntrain_steps = 2424186//GLOBAL_BATCH_SIZE\nBUFFER_SIZE = 10000\nTRAIN_IMAGE_MODEL = True\n\nprefetch = 50\nHEIGHT = 320\nWIDTH = 320","metadata":{"_uuid":"b7ecacfe-4498-4696-83ac-fd30f753ff76","_cell_guid":"c0202deb-8d97-41e4-acac-cb589c69f4ca","collapsed":false,"papermill":{"duration":0.034366,"end_time":"2021-03-22T14:43:26.197909","exception":false,"start_time":"2021-03-22T14:43:26.163543","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Feel free to change these parameters according to your system's configuration\nembedding_dim = 256\nvocab_size = len(tokenizer.index_word)+1\nattention_features_shape = 256\nrnn_units = 512","metadata":{"_uuid":"5d068df3-ae1a-42f6-9df5-3edbc306ff15","_cell_guid":"5e694eb5-7248-41a3-b364-f9aa32722aab","collapsed":false,"papermill":{"duration":0.034343,"end_time":"2021-03-22T14:43:26.259091","exception":false,"start_time":"2021-03-22T14:43:26.224748","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def data_augment(image):\n    p_rotation = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_noise = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_flip1 = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_flip2 = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n            \n    # Rotation\n    if p_rotation > .2:\n        image = rotation(image)\n        \n    # Flip\n    if p_flip1 > .4:\n        image = tf.image.random_flip_left_right(image, seed)\n        \n    # Flip\n    if p_flip2 > .4:\n        image = tf.image.random_flip_up_down(image, seed)\n        \n    # Resize \n    image = tf.image.resize(image,(WIDTH, HEIGHT))\n            \n    # Noise\n    if p_noise >= .4:\n        image = random_noise(image)\n        \n    return image\n\ndef rotation(img, rotation=0.2):\n    rotation = tf.random.uniform([], -1.0, 1.0, dtype=tf.float32)*rotation\n    shape = tf.shape(img)\n    h,w = shape[0],shape[1]\n    # Pad the image with zeros to avoid losing some pixels after rotation. \n    # This will double the image width and height\n    img = tf.image.pad_to_bounding_box(img,h//2, w//2,h*2, w*2)\n    img = tfa.image.rotate(img,rotation,fill_value=0)\n    # Now remove the zero pads\n    return remove_pad(img)\n\n\ndef remove_pad(arr,pad_value = 0.0):\n    arr_masked = tf.reduce_all(arr != pad_value , axis=-1)\n    #x\n    y = tf.argmax(arr_masked, axis=1)\n    y = tf.where(y)\n    y_min,y_max = y[0,0],y[-1,0]+1\n    #y\n    x = tf.argmax(arr_masked, axis=0)\n    x = tf.where(x)\n    x_min,x_max = x[0,0],x[-1,0]+1\n    arr = arr[y_min:y_max,x_min:x_max]\n    return arr\n\ndef random_noise(img,p=0.01):\n    shape = tf.shape(img)\n    choice = tf.random.categorical(tf.math.log([[p, 1-p]]), tf.size(img),dtype=tf.int32)\n    noise = tf.random.categorical(tf.math.log([[1., 1.]]), tf.size(img),dtype=tf.int32)\n    choice = tf.reshape(choice,shape)\n    noise = tf.reshape(noise,shape)\n    noise = tf.abs(choice-1)*noise\n    choice = tf.cast(choice,img.dtype)\n    noise = tf.cast(noise,img.dtype)\n    return (choice*img)+noise","metadata":{"_uuid":"b8d83835-471a-46da-87d5-8e22c664aba2","_cell_guid":"08e4fb14-7938-49e1-b822-f795c355119a","collapsed":false,"_kg_hide-input":true,"papermill":{"duration":0.042408,"end_time":"2021-03-22T14:43:26.328744","exception":false,"start_time":"2021-03-22T14:43:26.286336","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load the pretrained image model","metadata":{"_uuid":"56591733-4fdd-4d90-9440-3086d6685654","_cell_guid":"0adf7ab0-d243-4571-ae79-c79f2bbd4ebc","papermill":{"duration":0.026753,"end_time":"2021-03-22T14:43:26.38265","exception":false,"start_time":"2021-03-22T14:43:26.355897","status":"completed"},"tags":[],"trusted":true}},{"cell_type":"code","source":"with strategy.scope():\n    load_locally = tf.saved_model.LoadOptions(experimental_io_device='/job:localhost')\n    #image_model = tf.keras.models.load_model('../input/mt-pretraining/EfficientNetB0.h5',options=load_locally,compile=False)\n    image_input = tf.keras.layers.Input(shape=(WIDTH,HEIGHT,1))\n    image_model = tf.keras.applications.EfficientNetB1(include_top=False,weights=None,input_shape=(WIDTH,HEIGHT,1),)\n    image_model = image_model(image_input)\n    \n    image_model = tf.keras.layers.Reshape((attention_features_shape,-1))(image_model)\n    image_model = tf.keras.Model(image_input,image_model)\n    image_model.compile()\n    image_model.summary()","metadata":{"_uuid":"37ac148d-70b8-4b81-ab09-48fa670eb034","_cell_guid":"91d7ffea-eea1-4a6b-8a21-832f23f1fac8","papermill":{"duration":9.925843,"end_time":"2021-03-22T14:43:36.335817","exception":false,"start_time":"2021-03-22T14:43:26.409974","status":"completed"},"tags":[],"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_dataset(_):\n    builder = TrainDataset(data_dir=GCS_PATH)\n    # The following line download the dataset\n    builder.download_and_prepare()\n    dataset = builder.as_dataset()['train']\n\n    # normalize, shuffle and bacth\n    def preprecoss(x):\n        img,target = x['image'],x['target']\n        # Normalize : There are two pixels 0 and 255\n        img = tf.cast(img == 0,tf.float32)\n        \n        # label\n        target = tf.cast(target, tf.int32 )\n        return data_augment(img),target\n        \n        return data_augment(img),target\n    dataset = dataset.repeat().shuffle(BUFFER_SIZE).map(preprecoss,num_parallel_calls=AUTO)\n    dataset = dataset.batch(BATCH_SIZE_PER_REPLICA).prefetch(prefetch)\n    return dataset\n\nwith strategy.scope():\n    if tpu is None:\n        dataset = get_dataset(0)\n    else:\n        dataset = strategy.experimental_distribute_datasets_from_function(get_dataset)\n    train_iterator = iter(dataset)","metadata":{"_uuid":"bdbc6a26-24a5-4f77-89ad-b6bfb698c2df","_cell_guid":"2de8ee07-df00-4675-b614-634d243e253d","collapsed":false,"papermill":{"duration":5.613738,"end_time":"2021-03-22T14:43:41.979233","exception":false,"start_time":"2021-03-22T14:43:36.365495","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\ninputs = next(train_iterator)","metadata":{"_uuid":"50dd3f7f-b5c1-459f-a221-1b108606fb4e","_cell_guid":"e400140c-95e2-46da-8284-b7f5542b643d","collapsed":false,"papermill":{"duration":0.035854,"end_time":"2021-03-22T14:43:42.044863","exception":false,"start_time":"2021-03-22T14:43:42.009009","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{"_uuid":"0da68d82-c184-49aa-88f5-00447a22a991","_cell_guid":"edfca5e5-0bf8-4daf-9bde-7eca2d7f7e74","papermill":{"duration":0.028162,"end_time":"2021-03-22T14:43:42.102429","exception":false,"start_time":"2021-03-22T14:43:42.074267","status":"completed"},"tags":[],"trusted":true}},{"cell_type":"code","source":"class CNN_Encoder(tf.keras.Model):\n    # Since you have already extracted the features and dumped it using pickle\n    # This encoder passes those features through a Fully connected layer\n    def __init__(self, embedding_dim):\n        super(CNN_Encoder, self).__init__()\n        self.lstm_layer = tf.keras.layers.LSTM(rnn_units,\n                                   return_sequences=True,\n                                   return_state=True,\n                                   recurrent_initializer='glorot_uniform')\n\n\n    def call(self, x):\n        output, h, c = self.lstm_layer(x)\n        return output, h, c\n    \n    def get_config(self):\n        return {\n            'lstm_layer': self.lstm_layer,\n        }","metadata":{"_uuid":"7c8763c2-4b98-4613-860a-ad6b9876f927","_cell_guid":"dfd65b17-f3f1-4f3c-b698-8c6ffa6d4b82","collapsed":false,"papermill":{"duration":0.037612,"end_time":"2021-03-22T14:43:42.292575","exception":false,"start_time":"2021-03-22T14:43:42.254963","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class RNN_Decoder(tf.keras.Model):\n\n    def __init__(self, embedding_dim, vocab_size):\n        super(RNN_Decoder, self).__init__()\n\n        self.dec_units = rnn_units\n        self.attention_type = 'luong'\n\n        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)\n        \n        # Final Dense layer on which softmax will be applied\n        self.fc = tf.keras.Sequential([\n            tf.keras.layers.Dropout(0.5),\n            tf.keras.layers.Dense(vocab_size)\n        ])\n        \n        # Create attention mechanism with memory = None\n        memory_sequence_length = None# GLOBAL_BATCH_SIZE * [max_seq]\n        self.attention_mechanism = self.build_attention_mechanism(self.dec_units,None, memory_sequence_length,self.attention_type)\n\n        # Define the fundamental cell for decoder recurrent structure\n        self.decoder_rnn_cell = tf.keras.layers.LSTMCell(rnn_units)\n        # Wrap attention mechanism with the fundamental rnn cell of decoder\n        self.rnn_cell = self.build_rnn_cell()\n\n        # Sampler\n        self.sampler = tfa.seq2seq.sampler.TrainingSampler()\n\n        # Define the decoder with respect to fundamental rnn cell\n        self.decoder = tfa.seq2seq.BasicDecoder(self.rnn_cell, sampler=self.sampler, output_layer=self.fc,maximum_iterations=max_seq-1)\n\n\n    def call(self, inputs, initial_state):\n        x = self.embedding(inputs)\n        sequence_length = tf.repeat(max_seq - 1,tf.shape(x)[0])\n        outputs, _, _ = self.decoder(x, initial_state=initial_state,\n                                     sequence_length=sequence_length)\n        return outputs\n\n    def build_rnn_cell(self):\n        rnn_cell = tfa.seq2seq.AttentionWrapper(self.decoder_rnn_cell,\n                                                self.attention_mechanism, attention_layer_size=self.dec_units)\n        return rnn_cell\n\n    def build_attention_mechanism(self, dec_units, memory, memory_sequence_length=None, attention_type='luong'):\n        # ------------- #\n        # typ: Which sort of attention (Bahdanau, Luong)\n        # dec_units: final dimension of attention outputs \n        # memory: encoder hidden states of shape (batch_size, max_length_input, enc_units)\n        # memory_sequence_length: 1d array of shape (batch_size) with every element set to max_length_input (for masking purpose)\n\n        if (attention_type == 'bahdanau'):\n            return tfa.seq2seq.BahdanauAttention(units=dec_units, memory=memory,\n                                                 memory_sequence_length=memory_sequence_length)\n        else:\n            return tfa.seq2seq.LuongAttention(units=dec_units, memory=memory,\n                                              memory_sequence_length=memory_sequence_length)\n\n    def build_initial_state(self, batch_sz,encoder_state):\n        decoder_initial_state = self.rnn_cell.get_initial_state(batch_size=batch_sz, dtype=tf.float32)\n        decoder_initial_state = decoder_initial_state.clone(cell_state=encoder_state)\n        return decoder_initial_state\n\n    def reset_state(self, batch_size):\n        return tf.zeros((batch_size, rnn_units))\n\n    def get_config(self):\n        return {\n            'units': self.units,\n            'embedding': self.embedding,\n            'rnn': self.rnn,\n            'fc': self.fc,\n            'attention': self.attention, }\n","metadata":{"_uuid":"6995a60a-d922-4efb-bf80-bb64888cdefd","_cell_guid":"c4327769-696a-4f8a-ac21-dd14d19427e4","collapsed":false,"papermill":{"duration":0.04087,"end_time":"2021-03-22T14:43:42.362881","exception":false,"start_time":"2021-03-22T14:43:42.322011","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    optimizer = tf.keras.optimizers.Adam(0.0025)\n    crossentropy = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)\n\n    #@tf.function\n    def loss_function(real, pred):\n        mask = tf.math.logical_not(tf.math.equal(real, 0))\n        loss_ = crossentropy(real, pred)\n        \n        mask = tf.cast(mask, dtype=loss_.dtype)\n        loss_ *= mask\n      \n        return tf.nn.compute_average_loss(loss_, global_batch_size=GLOBAL_BATCH_SIZE)\n        #return tf.reduce_mean(loss_)","metadata":{"_uuid":"958f6998-1557-4254-a87f-4e265a3130e4","_cell_guid":"d36e92ba-a37c-49a9-8162-465007fedd5a","collapsed":false,"papermill":{"duration":0.041475,"end_time":"2021-03-22T14:43:42.433213","exception":false,"start_time":"2021-03-22T14:43:42.391738","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create Encoder","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    bs = 2\n    encoder = CNN_Encoder(embedding_dim,)\n    sample_output = image_model(tf.zeros((bs,WIDTH,HEIGHT,1),tf.float32))\n    sample_output, sample_h, sample_c = encoder(sample_output)\n    print(\"Encoder Outputs Shape: \", sample_output.shape)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create Decoder","metadata":{}},{"cell_type":"code","source":"with strategy.scope():\n    decoder = RNN_Decoder(embedding_dim, vocab_size)\n    sample_x = tf.random.uniform((bs, max_seq))\n    decoder.attention_mechanism.setup_memory(sample_output)\n    initial_state = decoder.build_initial_state(bs,[sample_h, sample_c],)\n\n    sample_decoder_outputs = decoder(sample_x, initial_state)\n\n    print(\"Decoder Outputs Shape: \", sample_decoder_outputs.rnn_output.shape)\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Checkpoint","metadata":{"_uuid":"14234cd3-4992-45cc-ae1c-8d8c3a75aa0a","_cell_guid":"15e35e06-94b9-4a37-94a2-4a1f9b438b56","papermill":{"duration":0.029148,"end_time":"2021-03-22T14:43:43.856845","exception":false,"start_time":"2021-03-22T14:43:43.827697","status":"completed"},"tags":[],"trusted":true}},{"cell_type":"code","source":"checkpoint_name = \"./model\"\n!cp -r ../input/mt-fast-distributed-training-tpu/*.dill .\n!cp -r ../input/mt-fast-distributed-training-tpu/*.h5 .","metadata":{"_uuid":"39f2b4dd-4bbb-4c99-97a7-679e8e88cbe1","_cell_guid":"084e8749-b53e-409e-a4fc-e7d685f9e1da","collapsed":false,"papermill":{"duration":0.037655,"end_time":"2021-03-22T14:43:43.923219","exception":false,"start_time":"2021-03-22T14:43:43.885564","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def save_model(name=checkpoint_name):\n    encoder.save_weights(f'{name}_encoder.h5', options=save_locally)\n    decoder.save_weights(f'{name}_decoder.h5', options=save_locally)\n    image_model.save_weights('image_model.h5', options=save_locally)\n    \ndef load_model(name=checkpoint_name):\n    encoder.load_weights(f'{name}_encoder.h5', options=load_locally)\n    decoder.load_weights(f'{name}_decoder.h5', options=load_locally)\n    image_model.load_weights('image_model.h5', options=load_locally)\n    return encoder, decoder,image_model","metadata":{"_uuid":"1f331bdb-bd2f-4f95-9a4e-6ccec79e1c33","_cell_guid":"0f41b370-bcda-4546-9604-3784fdc61685","collapsed":false,"_kg_hide-output":true,"papermill":{"duration":0.037667,"end_time":"2021-03-22T14:43:43.989838","exception":false,"start_time":"2021-03-22T14:43:43.952171","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    save_locally = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')\n    load_locally = tf.saved_model.LoadOptions(experimental_io_device='/job:localhost')\n    encoder, decoder, image_model = load_model()","metadata":{"_uuid":"58003073-26d4-4d42-9244-9ac51614e6c3","_cell_guid":"b7046f33-719c-42df-8367-b09f1bc7ab5c","papermill":{"duration":0.038456,"end_time":"2021-03-22T14:43:44.057157","exception":false,"start_time":"2021-03-22T14:43:44.018701","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@tf.function(experimental_relax_shapes=True)\ndef train_step(inputs):\n    img_tensor, target = inputs\n    loss = 0.\n    \n    batch_sz = tf.shape(img_tensor)[0]\n    \n    #image features\n    if not TRAIN_IMAGE_MODEL:\n        img_tensor = image_model(img_tensor,training=False)\n\n    with tf.GradientTape() as tape:\n        \n        if TRAIN_IMAGE_MODEL:\n                img_tensor = image_model(img_tensor,training=True)\n                \n        enc_output, enc_h, enc_c = encoder(img_tensor)\n\n        dec_input = target[:, :-1]  # Ignore <end> token\n        real = target[:, 1:]  # ignore <start> token\n\n        # Set the AttentionMechanism object with encoder_outputs\n        decoder.attention_mechanism.setup_memory(enc_output)\n\n        # Create AttentionWrapperState as initial_state for decoder\n        decoder_initial_state = decoder.build_initial_state(batch_sz, [enc_h, enc_c])\n        pred = decoder(dec_input, decoder_initial_state)\n        logits = pred.rnn_output\n        loss = loss_function(real, logits)\n\n    variables = encoder.trainable_variables + decoder.trainable_variables\n    gradients = tape.gradient(loss, variables)\n    optimizer.apply_gradients(zip(gradients, variables))\n\n    return loss\n\n@tf.function(experimental_relax_shapes=True)\ndef distributed_train_step(inputs,):\n    per_replica_losses = strategy.run(train_step, args=(inputs,))\n    return  strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses,axis=None)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"loss_plot = []\nif os.path.exists('losses.dill'):\n    loss_plot = dill.load(open('losses.dill','rb'))","metadata":{"_uuid":"fd4b5ad8-585b-48b3-ba3e-1f2eca5dbb9e","_cell_guid":"af277834-860e-4625-ab31-3ebd1fa9bd01","collapsed":false,"_kg_hide-input":true,"_kg_hide-output":true,"papermill":{"duration":0.036643,"end_time":"2021-03-22T14:43:44.250908","exception":false,"start_time":"2021-03-22T14:43:44.214265","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"EPOCHS = 10\nbest_loss = 0\nfor epoch in range(EPOCHS):\n    start = time.time()\n    total_loss = 0.0\n    print(f'Epoch {epoch + 1}/{EPOCHS}')\n    for batch in range(train_steps):\n        inputs = next(train_iterator)\n        loss = distributed_train_step(inputs).numpy()\n        total_loss += loss\n            \n        if batch % 1==0:\n            print(f'{batch+1}/{train_steps} Total Loss: {total_loss/(batch+1):.4f}  Batch Loss: {loss:.4f}',end='\\r')\n            \n    # storing the epoch end loss value to plot later\n    loss_plot.append(total_loss / train_steps)\n    print()\n    if best_loss >= total_loss:\n        print(\"Saving...\")\n        save_model()\n        best_loss = total_loss\n\n    print(f'Epoch {epoch + 1} Loss {total_loss/train_steps:.6f}')\n    print(f'Time taken for 1 epoch {time.time() - start} sec\\n')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"save_model()","metadata":{"_uuid":"13593346-69a4-46c4-bda5-6cf9a338b92e","_cell_guid":"1d42258e-d6eb-4080-9dce-92601730c708","collapsed":false,"papermill":{"duration":0.195995,"end_time":"2021-03-22T14:43:44.547128","exception":false,"start_time":"2021-03-22T14:43:44.351133","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.plot(loss_plot)\nplt.xlabel('Epochs')\nplt.ylabel('Loss')\nplt.title('Loss Plot')\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dill.dump(loss_plot,open('losses.dill','wb'))","metadata":{"_uuid":"8d82cc81-cde2-4600-a335-5b17c9fb1610","_cell_guid":"af0517c1-a79e-41eb-9316-84d9a1b9ee4b","collapsed":false,"_kg_hide-input":true,"_kg_hide-output":true,"papermill":{"duration":0.037931,"end_time":"2021-03-22T14:43:44.823708","exception":false,"start_time":"2021-03-22T14:43:44.785777","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Testing","metadata":{"_uuid":"d5222082-a6b3-428c-89cf-d720c0ca9f7f","_cell_guid":"f7e1badc-818b-4146-85e6-4d240b175c5c","papermill":{"duration":0.029819,"end_time":"2021-03-22T14:43:44.883856","exception":false,"start_time":"2021-03-22T14:43:44.854037","status":"completed"},"tags":[],"trusted":true}},{"cell_type":"code","source":"# Get test dataset\ndef get_test_dataset(_):\n    builder = TestDataset(data_dir=GCS_PATH)\n    # The following line download the dataset\n    builder.download_and_prepare()\n    dataset = builder.as_dataset()['test']\n\n    # normalize, shuffle and bacth\n    def preprecoss(x):\n        img = x['image']\n        \n        #Normalize\n        img = tf.cast(img == 0,tf.float32)\n        im_size = tf.shape(img)\n        w,h = im_size[0],im_size[1]\n        if h > w:\n            img = tf.image.transpose(img)\n            img = tf.image.flip_up_down(img)\n        \n        img = tf.image.resize(img,(WIDTH, HEIGHT))\n        return img,x['image_id']\n    \n    dataset = dataset.map(preprecoss,num_parallel_calls=AUTO)\n    dataset = dataset.batch(BATCH_SIZE_PER_REPLICA)\n    dataset = dataset.prefetch(prefetch)\n    return dataset\n\nwith strategy.scope():\n    test_dataset = strategy.experimental_distribute_datasets_from_function(get_test_dataset)\n    #test_iterator = iter(test_dataset)","metadata":{"_uuid":"625f21f6-6f49-45c5-ba39-c5291bc5da29","_cell_guid":"7c6346a8-7215-4080-8dd3-de4df58178be","collapsed":false,"papermill":{"duration":1.367738,"end_time":"2021-03-22T14:43:49.196748","exception":false,"start_time":"2021-03-22T14:43:47.82901","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@tf.function(experimental_relax_shapes=True)\ndef greedy_predict(img_tensor):\n    inference_batch_size = img_tensor.shape[0]\n    img_tensor = image_model(img_tensor,training=False)\n    \n    enc_out, enc_h, enc_c = encoder(img_tensor)\n    dec_h = enc_h\n    dec_c = enc_c\n    \n    start_tokens = tf.fill([inference_batch_size], start_index)\n    end_token = end_index\n    \n    greedy_sampler = tfa.seq2seq.GreedyEmbeddingSampler()\n    \n    # Instantiate BasicDecoder object\n    decoder_instance = tfa.seq2seq.BasicDecoder(cell=decoder.rnn_cell, sampler=greedy_sampler, output_layer=decoder.fc,maximum_iterations=max_seq)\n    # Setup Memory in decoder stack\n    decoder.attention_mechanism.setup_memory(enc_out)\n\n    # set decoder_initial_state\n    decoder_initial_state = decoder.build_initial_state(inference_batch_size, [enc_h, enc_c])\n\n    ### Since the BasicDecoder wraps around Decoder's rnn cell only, you have to ensure that the inputs to BasicDecoder \n    ### decoding step is output of embedding layer. tfa.seq2seq.GreedyEmbeddingSampler() takes care of this. \n    ### You only need to get the weights of embedding layer, which can be done by decoder.embedding.variables[0] and pass this callabble to BasicDecoder's call() function\n\n    decoder_embedding_matrix = decoder.embedding.variables[0]\n\n    outputs, _, _ = decoder_instance(decoder_embedding_matrix, start_tokens = start_tokens, end_token= end_token, initial_state=decoder_initial_state)\n    return outputs.sample_id#.numpy()","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@tf.function(experimental_relax_shapes=True)\ndef beam_predict(img_tensor,beam_width=3):\n    inference_batch_size = img_tensor.shape[0]\n    img_tensor = image_model(img_tensor,training=False)\n    \n    enc_out, enc_h, enc_c = encoder(img_tensor)\n\n    dec_h = enc_h\n    dec_c = enc_c\n\n    start_tokens = tf.fill([inference_batch_size], start_index)\n    end_token = end_index\n\n    # From official documentation\n    # NOTE If you are using the BeamSearchDecoder with a cell wrapped in AttentionWrapper, then you must ensure that:\n    # The encoder output has been tiled to beam_width via tfa.seq2seq.tile_batch (NOT tf.tile).\n    # The batch_size argument passed to the get_initial_state method of this wrapper is equal to true_batch_size * beam_width.\n    # The initial state created with get_initial_state above contains a cell_state value containing properly tiled final state from the encoder.\n\n    enc_out = tfa.seq2seq.tile_batch(enc_out, multiplier=beam_width)\n    decoder.attention_mechanism.setup_memory(enc_out)\n    print(\"beam_with * [batch_size, max_length_input, rnn_units] :  3 * [1, 16, 1024]] :\", enc_out.shape)\n\n    # set decoder_inital_state which is an AttentionWrapperState considering beam_width\n    hidden_state = tfa.seq2seq.tile_batch([enc_h, enc_c], multiplier=beam_width)\n    print(beam_width*inference_batch_size)\n    decoder_initial_state = decoder.rnn_cell.get_initial_state(batch_size=beam_width*inference_batch_size,dtype=tf.float32)\n    decoder_initial_state = decoder_initial_state.clone(cell_state=hidden_state)\n\n    # Instantiate BeamSearchDecoder\n    decoder_instance = tfa.seq2seq.BeamSearchDecoder(decoder.rnn_cell,beam_width=beam_width, output_layer=decoder.fc,maximum_iterations=max_seq)\n    decoder_embedding_matrix = decoder.embedding.variables[0]\n\n    # The BeamSearchDecoder object's call() function takes care of everything.\n    outputs, final_state, sequence_lengths = decoder_instance(decoder_embedding_matrix, start_tokens=start_tokens, end_token=end_token, initial_state=decoder_initial_state)\n    # outputs is tfa.seq2seq.FinalBeamSearchDecoderOutput object. \n    # The final beam predictions are stored in outputs.predicted_id\n    # outputs.beam_search_decoder_output is a tfa.seq2seq.BeamSearchDecoderOutput object which keep tracks of beam_scores and parent_ids while performing a beam decoding step\n    # final_state = tfa.seq2seq.BeamSearchDecoderState object.\n    # Sequence Length = [inference_batch_size, beam_width] details the maximum length of the beams that are generated\n\n\n    # outputs.predicted_id.shape = (inference_batch_size, time_step_outputs, beam_width)\n    # outputs.beam_search_decoder_output.scores.shape = (inference_batch_size, time_step_outputs, beam_width)\n    # Convert the shape of outputs and beam_scores to (inference_batch_size, beam_width, time_step_outputs)\n    final_outputs = tf.transpose(outputs.predicted_ids, perm=(0,2,1))\n    beam_scores = tf.transpose(outputs.beam_search_decoder_output.scores, perm=(0,2,1))\n\n    return final_outputs, beam_scores","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@tf.function(experimental_relax_shapes=True)\ndef distributed_greedy_predict(inputs):\n    return strategy.run(greedy_predict, args=(inputs,))\n\n\n@tf.function(experimental_relax_shapes=True)\ndef distributed_beam_predict(inputs):\n    return strategy.run(beam_predict, args=(inputs,))","metadata":{"_uuid":"1bdcb2d5-cc57-49a5-afd3-1dbb284cc3df","_cell_guid":"8676b3dd-c93a-4ef8-adee-41ddad0bba3f","collapsed":false,"papermill":{"duration":0.042847,"end_time":"2021-03-22T14:43:50.096673","exception":false,"start_time":"2021-03-22T14:43:50.053826","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def process_predict(x):\n    x = x.numpy()\n    \n    res = []\n    for i in range(len(x)):\n        tmp = x[i]\n        idx = np.argwhere(tmp == end_index)\n        # Remove sos, eos and pad \n        if idx.size != 0:\n            idx = idx[0,0]\n            tmp = tmp[:idx]\n            \n        tmp = tmp[tmp!=0]\n        tmp = tmp[tmp!=start_index]\n        tmp = tmp[tmp!=end_index]\n        res.append(tmp)\n    captions = tokenizer.detokenize(res)\n    captions = [f'InChI=1S/{c}'  for c in captions]\n    return captions","metadata":{"_uuid":"6f8f8a35-f626-4832-90e9-57a17c083d90","_cell_guid":"ce08bfcb-5a29-4af1-963c-d0258aaae524","collapsed":false,"papermill":{"duration":0.040931,"end_time":"2021-03-22T14:43:50.239386","exception":false,"start_time":"2021-03-22T14:43:50.198455","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submit","metadata":{"_uuid":"20dea646-cb5f-430d-8c8e-6a43ae27cce7","_cell_guid":"6302109c-e772-4523-a723-0d1024769d11","papermill":{"duration":0.031004,"end_time":"2021-03-22T14:43:50.30245","exception":false,"start_time":"2021-03-22T14:43:50.271446","status":"completed"},"tags":[],"trusted":true}},{"cell_type":"code","source":"%%skip not SUBMIT\nresults = []\nimages = []\n\nfor (image,iids) in tqdm(test_dataset,total=1616107//GLOBAL_BATCH_SIZE):\n    if type(image) == value_lib.PerReplica:\n        stop = False\n        for v in image.values:\n            if v.shape[0] == 0:\n                stop = True\n        if stop:\n            print('Stopping...')\n            break\n        res = distributed_greedy_predict(image)\n        res = tf.concat(strategy.unwrap(res),axis=0)\n        iids = tf.concat(strategy.unwrap(iids),axis=0)\n    else:\n        res = greedy_predict(image)\n    \n    res = process_predict(res)\n    results.extend(res)\n    images.extend([s.decode('utf-8') for s in iids.numpy()])\n\n    \n# Process the last batch\nif type(image) == value_lib.PerReplica:\n    image_values,iids_values =  image.values,iids.values\n\n    for i in range(strategy.num_replicas_in_sync):\n        if image_values[i].shape[0] == 0:\n            continue\n        res = greedy_predict(image_values[i])\n        res = process_predict(res)\n        results.extend(res)\n        images.extend([s.decode('utf-8') for s in iids_values[i].numpy()])","metadata":{"_uuid":"6b807982-fde2-47a8-9e67-fc29501d9526","_cell_guid":"c45f6f7e-7ce7-4453-81a4-bfcf8070d134","papermill":{"duration":7870.577112,"end_time":"2021-03-22T16:55:00.911071","exception":false,"start_time":"2021-03-22T14:43:50.333959","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%skip not SUBMIT\nsubmission_df = pd.read_csv('../input/bms-molecular-translation/sample_submission.csv',index_col=0)\nsubmission_df.loc[images,'InChI']=results\nsubmission_df.to_csv('submission.csv')","metadata":{"_uuid":"60d71920-828a-4064-8bf8-2ace027bdb2c","_cell_guid":"a7114e1a-465f-4177-ac4a-695923fa4e4d","collapsed":false,"papermill":{"duration":26.904573,"end_time":"2021-03-22T16:55:27.903903","exception":false,"start_time":"2021-03-22T16:55:00.99933","status":"completed"},"tags":[],"jupyter":{"outputs_hidden":false},"trusted":true},"execution_count":null,"outputs":[]}]}