{"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":"# Version\nversion4 - Modified to run on kaggle kernel without gcs\n\n# Reference\n\nTokenization part is from these great starter kit.  \nhttps://www.kaggle.com/yasufuminakama/inchi-resnet-lstm-with-attention-starter\nhttps://www.kaggle.com/yasufuminakama/inchi-resnet-lstm-with-attention-inference\nhttps://www.kaggle.com/yasufuminakama/inchi-preprocess-2\n\nTPU training pipeline is from this kernel, it saves me a lot.  \nhttps://www.kaggle.com/markwijkhuizen/tensorflow-tpu-training-baseline-lb-16-92\n\nAnd transformer part is from here.  \nhttps://www.kaggle.com/aditya08/imagecaptioning-show-attnd-tell-w-transformer","metadata":{}},{"cell_type":"markdown","source":"# LIBRARIES","metadata":{"execution":{"iopub.status.busy":"2021-06-05T04:37:25.660261Z","iopub.execute_input":"2021-06-05T04:37:25.660679Z","iopub.status.idle":"2021-06-05T04:37:25.664993Z","shell.execute_reply.started":"2021-06-05T04:37:25.660596Z","shell.execute_reply":"2021-06-05T04:37:25.664282Z"}}},{"cell_type":"code","source":"!pip install -q efficientnet tensorflow_addons Levenshtein","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-06-09T13:33:53.149786Z","iopub.execute_input":"2021-06-09T13:33:53.150173Z","iopub.status.idle":"2021-06-09T13:34:02.034987Z","shell.execute_reply.started":"2021-06-09T13:33:53.150094Z","shell.execute_reply":"2021-06-09T13:34:02.033936Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import io\nimport json\nimport math\nimport os\nimport random\nimport re\nimport time\nfrom dataclasses import dataclass\nfrom typing import List, Optional, Union\n\nimport efficientnet.tfkeras as efn\nimport Levenshtein\nimport numpy as np\nimport tensorflow as tf\nimport tensorflow_addons as tfa\nfrom tensorflow_addons import optimizers as tfa_optimizers\nfrom tensorflow.keras.mixed_precision import experimental as mixed_precision\nfrom tqdm.notebook import tqdm\n\nfrom kaggle_datasets import KaggleDatasets\n\n# seed everything\nSEED = 42\nos.environ['PYTHONHASHSEED'] = str(SEED)\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntf.random.set_seed(SEED)\n\n\nGCS_PATHS = {\n    '416x736_no_pad': KaggleDatasets().get_gcs_path('tfrecords018'),\n    '416x736_no_pad_test': KaggleDatasets().get_gcs_path('tfrecords016'),\n}\nYOUR_GCS_DIR = None","metadata":{"execution":{"iopub.status.busy":"2021-06-09T13:34:02.037730Z","iopub.execute_input":"2021-06-09T13:34:02.038055Z","iopub.status.idle":"2021-06-09T13:34:09.110688Z","shell.execute_reply.started":"2021-06-09T13:34:02.038017Z","shell.execute_reply":"2021-06-09T13:34:09.109682Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# DEVICE SETUP","metadata":{}},{"cell_type":"code","source":"try:\n    TPU = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print('Running on TPU ', TPU.master())\nexcept ValueError:\n    print('Running on GPU')\n    TPU = None\n\nif TPU:\n    tf.config.experimental_connect_to_cluster(TPU)\n    tf.tpu.experimental.initialize_tpu_system(TPU)\n    strategy = tf.distribute.experimental.TPUStrategy(TPU)\nelse:\n    strategy = tf.distribute.get_strategy()\n\nREPLICAS = strategy.num_replicas_in_sync\nprint(f'NUM REPLICAS: {REPLICAS}')\n\nmixed_precision.set_policy('mixed_bfloat16' if TPU else 'float32')\ntf.config.optimizer.set_jit(True)\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-09T13:34:09.113030Z","iopub.execute_input":"2021-06-09T13:34:09.113506Z","iopub.status.idle":"2021-06-09T13:34:14.869619Z","shell.execute_reply.started":"2021-06-09T13:34:09.113440Z","shell.execute_reply":"2021-06-09T13:34:14.868589Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# CONFIG","metadata":{}},{"cell_type":"code","source":"@dataclass\nclass Config:\n    # general\n    exp_id = 'exp001'\n    save_dir: str = f'{YOUR_GCS_DIR}/{exp_id}'\n    debug: bool = False\n    inference: bool = False\n    resume: bool = False\n\n    # solver\n    steps_per_epoch: int = 40000\n    batch_size_base: int = 16\n    batch_size: int = batch_size_base * REPLICAS\n    test_batch_size_base: int = 16\n    test_batch_size: int = batch_size_base * REPLICAS\n    num_epochs: int = 50\n    eval_freq: int = 5\n    warmup_steps: int = 500\n    verbose_freq: int = 100\n    save_freq: int = 5000\n    total_steps: int = num_epochs * steps_per_epoch\n\n    # data\n    image_height: int = 416\n    image_width: int = 736\n    row_size: int = (image_height - 1) // 32 + 1\n    col_size: int = (image_width - 1) // 32 + 1\n    vocab_size: int = 193\n    seq_len: int = 200\n    dtype: str = tf.bfloat16 if TPU else tf.float32\n\n    train_gcs_dir: str = GCS_PATHS['416x736_no_pad']\n    val_gcs_dir: str = GCS_PATHS['416x736_no_pad']\n    test_gcs_dir: str = GCS_PATHS['416x736_no_pad_test']\n    val_size: int = 121210\n    val_steps: int = val_size // batch_size\n    rotate_angle: int = 5\n    zoom_range: float = 0.1\n\n    # configure model\n    encoder_dim: int = 1792\n    start_token: int = 191\n    end_token: int = 192\n    pad_token: int = 0\n    num_layers: int = 2\n    d_model: int = 512\n    num_heads: int = 8\n    dff: int = 2048\n    encoder_drop_rate: float = 0.1\n    decoder_drop_rate: float = 0.1\n\n\nCFG = Config()\nos.makedirs(CFG.save_dir, exist_ok=True)\n\nPAD_TOKEN = tf.constant(CFG.pad_token, dtype=tf.int64)\nSTART_TOKEN = tf.constant(CFG.start_token, dtype=tf.int64)\nEND_TOKEN = tf.constant(CFG.end_token, dtype=tf.int64)\n\nif CFG.debug:\n    CFG.steps_per_epoch = 10\n    CFG.val_steps = 10\n    CFG.num_epochs = 5","metadata":{"execution":{"iopub.status.busy":"2021-06-09T13:34:14.873222Z","iopub.execute_input":"2021-06-09T13:34:14.873538Z","iopub.status.idle":"2021-06-09T13:34:14.892118Z","shell.execute_reply.started":"2021-06-09T13:34:14.873509Z","shell.execute_reply":"2021-06-09T13:34:14.890840Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# DATASET","metadata":{"execution":{"iopub.status.busy":"2021-06-05T04:38:08.680147Z","iopub.execute_input":"2021-06-05T04:38:08.680473Z","iopub.status.idle":"2021-06-05T04:38:08.683585Z","shell.execute_reply.started":"2021-06-05T04:38:08.680447Z","shell.execute_reply":"2021-06-05T04:38:08.68295Z"}}},{"cell_type":"code","source":"def random_rotate(img, angle=CFG.rotate_angle):\n    angle *= np.pi/180\n    angle = tf.random.uniform(\n        shape=[CFG.batch_size], minval=-angle, maxval=angle)\n    return tfa.image.rotate(img, angle, fill_value=255)\n\n\ndef read_tfrecord(example):\n    tfrec_format = {\n        'image': tf.io.FixedLenFeature([], tf.string),\n        'image_id': tf.io.FixedLenFeature([], tf.string),\n        'label': tf.io.FixedLenFeature([], tf.string)\n    }\n    example = tf.io.parse_single_example(example, tfrec_format)\n    img = tf.image.decode_png(example['image'])\n    img = tf.reshape(img, (CFG.image_height, CFG.image_width, 3))\n    label = tf.io.decode_raw(example['label'], tf.int64)\n    label = tf.reshape(label, (277,))\n    label = label[:CFG.seq_len]\n    label = (label + 1) % 193 # to set pad token as 0\n    return img, label\n\n\ndef read_test_tfrecord(example):\n    tfrec_format = {\n        'image': tf.io.FixedLenFeature([], tf.string),\n        'image_id': tf.io.FixedLenFeature([], tf.string),\n    }\n    example = tf.io.parse_single_example(example, tfrec_format)\n    img = tf.image.decode_png(example['image'])\n    img = tf.reshape(img, (CFG.image_height, CFG.image_width, 3))\n    image_id = example['image_id']\n    return img, image_id\n\n\ndef imagenet_normalize(img, labels):\n    IMAGENET_MEAN = tf.constant([0.485, 0.456, 0.406], dtype=tf.float32)\n    IMAGENET_STD = tf.constant([0.229, 0.224, 0.225], dtype=tf.float32)\n    img = tf.cast(img, tf.float32) / 255.0\n    img = (img - IMAGENET_MEAN) / IMAGENET_STD\n    img = tf.cast(img, CFG.dtype)\n    return img, labels\n\n\ndef get_dataset(mode, batch_size=64, data_root='./', fold=0):\n    lengths = [\n        121210, 121210, 121210, 121210, 121210, 121210,\n        121209, 121209, 121209, 121209, 121209, 121209,\n        121209, 121209, 121209, 121209, 121209, 121209,\n        121209, 121209\n    ]\n\n    if mode == 'train':\n        files = tf.io.gfile.glob(f'{CFG.train_gcs_dir}/*.tfrec')\n        files = [f for f in files if f'fold{fold}' not in f]\n        length = sum([length for i, length in enumerate(lengths) if i != fold])\n    else:\n        files = tf.io.gfile.glob(f'{CFG.val_gcs_dir}/*.tfrec')\n        files = [f for f in files if f'fold{fold}' in f]\n        length = lengths[fold]\n\n    AUTO = tf.data.experimental.AUTOTUNE\n    ds = tf.data.TFRecordDataset(files, num_parallel_reads=AUTO)\n    ds = ds.prefetch(AUTO)\n    ds = ds.map(read_tfrecord, num_parallel_calls=AUTO)\n\n    if mode == 'train':\n        ignore_order = tf.data.Options()\n        ignore_order.experimental_deterministic = False\n\n        ds = ds.batch(batch_size, drop_remainder=True)\n        zoom = tf.keras.layers.experimental.preprocessing.RandomZoom(\n            (-CFG.zoom_range, CFG.zoom_range))\n        ds = ds.map(lambda x, y: (zoom(x), y),\n                    num_parallel_calls=AUTO)\n        ds = ds.map(lambda x, y: (random_rotate(x), y),\n                    num_parallel_calls=AUTO)\n\n        ds = ds.map(imagenet_normalize, num_parallel_calls=AUTO)\n        ds = ds.with_options(ignore_order)\n        ds = ds.shuffle(512, reshuffle_each_iteration=True)\n        ds = ds.repeat()\n    else:\n        ds = ds.batch(batch_size, drop_remainder=True)\n        ds = ds.map(imagenet_normalize, num_parallel_calls=AUTO)\n    ds = ds.prefetch(1)\n    return ds, length\n\n\ndef get_test_dataset(batch_size=64):\n    length = 1616107\n    files = tf.io.gfile.glob(f'{CFG.test_gcs_dir}/*.tfrec')\n    \n    AUTO = tf.data.experimental.AUTOTUNE\n    ds = tf.data.TFRecordDataset(files, num_parallel_reads=AUTO)\n    ds = ds.prefetch(AUTO)\n    ds = ds.repeat()\n    ds = ds.map(read_test_tfrecord, num_parallel_calls=AUTO)\n    ds = ds.batch(batch_size, drop_remainder=False)\n    ds = ds.map(imagenet_normalize, num_parallel_calls=AUTO)\n    ds = ds.prefetch(1)\n    return ds, length","metadata":{"execution":{"iopub.status.busy":"2021-06-09T13:34:14.893435Z","iopub.execute_input":"2021-06-09T13:34:14.893825Z","iopub.status.idle":"2021-06-09T13:34:14.924746Z","shell.execute_reply.started":"2021-06-09T13:34:14.893794Z","shell.execute_reply":"2021-06-09T13:34:14.923696Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" # MODEL","metadata":{}},{"cell_type":"code","source":"def get_angles(pos, i, d_model):\n    angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))\n    return pos * angle_rates\n\n\ndef positional_encoding_1d(position, d_model):\n    angle_rads = get_angles(np.arange(position)[:, np.newaxis],\n                            np.arange(d_model)[np.newaxis, :],\n                            d_model)\n\n    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n    pos_encoding = angle_rads[np.newaxis, ...]\n    return tf.cast(pos_encoding, dtype=CFG.dtype)\n\n\ndef positional_encoding_2d(row, col, d_model):\n    assert d_model % 2 == 0\n    row_pos = np.repeat(np.arange(row), col)[:, np.newaxis]\n    col_pos = np.repeat(np.expand_dims(np.arange(col), 0),\n                        row, axis=0).reshape(-1, 1)\n\n    angle_rads_row = get_angles(row_pos, np.arange(\n        d_model//2)[np.newaxis, :], d_model//2)\n    angle_rads_col = get_angles(col_pos, np.arange(\n        d_model//2)[np.newaxis, :], d_model//2)\n\n    angle_rads_row[:, 0::2] = np.sin(angle_rads_row[:, 0::2])\n    angle_rads_row[:, 1::2] = np.cos(angle_rads_row[:, 1::2])\n    angle_rads_col[:, 0::2] = np.sin(angle_rads_col[:, 0::2])\n    angle_rads_col[:, 1::2] = np.cos(angle_rads_col[:, 1::2])\n    pos_encoding = np.concatenate([angle_rads_row, angle_rads_col], axis=1)[\n        np.newaxis, ...]\n    return tf.cast(pos_encoding, dtype=CFG.dtype)\n\n\ndef create_padding_mask(seq):\n    seq = tf.cast(tf.math.equal(seq, 0), CFG.dtype)\n    return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)\n\n\ndef create_look_ahead_mask(size):\n    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n    mask = tf.cast(mask, CFG.dtype)\n    return mask  # (seq_len, seq_len)\n\n\ndef scaled_dot_product_attention(q, k, v, mask):\n    # (..., seq_len_q, seq_len_k)\n    matmul_qk = tf.matmul(q, k, transpose_b=True)\n    dk = tf.cast(tf.shape(k)[-1], CFG.dtype)\n    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)\n\n    if mask is not None:\n        scaled_attention_logits += (mask * -1e9)\n\n    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)\n    output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)\n\n    return output, attention_weights\n\n\ndef create_masks_decoder(tar):\n    look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])\n    dec_target_padding_mask = create_padding_mask(tar)\n    combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)\n    return combined_mask\n\n\ndef point_wise_feed_forward_network(d_model, dff):\n    return tf.keras.Sequential([\n        tf.keras.layers.Dense(dff, activation='relu'),\n        tf.keras.layers.Dense(d_model)])\n\n\nclass MultiHeadAttention(tf.keras.layers.Layer):\n    def __init__(self, d_model, num_heads):\n        super(MultiHeadAttention, self).__init__()\n        self.num_heads = num_heads\n        self.d_model = d_model\n        assert d_model % self.num_heads == 0\n        self.depth = d_model // self.num_heads\n        self.wq = tf.keras.layers.Dense(d_model)\n        self.wk = tf.keras.layers.Dense(d_model)\n        self.wv = tf.keras.layers.Dense(d_model)\n        self.dense = tf.keras.layers.Dense(d_model)\n\n    def split_heads(self, x, batch_size):\n        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n        return tf.transpose(x, perm=[0, 2, 1, 3])\n\n    def call(self, q, k, v, q_pos=None, k_pos=None, mask=None):\n        batch_size = tf.shape(q)[0]\n        q = self.wq(q)\n        k = self.wk(k)\n        v = self.wv(v)\n\n        if q_pos is not None:\n            q = q + q_pos\n        if k_pos is not None:\n            k = k + k_pos\n\n        q = self.split_heads(q, batch_size)\n        k = self.split_heads(k, batch_size)\n        v = self.split_heads(v, batch_size)\n\n        scaled_attention, attention_weights = scaled_dot_product_attention(\n            q, k, v, mask)\n        scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])\n        scaled_attention = tf.reshape(scaled_attention,\n                                      (batch_size, -1, self.d_model))\n        output = self.dense(scaled_attention)\n        return output, attention_weights\n\n\nclass Encoder(tf.keras.Model):\n    def __init__(self, d_model, drop_rate):\n        super(Encoder, self).__init__()\n        self.d_model = d_model\n\n        self.backbone = efn.EfficientNetB4(\n            include_top=False, weights='noisy-student')\n        self.reshape = tf.keras.layers.Reshape(\n            [-1, self.d_model], name='reshape_featuere_maps')\n\n        self.embedding = tf.keras.layers.Dense(self.d_model, activation='relu')\n        self.dropout = tf.keras.layers.Dropout(drop_rate)\n\n    def call(self, x, training):\n        x = self.backbone(x, training=training)  # (B, H, W, 1792)\n        x = self.embedding(x, training=training)  # (B, H, W, 512)\n        x = self.reshape(x, training=training)  # (B, H*W, 512)\n        x = self.dropout(x, training=training)\n        return x\n\n\nclass DecoderLayer(tf.keras.layers.Layer):\n    def __init__(self, d_model, num_heads, dff, max_len, rate=0.1):\n        super(DecoderLayer, self).__init__()\n        self.mha1 = MultiHeadAttention(d_model, num_heads)\n        self.mha2 = MultiHeadAttention(d_model, num_heads)\n\n        self.ffn = point_wise_feed_forward_network(d_model, dff)\n\n        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n        self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n\n        self.dropout1 = tf.keras.layers.Dropout(rate)\n        self.dropout2 = tf.keras.layers.Dropout(rate)\n        self.dropout3 = tf.keras.layers.Dropout(rate)\n\n    def call(self, x, enc_output, enc_pos, dec_pos, training, look_ahead_mask=None, padding_mask=None):\n        # (batch_size, target_seq_len, d_model)\n        attn1, attn_weights_block1 = self.mha1(\n            x, x, x, q_pos=dec_pos, k_pos=dec_pos, mask=look_ahead_mask)\n        attn1 = self.dropout1(attn1, training=training)\n        out1 = self.layernorm1(attn1 + x)\n\n        attn2, attn_weights_block2 = self.mha2(\n            out1, enc_output, enc_output, q_pos=dec_pos, k_pos=enc_pos)\n        attn2 = self.dropout2(attn2, training=training)\n        out2 = self.layernorm2(attn2 + out1)\n\n        ffn_output = self.ffn(out2)\n        ffn_output = self.dropout3(ffn_output, training=training)\n        out3 = self.layernorm3(ffn_output + out2)\n\n        return out3, attn_weights_block1, attn_weights_block2\n\n\nclass Decoder(tf.keras.Model):\n    def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size, max_len, rate=0.1):\n        super(Decoder, self).__init__()\n        self.d_model = d_model\n        self.num_layers = num_layers\n\n        self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)\n        self.pos_encoding_1d = positional_encoding_1d(max_len, d_model)\n        self.pos_encoding_2d = positional_encoding_2d(\n            CFG.row_size, CFG.col_size, self.d_model)\n\n        self.dec_layers = [DecoderLayer(d_model, num_heads, dff, max_len, rate)\n                           for _ in range(num_layers)]\n        self.dropout = tf.keras.layers.Dropout(rate)\n        self.final_layer = tf.keras.layers.Dense(target_vocab_size)\n\n    def call(self, x, enc_output, training, look_ahead_mask=None, padding_mask=None):\n        seq_len = tf.shape(x)[1]\n        dec_pos = self.pos_encoding_1d[:, :seq_len, :]\n        x = self.embedding(x)  # (batch_size, target_seq_len, d_model)\n        x *= tf.math.sqrt(tf.cast(self.d_model, CFG.dtype))\n\n        for i in range(self.num_layers):\n            x, block1, block2 = self.dec_layers[i](\n                x, enc_output, self.pos_encoding_2d, dec_pos, training, look_ahead_mask, padding_mask)\n\n        predictions = self.final_layer(x)\n        return predictions","metadata":{"execution":{"iopub.status.busy":"2021-06-09T13:34:14.926126Z","iopub.execute_input":"2021-06-09T13:34:14.926470Z","iopub.status.idle":"2021-06-09T13:34:15.143616Z","shell.execute_reply.started":"2021-06-09T13:34:14.926439Z","shell.execute_reply":"2021-06-09T13:34:15.142459Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# LOSS","metadata":{}},{"cell_type":"code","source":"class FocalLoss(tf.keras.losses.Loss):\n    def __init__(self,\n                 alpha=0.25,\n                 gamma=2.0,\n                 reduction=tf.keras.losses.Reduction.AUTO,\n                 name=None):\n        \"\"\"Initializes `FocalLoss`.\n        Args:\n            alpha: The `alpha` weight factor for binary class imbalance.\n            gamma: The `gamma` focusing parameter to re-weight loss.\n            reduction: (Optional) Type of `tf.keras.losses.Reduction` to apply to\n                loss. Default value is `AUTO`. `AUTO` indicates that the reduction\n                option will be determined by the usage context. For almost all cases\n                this defaults to `SUM_OVER_BATCH_SIZE`. When used with\n                `tf.distribute.Strategy`, outside of built-in training loops such as\n                `tf.keras` `compile` and `fit`, using `AUTO` or `SUM_OVER_BATCH_SIZE`\n                will raise an error. Please see this custom training [tutorial](\n                https://www.tensorflow.org/tutorials/distribute/custom_training) for\n                more details.\n          name: Optional name for the op. Defaults to 'retinanet_class_loss'.\n        \"\"\"\n        self._alpha = alpha\n        self._gamma = gamma\n        super(FocalLoss, self).__init__(reduction=reduction, name=name)\n\n    def call(self, y_true, y_pred):\n        \"\"\"Invokes the `FocalLoss`.\n        Args:\n            y_true: A tensor of size [batch, num_anchors, num_classes]\n            y_pred: A tensor of size [batch, num_anchors, num_classes]\n        Returns:\n            Summed loss float `Tensor`.\n        \"\"\"\n        with tf.name_scope('focal_loss'):\n            y_true = tf.one_hot(y_true, CFG.vocab_size)\n            y_true = tf.cast(y_true, dtype=tf.float32)\n            y_pred = tf.cast(y_pred, dtype=tf.float32)\n            positive_label_mask = tf.equal(y_true, 1.0)\n            cross_entropy = (\n                tf.nn.sigmoid_cross_entropy_with_logits(labels=y_true, logits=y_pred))\n            probs = tf.sigmoid(y_pred)\n            probs_gt = tf.where(positive_label_mask, probs, 1.0 - probs)\n            # With small gamma, the implementation could produce NaN during back prop.\n            modulator = tf.pow(1.0 - probs_gt, self._gamma)\n            loss = modulator * cross_entropy\n            weighted_loss = tf.where(positive_label_mask, self._alpha * loss,\n                                     (1.0 - self._alpha) * loss)\n\n        return weighted_loss\n\n    def get_config(self):\n        config = {\n            'alpha': self._alpha,\n            'gamma': self._gamma,\n        }\n        base_config = super(FocalLoss, self).get_config()\n        return dict(list(base_config.items()) + list(config.items()))","metadata":{"execution":{"iopub.status.busy":"2021-06-09T13:34:15.145216Z","iopub.execute_input":"2021-06-09T13:34:15.145712Z","iopub.status.idle":"2021-06-09T13:34:15.161154Z","shell.execute_reply.started":"2021-06-09T13:34:15.145663Z","shell.execute_reply":"2021-06-09T13:34:15.160007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# METRICS","metadata":{}},{"cell_type":"code","source":"class Tokenizer(object):\n\n    def __init__(self):\n        self.stoi = {\n            '(': 0, ')': 1, '+': 2, ',': 3, '-': 4, '/b': 5, '/c': 6, '/h': 7, '/i': 8, '/m': 9, '/s': 10,\n            '/t': 11, '0': 12, '1': 13, '10': 14, '100': 15, '101': 16, '102': 17, '103': 18, '104': 19, '105': 20,\n            '106': 21, '107': 22, '108': 23, '109': 24, '11': 25, '110': 26, '111': 27, '112': 28, '113': 29, '114': 30,\n            '115': 31, '116': 32, '117': 33, '118': 34, '119': 35, '12': 36, '120': 37, '121': 38, '122': 39, '123': 40,\n            '124': 41, '125': 42, '126': 43, '127': 44, '128': 45, '129': 46, '13': 47, '130': 48, '131': 49, '132': 50,\n            '133': 51, '134': 52, '135': 53, '136': 54, '137': 55, '138': 56, '139': 57, '14': 58, '140': 59, '141': 60,\n            '142': 61, '143': 62, '144': 63, '145': 64, '146': 65, '147': 66, '148': 67, '149': 68, '15': 69, '150': 70,\n            '151': 71, '152': 72, '153': 73, '154': 74, '155': 75, '156': 76, '157': 77, '158': 78, '159': 79, '16': 80,\n            '161': 81, '163': 82, '165': 83, '167': 84, '17': 85, '18': 86, '19': 87, '2': 88, '20': 89, '21': 90,\n            '22': 91, '23': 92, '24': 93, '25': 94, '26': 95, '27': 96, '28': 97, '29': 98, '3': 99, '30': 100,\n            '31': 101, '32': 102, '33': 103, '34': 104, '35': 105, '36': 106, '37': 107, '38': 108, '39': 109, '4': 110,\n            '40': 111, '41': 112, '42': 113, '43': 114, '44': 115, '45': 116, '46': 117, '47': 118, '48': 119, '49': 120,\n            '5': 121, '50': 122, '51': 123, '52': 124, '53': 125, '54': 126, '55': 127, '56': 128, '57': 129, '58': 130,\n            '59': 131, '6': 132, '60': 133, '61': 134, '62': 135, '63': 136, '64': 137, '65': 138, '66': 139, '67': 140,\n            '68': 141, '69': 142, '7': 143, '70': 144, '71': 145, '72': 146, '73': 147, '74': 148, '75': 149, '76': 150,\n            '77': 151, '78': 152, '79': 153, '8': 154, '80': 155, '81': 156, '82': 157, '83': 158, '84': 159, '85': 160,\n            '86': 161, '87': 162, '88': 163, '89': 164, '9': 165, '90': 166, '91': 167, '92': 168, '93': 169, '94': 170,\n            '95': 171, '96': 172, '97': 173, '98': 174, '99': 175, 'B': 176, 'Br': 177, 'C': 178, 'Cl': 179, 'D': 180,\n            'F': 181, 'H': 182, 'I': 183, 'N': 184, 'O': 185, 'P': 186, 'S': 187, 'Si': 188, 'T': 189, '<sos>': 190,\n            '<eos>': 191, '<pad>': 192}\n        self.itos = {v: k for k, v in self.stoi.items()}\n\n    def __len__(self):\n        return len(self.stoi)\n\n    def fit_on_texts(self, texts):\n        vocab = set()\n        for text in texts:\n            vocab.update(text.split(' '))\n        vocab = sorted(vocab)\n        vocab.append('<sos>')\n        vocab.append('<eos>')\n        vocab.append('<pad>')\n        for i, s in enumerate(vocab):\n            self.stoi[s] = i\n        self.itos = {item[1]: item[0] for item in self.stoi.items()}\n\n    def text_to_sequence(self, text):\n        sequence = []\n        sequence.append(self.stoi['<sos>'])\n        for s in text.split(' '):\n            sequence.append(self.stoi[s])\n        sequence.append(self.stoi['<eos>'])\n        return sequence\n\n    def texts_to_sequences(self, texts):\n        sequences = []\n        for text in texts:\n            sequence = self.text_to_sequence(text)\n            sequences.append(sequence)\n        return sequences\n\n    def sequence_to_text(self, sequence):\n        return ''.join(list(map(lambda i: self.itos[i], sequence)))\n\n    def sequences_to_texts(self, sequences):\n        texts = []\n        for sequence in sequences:\n            text = self.sequence_to_text(sequence)\n            texts.append(text)\n        return texts\n\n    def predict_caption(self, sequence):\n        caption = ''\n        for i in sequence:\n            i = (i + 192) % 193\n            if i == self.stoi['<eos>'] or i == self.stoi['<pad>']:\n                break\n            elif i == self.stoi['<sos>']:\n                continue\n            caption += self.itos[i]\n        return caption\n\n    def predict_captions(self, sequences):\n        captions = []\n        for sequence in sequences:\n            caption = self.predict_caption(sequence)\n            captions.append(caption)\n        return captions\n\n\ntokenizer = Tokenizer()\n\ndef get_levenshtein_distance(preds, lbls, return_preds=True):\n    preds = tf.cast(preds, tf.int64)\n\n    lbls = strategy.gather(lbls, axis=0)\n    lbls = tf.cast(lbls, tf.int64)\n\n    y_trues = tokenizer.predict_captions(lbls.numpy())\n    y_preds = tokenizer.predict_captions(preds.numpy())\n\n    scores = []\n    for true, pred in zip(y_trues, y_preds):\n        score = Levenshtein.distance(true, pred)\n        scores.append(score)\n    avg_score = np.mean(scores)\n\n    if return_preds:\n        return avg_score, y_trues, y_preds, scores\n    return avg_score\n\n\ndef check_preds(preds, lbls):\n    preds = tf.cast(preds, tf.int64)\n\n    lbls = strategy.gather(lbls, axis=0)\n    lbls = tf.cast(lbls, tf.int64)\n\n    y_true = tokenizer.predict_captions(lbls.numpy())\n    y_pred = tokenizer.predict_captions(preds.numpy())\n    for i, (true, pred) in enumerate(zip(y_true, y_pred)):\n        print('*'*100)\n        print('preds :', pred)\n        print('labels:', true)\n        print('score:', Levenshtein.distance(true, pred))\n        if i == 30:\n            break\n\n\ndef check_test_preds(preds):\n    preds = tf.cast(preds, tf.int64)\n    y_pred = tokenizer.predict_captions(preds.numpy())\n    for i, pred in enumerate(y_pred):\n        print('*'*100)\n        print('preds :', pred)\n        if i == 30:\n            break","metadata":{"execution":{"iopub.status.busy":"2021-06-09T13:34:15.163315Z","iopub.execute_input":"2021-06-09T13:34:15.163627Z","iopub.status.idle":"2021-06-09T13:34:15.201529Z","shell.execute_reply.started":"2021-06-09T13:34:15.163598Z","shell.execute_reply":"2021-06-09T13:34:15.200610Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# SCHEDULER","metadata":{}},{"cell_type":"code","source":"def lrfn(step, WARMUP_LR_START, LR_START, LR_FINAL, DECAYS):\n    # exponential warmup\n    if step < CFG.warmup_steps:\n        warmup_factor = (step / CFG.warmup_steps) ** 2\n        lr = WARMUP_LR_START + (LR_START - WARMUP_LR_START) * warmup_factor\n    # staircase decay\n    else:\n        power = (step - CFG.warmup_steps) // ((CFG.total_steps -\n                                               CFG.warmup_steps) / (DECAYS + 1))\n        decay_factor = ((LR_START / LR_FINAL) ** (1 / DECAYS)) ** power\n        lr = LR_START / decay_factor\n\n    return round(lr, 8)\n\n\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, loss=None):\n        self.lr = self.lr_schedule[step]\n        # assign learning rate to optimizer\n        self.opt.learning_rate.assign(self.lr)\n\n    def get_counter(self):\n        return self.c\n\n    def get_lr(self):\n        return self.lr","metadata":{"execution":{"iopub.status.busy":"2021-06-09T13:34:15.202822Z","iopub.execute_input":"2021-06-09T13:34:15.203224Z","iopub.status.idle":"2021-06-09T13:34:15.219625Z","shell.execute_reply.started":"2021-06-09T13:34:15.203196Z","shell.execute_reply":"2021-06-09T13:34:15.218800Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# TRAINER","metadata":{"execution":{"iopub.status.busy":"2021-06-05T04:40:13.520776Z","iopub.execute_input":"2021-06-05T04:40:13.521152Z","iopub.status.idle":"2021-06-05T04:40:13.525189Z","shell.execute_reply.started":"2021-06-05T04:40:13.521121Z","shell.execute_reply":"2021-06-05T04:40:13.524006Z"}}},{"cell_type":"code","source":"def log(batch, loss, t_start_batch, val_loss, val_ls_distance, val_acc, lr):\n    # training metrics\n    print(\n        f'Step %s|' % f'{batch}/{CFG.steps_per_epoch}'.ljust(9, ' '),\n        f'loss: %.3f,' % loss,\n        f'acc: %.3f, ' % train_accuracy.result(),\n        end='')\n\n    # plot validation metrics if given\n    if val_loss is not None and val_ls_distance is not None and val_acc is not None:\n        print(\n            f'val_loss: %.3f, ' % val_loss,\n            f'val lsd: %s,' % ('%.3f' % val_ls_distance).ljust(5, ' '),\n            f'val_acc: %.3f, ' % val_acc,\n            end='')\n    # always end with learning rate, batch duration and line break\n    print(\n        f'lr: %s,' % ('%.3E' % lr).ljust(7),\n        f't: %s sec' % int(time.time() - t_start_batch),\n    )\n\nclass Trainer:\n    def __init__(self, encoder, decoder, optimizer, scheduler, loss_fn, metric_fn, num_epochs=100, resume=False, resume_epoch=0, steps_per_epoch=None):\n        self.encoder = encoder\n        self.decoder = decoder\n        self.optimizer = optimizer\n        self.scheduler = scheduler\n        self.loss_fn = loss_fn\n        self.metric_fn = metric_fn\n        self.num_epochs = num_epochs\n        self.resume = resume\n\n        self.init_epoch = 1\n        self.total_steps = 0\n        if self.resume:\n            files = tf.io.gfile.glob(f'{CFG.save_dir}/encoder_step*')\n            last_step = 0\n            if len(files) > 0:\n                last_step = int(files[-1].strip('').split('step')[1][:6])\n                self.init_epoch = last_step // CFG.steps_per_epoch + 1\n                self.total_steps = last_step\n                encoder_path = f'{CFG.save_dir}/encoder_step{self.total_steps:06}.ckpt'\n                decoder_path = f'{CFG.save_dir}/decoder_step{self.total_steps:06}.ckpt'\n                self.encoder.load_weights(encoder_path)\n                self.decoder.load_weights(decoder_path)\n                self.scheduler.step(self.total_steps)\n                print('load encoder from:', encoder_path)\n                print('load decoder from:', decoder_path)\n\n\n    def train_step(self, images, labels):\n        labels_input = labels[:, :-1]\n        labels_target = labels[:, 1:]\n        dec_mask = create_masks_decoder(labels_target)\n\n        with tf.GradientTape() as tape:\n            enc_output = self.encoder(images, training=True)\n            predictions = self.decoder(\n                labels_input, enc_output, training=True, look_ahead_mask=dec_mask)\n            loss = self.loss_fn(labels_target, predictions)\n            self.metric_fn.update_state(labels_target, predictions)\n\n        variables = self.encoder.trainable_variables + self.decoder.trainable_variables\n        gradients = tape.gradient(loss, variables)\n        gradients, _ = tf.clip_by_global_norm(gradients, 5.0)\n        self.optimizer.apply_gradients(zip(gradients, variables))\n\n        return loss\n\n    @tf.function\n    def distributed_train_step(self, images, labels):\n        per_replica_losses = strategy.run(\n            self.train_step, args=(images, labels))\n        loss = strategy.reduce(tf.distribute.ReduceOp.SUM,\n                               per_replica_losses, axis=None)\n        return loss\n\n    def validation_step(self, images, labels):\n        total_loss = 0.0\n\n        enc_output = self.encoder(images, training=False)\n\n        batch_size = tf.shape(images)[0]\n        output = tf.fill([batch_size, 1], value=START_TOKEN)\n        output = tf.cast(output, tf.int32)\n\n        for t in tqdm(range(1, CFG.seq_len)):\n            dec_mask = create_masks_decoder(output)\n            predictions = self.decoder(\n                output, enc_output, False, dec_mask)\n            predictions = predictions[:, -1:, :]\n\n            loss = self.loss_fn(labels[:, t], tf.squeeze(predictions))\n            total_loss += loss\n            self.metric_fn.update_state(labels[:, t], tf.squeeze(predictions))\n\n            dec_input = tf.math.argmax(\n                predictions, axis=-1, output_type=tf.int32)\n            output = tf.concat([output, dec_input], axis=1)\n\n        return total_loss, output\n\n    @tf.function\n    def distributed_val_step(self, images_val, labels_val):\n        per_replica_losses, per_replica_predictions_seq = strategy.run(\n            self.validation_step, args=(images_val, labels_val))\n        loss = strategy.reduce(tf.distribute.ReduceOp.SUM,\n                               per_replica_losses, axis=None)\n        predictions_seq = strategy.gather(per_replica_predictions_seq, axis=0)\n\n        return loss, predictions_seq\n\n    def evaluate(self, val_dist_dataset):\n        total_loss = 0.0\n        total_ls_distance = 0.0\n        total_acc = 0.0\n        results = {}\n        results['preds'] = []\n        results['labels'] = []\n        results['scores'] = []\n\n        for step, (images, labels) in tqdm(enumerate(val_dist_dataset)):\n            batch_loss, predictions_seq = self.distributed_val_step(\n                images, labels)\n            levenshtein_distance, text_labels, text_preds, scores = get_levenshtein_distance(\n                predictions_seq, labels)\n\n            results['preds'].extend(text_preds)\n            results['labels'].extend(text_labels)\n            results['scores'].extend(scores)\n\n            if step == 0:\n                check_preds(predictions_seq, labels)\n\n            total_loss += batch_loss / CFG.val_steps\n            total_ls_distance += levenshtein_distance / CFG.val_steps\n            total_acc += self.metric_fn.result() / CFG.val_steps\n            self.metric_fn.reset_states()\n\n            if step + 1 == CFG.val_steps:\n                return total_loss, total_ls_distance, total_acc, results\n\n    def test_step(self, images):\n        enc_output = self.encoder(images, training=False)\n\n        batch_size = tf.shape(images)[0]\n        output = tf.fill([batch_size, 1], value=START_TOKEN)\n        output = tf.cast(output, tf.int32)\n\n        for _ in tqdm(range(1, CFG.seq_len)):\n            dec_mask = create_masks_decoder(output)\n            predictions = self.decoder(\n                output, enc_output, False, dec_mask)\n            predictions = predictions[:, -1:, :]\n\n            dec_input = tf.math.argmax(\n                predictions, axis=-1, output_type=tf.int32)\n            output = tf.concat([output, dec_input], axis=1)\n\n        return output\n\n    @tf.function\n    def distributed_test_step(self, images):\n        per_replica_predictions_seq = strategy.run(\n            self.test_step, args=(images,))\n        predictions_seq = strategy.gather(per_replica_predictions_seq, axis=0)\n\n        return predictions_seq\n\n    def predict(self, test_dataset, num_test_steps):\n        test_dist_dataset = strategy.experimental_distribute_dataset(\n            test_dataset)\n\n        all_predictions = {}\n        for (step, (images, image_ids)) in tqdm(enumerate(test_dist_dataset), total=num_test_steps):\n            predictions_seq = self.distributed_test_step(images)\n            predictions_text = tokenizer.predict_captions(\n                predictions_seq.numpy())\n            image_ids = strategy.gather(image_ids, axis=0)\n\n            if step == 0:\n                check_test_preds(predictions_seq)\n\n            for text, image_id in zip(predictions_text, image_ids):\n                image_id = image_id.numpy().decode()\n                all_predictions[image_id] = text\n\n            if step == num_test_steps - 1:\n                return all_predictions\n\n    def fit(self, train_dataset, val_dataset):\n        best_metric = 10e6\n        for epoch in range(self.init_epoch, self.num_epochs + 1):\n            print(f'***** EPOCH {epoch} *****')\n            t_start = time.time()\n            t_start_batch = time.time()\n            total_loss = 0.0\n\n            train_dist_dataset = strategy.experimental_distribute_dataset(\n                train_dataset)\n            val_dist_dataset = strategy.experimental_distribute_dataset(\n                val_dataset)\n\n            for (step, (images, labels)) in enumerate(train_dist_dataset):\n                self.total_steps += 1\n                step += 1\n                batch_loss = self.distributed_train_step(images, labels)\n                batch_loss = tf.cast(batch_loss, tf.float32)\n\n                if step == CFG.steps_per_epoch and epoch % CFG.eval_freq == 0:\n                    val_loss, val_ls_distance, val_acc, results = self.evaluate(\n                        val_dist_dataset)\n\n                    val_loss = tf.cast(val_loss, tf.float32)\n                    log(step, batch_loss, t_start_batch,\n                        val_loss, val_ls_distance, val_acc, self.scheduler.get_lr())\n                    self.metric_fn.reset_states()\n\n                    if val_ls_distance < best_metric:\n                        print('best updated to ', val_ls_distance)\n                        best_metric = val_ls_distance\n                        if YOUR_GCS_DIR is not None:\n                            self.encoder.save_weights(\n                                f'{CFG.save_dir}/best_encoder.ckpt')\n                            self.decoder.save_weights(\n                                f'{CFG.save_dir}/best_decoder.ckpt')\n\n                elif step % CFG.verbose_freq == 0:\n                    log(step, batch_loss, t_start_batch,\n                        None, None, None, self.scheduler.get_lr())\n                    self.metric_fn.reset_states()\n                    t_start_batch = time.time()\n\n                total_loss += batch_loss\n\n                if step == CFG.steps_per_epoch:\n                    break\n\n                self.scheduler.step(self.total_steps)\n\n            print(\n                f'Epoch {epoch} Loss {round(total_loss.numpy() / CFG.steps_per_epoch, 3)}, time: {int(time.time() - t_start)} sec\\n')\n\n","metadata":{"execution":{"iopub.status.busy":"2021-06-09T13:34:15.220725Z","iopub.execute_input":"2021-06-09T13:34:15.221119Z","iopub.status.idle":"2021-06-09T13:34:15.265914Z","shell.execute_reply.started":"2021-06-09T13:34:15.221090Z","shell.execute_reply":"2021-06-09T13:34:15.264571Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# MAIN","metadata":{"execution":{"iopub.status.busy":"2021-06-05T04:40:38.226461Z","iopub.execute_input":"2021-06-05T04:40:38.226792Z","iopub.status.idle":"2021-06-05T04:40:38.231057Z","shell.execute_reply.started":"2021-06-05T04:40:38.226765Z","shell.execute_reply":"2021-06-05T04:40:38.229954Z"}}},{"cell_type":"code","source":"with strategy.scope():\n    loss_object = FocalLoss(reduction=tf.keras.losses.Reduction.NONE)\n\n    def loss_function(real, pred):\n        per_example_loss = loss_object(real, pred)\n        return tf.nn.compute_average_loss(per_example_loss, global_batch_size=CFG.batch_size)\n\n    # Metrics\n    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n        name='train_accuracy')\n\n    # Encoder\n    encoder = Encoder(\n        CFG.d_model,\n        CFG.encoder_drop_rate\n    )\n\n    # Decoder\n    decoder = Decoder(\n        CFG.num_layers,\n        CFG.d_model,\n        CFG.num_heads,\n        CFG.dff,\n        CFG.vocab_size,\n        CFG.seq_len,\n        CFG.decoder_drop_rate\n    )\n\n    # Adam Optimizer\n    optimizer = tfa_optimizers.AdamW(weight_decay=1.0e-4)\n\n    lr_fn = [lrfn(step, 1e-6, 5e-4, 1e-5, CFG.num_epochs)\n             for step in range(CFG.total_steps)]\n    scheduler = LRReduce(optimizer, lr_fn)\n\n    trainer = Trainer(\n        encoder=encoder,\n        decoder=decoder,\n        optimizer=optimizer,\n        scheduler=scheduler,\n        loss_fn=loss_function,\n        metric_fn=train_accuracy,\n        num_epochs=CFG.num_epochs,\n        resume=CFG.resume,\n        steps_per_epoch=CFG.steps_per_epoch\n    )\n\nif CFG.inference:\n    test_dataset, test_length = get_test_dataset(CFG.test_batch_size)\n    num_test_steps = test_length // CFG.batch_size + 1\n    all_predictions = trainer.predict(test_dataset, num_test_steps)\n    if YOUR_GCS_DIR is not None:\n        with tf.io.gfile.GFile(f'{CFG.save_dir}/test_results.json', 'w') as f:\n            json.dump(all_predictions, f)\n\nelse:\n    train_dataset, train_length = get_dataset('train', CFG.batch_size)\n    val_dataset, val_length = get_dataset('val', CFG.batch_size)\n    print('train samples:', train_length)\n    print('val samples:', val_length)\n\n    trainer.fit(train_dataset, val_dataset)","metadata":{"execution":{"iopub.status.busy":"2021-06-09T13:34:15.267352Z","iopub.execute_input":"2021-06-09T13:34:15.267691Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}