{"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":"# Overview\n\nthis is a vision transformer baseline, is far from params fine tuning well \n\nI have not enough time to experiment before submit deadline maybe\n\nso I will be very glad if you find it usefull\n\n\n# Thanks\n\nThis notebook is a bit too late because I really spend too much time (about three weeks) dealing with transformer learning rate problem.\n\nYou change checkout [this](https://www.kaggle.com/c/bms-molecular-translation/discussion/241716#1326759) for detail of spike learning rate.\n\nSpecial thank given to [Darien Schettler](https://www.kaggle.com/dschettler8845) \n\n# Refference\n\n1. transformer https://www.kaggle.com/aditya08/imagecaptioning-show-attnd-tell-w-transformer/ this is from https://www.tensorflow.org/tutorials/text/transformer?hl=zh-cn\n2. tpu efficientnetV2 rnn https://www.kaggle.com/dschettler8845/bms-efficientnetv2-tpu-e2e-pipeline-in-3hrs/execution\n3. TPU CNN + Transformer https://www.kaggle.com/dschettler8845/3-71-cv-bms-efficientnetv2-transformer-e2e\n4. Spike learning rate is all you need without it you will be stuck in about 20% accuray and very unstable https://www.kaggle.com/c/bms-molecular-translation/discussion/241716#1326759\n5. vision transormer params detail https://github.com/faustomorales/vit-keras\n\n# ToubleShooting\n1. this note in colab is much slower than in kaggle, I dont know why\n2. if you valid lsd in very earlier get_levenshtein_distance will be very slow if you get a low lsd score , this may because it may reach the worst case of algorithm (BIG O) https://en.wikipedia.org/wiki/Levenshtein_distance#Computational_complexity","metadata":{}},{"cell_type":"markdown","source":"# Dependence","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\n\nimport tensorflow as tf\nfrom kaggle_datasets import KaggleDatasets\n\nfrom datetime import datetime\nimport random\nimport shutil\nimport math\nimport time\nimport sys\nimport io\nimport os\nimport gc\n\nimport matplotlib.pyplot as plt\n\nimport cv2","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:31.454341Z","iopub.execute_input":"2021-06-05T03:14:31.454719Z","iopub.status.idle":"2021-06-05T03:14:31.460356Z","shell.execute_reply.started":"2021-06-05T03:14:31.454686Z","shell.execute_reply":"2021-06-05T03:14:31.459314Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Setup TPU","metadata":{}},{"cell_type":"code","source":"print(f\"\\n... ACCELERATOR SETUP STARTING ...\\n\")\n\n# Detect hardware, return appropriate distribution strategy\ntry:\n    # TPU detection. No parameters necessary if TPU_NAME environment variable is set. On Kaggle this is always the case.\n    TPU=tf.distribute.cluster_resolver.TPUClusterResolver()  \nexcept ValueError:\n    TPU=None\n\nif TPU:\n    print(f\"\\n... RUNNING ON TPU - {TPU.master()}...\")\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    print(f\"\\n... RUNNING ON CPU/GPU ...\")\n    # Yield the default distribution strategy in Tensorflow\n    #   --> Works on CPU and single GPU.\n    strategy=tf.distribute.get_strategy() \n\n# What Is a Replica?\n#    --> A single Cloud TPU device consists of FOUR chips, each of which has TWO TPU cores. \n#    --> Therefore, for efficient utilization of Cloud TPU, a program should make use of each of the EIGHT (4x2) cores. \n#    --> Each replica is essentially a copy of the training graph that is run on each core and \n#        trains a mini-batch containing 1/8th of the overall batch size\nN_REPLICAS=strategy.num_replicas_in_sync\n    \nprint(f\"... # OF REPLICAS: {N_REPLICAS} ...\\n\")\n\nprint(f\"\\n... ACCELERATOR SETUP COMPLTED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:31.468847Z","iopub.execute_input":"2021-06-05T03:14:31.469363Z","iopub.status.idle":"2021-06-05T03:14:36.731541Z","shell.execute_reply.started":"2021-06-05T03:14:31.469332Z","shell.execute_reply":"2021-06-05T03:14:36.730791Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... DATA ACCESS SETUP STARTED ...\\n\")\n\n#if TPU:\n    # Google Cloud Dataset path to training and validation images\n# gcs Input tfrecord \nDATA_DIR=KaggleDatasets().get_gcs_path('bms-train-tfrecords-half-length')\n\n#DATA_DIR= './imnput/bms-train-tfrecords-half-length'\n\nprint(f\"\\n... DATA DIRECTORY PATH IS:\\n\\t--> {DATA_DIR}\")\n\nprint(f\"\\n... IMMEDIATE CONTENTS OF DATA DIRECTORY IS:\")\nfor file in tf.io.gfile.glob(os.path.join(DATA_DIR, \"*\")): print(f\"\\t--> {file}\")\n\nprint(\"\\n\\n... DATA ACCESS SETUP COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:36.732591Z","iopub.execute_input":"2021-06-05T03:14:36.733004Z","iopub.status.idle":"2021-06-05T03:14:37.11269Z","shell.execute_reply.started":"2021-06-05T03:14:36.73297Z","shell.execute_reply":"2021-06-05T03:14:37.111505Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"\\n... MIXED PRECISION SETUP STARTING ...\\n\")\nprint(\"\\n... SET TF TO OPERATE IN MIXED PRECISION – `bfloat16` – IF ON TPU ...\")\n\n# Set Mixed Precision Global Policy\n#     ---> To use mixed precision in Keras, you need to create a `tf.keras.mixed_precision.Policy`\n#          typically referred to as a dtype policy. \n#     ---> Dtype policies specify the dtypes layers will run in\ntf.keras.mixed_precision.set_global_policy('mixed_bfloat16' if TPU else 'float32')\n\n# target data type, bfloat16 when using TPU to improve throughput\nTARGET_DTYPE=tf.bfloat16 if TPU else tf.float32\nprint(f\"\\t--> THE TARGET DTYPE HAS BEEN SET TO {TARGET_DTYPE} ...\")\n\n# The policy specifies two important aspects of a layer: \n#     1. The dtype the layer's computations are done in\n#     2. The dtype of a layer's variables. \nprint(f\"\\n... TWO IMPORTANT ASPECTS OF THE GLOBAL MIXED PRECISION POLICY:\")\nprint(f'\\t--> COMPUTE DTYPE  : {tf.keras.mixed_precision.global_policy().compute_dtype}')\nprint(f'\\t--> VARIABLE DTYPE : {tf.keras.mixed_precision.global_policy().variable_dtype}')\n\nprint(f\"\\n\\n... MIXED PRECISION SETUP COMPLTED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:37.115011Z","iopub.execute_input":"2021-06-05T03:14:37.115452Z","iopub.status.idle":"2021-06-05T03:14:37.124094Z","shell.execute_reply.started":"2021-06-05T03:14:37.115408Z","shell.execute_reply":"2021-06-05T03:14:37.123092Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"\\n... XLA OPTIMIZATIONS STARTING ...\\n\")\n\nprint(f\"\\n... CONFIGURE JIT (JUST IN TIME) COMPILATION ...\\n\")\n# enable XLA optmizations (10% speedup when using @tf.function calls)\ntf.config.optimizer.set_jit(True)\n\nprint(f\"\\n... XLA OPTIMIZATIONS COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:37.125606Z","iopub.execute_input":"2021-06-05T03:14:37.125884Z","iopub.status.idle":"2021-06-05T03:14:37.138954Z","shell.execute_reply.started":"2021-06-05T03:14:37.125857Z","shell.execute_reply":"2021-06-05T03:14:37.137637Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... BASIC DATA SETUP STARTING ...\\n\")\n\n# All the possible tokens in our InChI 'language'\nTOKEN_LIST=[\"<PAD>\", \"InChI=1S/\", \"<END>\", \"/c\", \"/h\", \"/m\", \"/t\", \"/b\", \"/s\", \"/i\"] +\\\n             ['Si', 'Br', 'Cl', 'F', 'I', 'N', 'O', 'P', 'S', 'C', 'H', 'B', ] +\\\n             [str(i) for i in range(167,-1,-1)] +\\\n             [\"\\+\", \"\\(\", \"\\)\", \"\\-\", \",\", \"D\", \"T\"]\nprint(f\"\\n... TOKEN LIST:\")\nfor i, tok in enumerate(TOKEN_LIST): print(f\"\\t--> INTEGER-IDX={i:<3}  –––  STRING={tok}\")\n\n# The start/end/pad tokens will be removed from the string when computing the Levenshtein distance\n# We want them as tf.constant's so they will operate properly within the @tf.function context\nSTART_TOKEN=tf.constant(TOKEN_LIST.index(\"InChI=1S/\"), dtype=tf.uint8)\nEND_TOKEN=tf.constant(TOKEN_LIST.index(\"<END>\"), dtype=tf.uint8)\nPAD_TOKEN=tf.constant(TOKEN_LIST.index(\"<PAD>\"), dtype=tf.uint8)\n\n# Prefixes and Their Respective Ordering/Format\n#      -- ORDERING --> {c}{h/None}{b/None}{t/None}{m/None}{s/None}{i/None}{h/None}{t/None}{m/None}\nPREFIX_ORDERING=\"chbtmsihtm\"\nprint(f\"\\n... PREFIX ORDERING IS {PREFIX_ORDERING} ...\")\n\n# Paths to Respective Image Directories\nTRAIN_DIR=os.path.join(DATA_DIR, \"train_records\")\nVAL_DIR=os.path.join(DATA_DIR, \"val_records\")\n#TEST_DIR=os.path.join(TEST_DATA_DIR, \"test_records\")\n\n# Get the Full Paths to The Individual TFRecord Files\nTRAIN_TFREC_PATHS=sorted(\n    tf.io.gfile.glob(os.path.join(TRAIN_DIR, \"*.tfrec\")), \n    key=lambda x: int(x.rsplit(\"_\", 2)[1]))\nVAL_TFREC_PATHS=sorted(\n    tf.io.gfile.glob(os.path.join(VAL_DIR, \"*.tfrec\")), \n    key=lambda x: int(x.rsplit(\"_\", 2)[1]))\n\n\"\"\"\nTEST_TFREC_PATHS=sorted(\n    tf.io.gfile.glob(os.path.join(TEST_DIR, \"*.tfrec\")), \n    key=lambda x: int(x.rsplit(\"_\", 2)[1]))\n\"\"\"\n\n\nprint(f\"\\n... TFRECORD INFORMATION:\")\nfor SPLIT, TFREC_PATHS in zip([\"TRAIN\", \"VAL\"\n                              # , \"TEST\"\n                              ], [TRAIN_TFREC_PATHS, VAL_TFREC_PATHS\n                              #    , TEST_TFREC_PATHS\n                                 ]):\n    print(f\"\\t--> {len(TFREC_PATHS):<3} {SPLIT:<5} TFRECORDS\")\n\n# Paths to relevant CSV files containing training and submission information\nTRAIN_CSV_PATH=os.path.join(\"/kaggle/input\", \"bms-csvs-w-extra-metadata\", \"train_labels_w_extra.csv\")\nSS_CSV_PATH   =os.path.join(\"/kaggle/input\", \"bms-csvs-w-extra-metadata\", \"sample_submission_w_extra.csv\")\nprint(f\"\\n... PATHS TO CSVS:\")\nprint(f\"\\t--> TRAIN CSV: {TRAIN_CSV_PATH}\")\nprint(f\"\\t--> SS CSV   : {SS_CSV_PATH}\")\n\n# When debug is true we use a smaller batch size and smaller model\nDEBUG=False\n\nprint(\"\\n\\n... BASIC DATA SETUP COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:37.140265Z","iopub.execute_input":"2021-06-05T03:14:37.14057Z","iopub.status.idle":"2021-06-05T03:14:37.318796Z","shell.execute_reply.started":"2021-06-05T03:14:37.140541Z","shell.execute_reply":"2021-06-05T03:14:37.317863Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... INITIAL DATAFRAME INSTANTIATION STARTING ...\\n\")\n\n# Load the train and submission dataframes\ntrain_df=pd.read_csv(TRAIN_CSV_PATH)\nss_df   =pd.read_csv(SS_CSV_PATH)\n\n# --- Distribution Information ---\nN_EX   =len(train_df)\nN_TEST =len(ss_df)\nN_VAL  =80_000 # Fixed from dataset creation information\nN_TRAIN=N_EX-N_VAL\n\n# --- Batching Information ---\nBATCH_SIZE_DEBUG=2\n#REPLICA_BATCH_SIZE=64 # Could probably be 128\n\nREPLICA_BATCH_SIZE=32\n\nif DEBUG:\n    REPLICA_BATCH_SIZE=BATCH_SIZE_DEBUG\nOVERALL_BATCH_SIZE=REPLICA_BATCH_SIZE*N_REPLICAS\n\n\n# --- Input Image Information ---\nIMG_SHAPE=(192,384,3)\n\n# --- Autocalculate Training/Validation/Testing Information ---\nTRAIN_STEPS=N_TRAIN  // OVERALL_BATCH_SIZE\nVAL_STEPS  =N_VAL    // OVERALL_BATCH_SIZE\nTEST_STEPS =int(np.ceil(N_TEST/OVERALL_BATCH_SIZE))\n\n# This is for padding our test dataset so we only have whole batches\nREQUIRED_DATASET_PAD=OVERALL_BATCH_SIZE-N_TEST%OVERALL_BATCH_SIZE\n\n\nprint(f\"\\n... # OF TRAIN+VAL EXAMPLES  : {N_EX:<7} ...\")\nprint(f\"... # OF TRAIN EXAMPLES      : {N_TRAIN:<7} ...\")\nprint(f\"... # OF VALIDATION EXAMPLES : {N_VAL:<7} ...\")\nprint(f\"... # OF TEST EXAMPLES       : {N_TEST:<7} ...\\n\")\n\nprint(f\"\\n... REPLICA BATCH SIZE    : {REPLICA_BATCH_SIZE} ...\")\nprint(f\"... OVERALL BATCH SIZE    : {OVERALL_BATCH_SIZE} ...\\n\")\n\nprint(f\"\\n... IMAGE SHAPE           : {IMG_SHAPE} ...\\n\")\n\nprint(f\"\\n... TRAIN STEPS PER EPOCH : {TRAIN_STEPS:<5} ...\")\nprint(f\"... VAL STEPS PER EPOCH   : {VAL_STEPS:<5} ...\")\nprint(f\"... TEST STEPS PER EPOCH  : {TEST_STEPS:<5} ...\\n\")\n\nprint(\"\\n... TRAIN DATAFRAME ...\\n\")\ndisplay(train_df.head(3))\n\nprint(\"\\n... SUBMISSION DATAFRAME ...\\n\")\ndisplay(ss_df.head(3))\n\nprint(\"\\n... INITIAL DATAFRAME INSTANTIATION COMPLETED...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:37.319978Z","iopub.execute_input":"2021-06-05T03:14:37.320244Z","iopub.status.idle":"2021-06-05T03:14:51.807536Z","shell.execute_reply.started":"2021-06-05T03:14:37.320217Z","shell.execute_reply":"2021-06-05T03:14:51.805915Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prepare Dataset","metadata":{}},{"cell_type":"code","source":"print(\"\\n... SPECIAL VARIABLE SETUP STARTING ...\\n\")\n\n\n# Whether to start training using previously checkpointed model\nLOAD_MODEL       =False\nENCODER_CKPT_PATH=\"\"\nDECODER_CKPT_PATH=\"\"\n\nif LOAD_MODEL:\n    print(f\"\\n... ENCODER MODEL TRAINING WILL RESUME FROM PREVIOUS CHECKPOINT:\\n\\t-->{ENCODER_CKPT_PATH}\\n\")\n    print(f\"... DECODER MODEL TRAINING WILL RESUME FROM PREVIOUS CHECKPOINT:\\n\\t-->{DECODER_CKPT_PATH}\\n\")\nelse:\n    print(f\"\\n... MODEL TRAINING WILL START FROM SCRATCH ...\\n\")\n\n    \nprint(\"\\n... SPECIAL VARIABLE SETUP COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:51.809192Z","iopub.execute_input":"2021-06-05T03:14:51.809646Z","iopub.status.idle":"2021-06-05T03:14:51.817255Z","shell.execute_reply.started":"2021-06-05T03:14:51.8096Z","shell.execute_reply":"2021-06-05T03:14:51.816335Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def flatten_l_o_l(nested_list):\n    \"\"\" Flatten a list of lists \"\"\"\n    return [item for sublist in nested_list for item in sublist]\n\n\ndef tf_load_image(path, img_size=(192,384,3), invert=False):\n    \"\"\" Load an image with the correct size and shape \n    \n    Args:\n        path (tf.string): Path to the image to be loaded\n        img_size (tuple, optional): Size to reshape image to (required for TPU)\n        tile_to_3_channel (bool, optional): Whether to tile the single channel\n            image to 3 channels which will be required for most off-the-shelf models\n        invert (bool, optional): Whether or not to invert the background/foreground\n    \n    Returns:\n        3 channel tf.Constant image ready for training/inference\n    \n    \"\"\"\n    img=decode_img(tf.io.read_file(path), img_size, n_channels=3, invert=invert)        \n    return img\n    \n    \ndef decode_image(image_data, resize_to=(192,384,3)):\n    \"\"\" Function to decode the tf.string containing image information \n    \n    \n    Args:\n        image_data (tf.string): String containing encoded image data from tf.Example\n        resize_to (tuple, optional): Size that we will reshape the tensor to (required for TPU)\n    \n    Returns:\n        Tensor containing the resized single-channel image in the appropriate dtype\n    \"\"\"\n    image=tf.image.decode_png(image_data, channels=3)\n    image=tf.reshape(image, resize_to)\n    return tf.cast(image, TARGET_DTYPE)\n    \n    \n# sparse tensors are required to compute the Levenshtein distance\ndef dense_to_sparse(dense):\n    \"\"\" Convert a dense tensor to a sparse tensor \n    \n    Args:\n        dense (Tensor): TBD\n        \n    Returns:\n        A sparse tensor    \n    \"\"\"\n    indices=tf.where(tf.ones_like(dense))\n    values=tf.reshape(dense, (MAX_LEN*OVERALL_BATCH_SIZE,))\n    sparse=tf.SparseTensor(indices, values, dense.shape)\n    return sparse\n\ndef get_levenshtein_distance(preds, lbls):\n    \"\"\" Computes the Levenshtein distance between the predictions and labels \n    \n    Args:\n        preds (tensor): Batch of predictions\n        lbls (tensor): Batch of labels\n        \n    Returns:\n        The mean Levenshtein distance calculated across the batch\n    \"\"\"\n    preds=tf.where(tf.not_equal(lbls, END_TOKEN) & tf.not_equal(lbls, PAD_TOKEN), preds, 0)\n    lbls=tf.where(tf.not_equal(lbls, END_TOKEN), lbls, 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-05T03:14:51.820364Z","iopub.execute_input":"2021-06-05T03:14:51.820661Z","iopub.status.idle":"2021-06-05T03:14:51.833914Z","shell.execute_reply.started":"2021-06-05T03:14:51.820627Z","shell.execute_reply":"2021-06-05T03:14:51.832722Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n\\n... STARTING PREPARING VARIABLES FOR DATASET ...\\n\")\n\ntok_2_int={c.strip(\"\\\\\"):i for i,c in enumerate(TOKEN_LIST)}\nint_2_tok={v:k for k,v in tok_2_int.items()}\n\n# Max Length Was Determined Previously Using... \n#     >>> MAX_LEN=train_df.InChI.progress_apply(lambda x: len(re.findall(\"|\".join(TOKEN_LIST), x))).max()+1\nMAX_LEN=((train_df.inchi_token_len.max()+1)//2) # //2 yields 138... which is half of max length (speeds up training)\nVOCAB_LEN=len(int_2_tok)\n\nprint(f\"\\t--> TOKEN TO INTEGER MAP     : {tok_2_int}\")\nprint(f\"\\t--> INTEGER TO TOKEN MAP     : {int_2_tok}\")\nprint(f\"\\t--> MAX # OF TOKENS IN INCHI : {MAX_LEN}\")\nprint(f\"\\t--> LENGTH OF VOCAB          : {VOCAB_LEN}\")\n\nprint(f\"\\n\\n\\t--> CONVERTED INCHI STRINGS  :\")\nfor i, row in train_df.iloc[:N_VAL].sample(3).iterrows():\n    print(f\"\\n\\t\\t--> EXAMPLE #{i} FROM THE VALIDATION DATASET\")\n    print(\"\\t\\t\\t--> RAW INCHI : \", row[\"InChI\"])\n\nprint(\"\\n\\n... PREPARING VARIABLES FOR DATASET COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:51.835971Z","iopub.execute_input":"2021-06-05T03:14:51.836316Z","iopub.status.idle":"2021-06-05T03:14:51.8589Z","shell.execute_reply.started":"2021-06-05T03:14:51.836287Z","shell.execute_reply":"2021-06-05T03:14:51.858111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... CREATE TFRECORD RAW DATASETS STARTING ...\\n\")\n\n# Create tf.data.Dataset from filepaths for conversion later\nraw_train_ds=tf.data.TFRecordDataset(TRAIN_TFREC_PATHS, num_parallel_reads=None)\nraw_val_ds=tf.data.TFRecordDataset(VAL_TFREC_PATHS, num_parallel_reads=None)\n#raw_test_ds=tf.data.TFRecordDataset(TEST_TFREC_PATHS, num_parallel_reads=None)\n\n# raw_test_ds=tf.data.TFRecordDataset(TEST_TFREC_PATHS, num_parallel_reads=None)\n\nprint(f\"\\n... THE RAW TF.DATA.TFRECORDDATASET OBJECT:\\n\\t--> {raw_train_ds}\\n\")\n\nprint(\"\\n... CREATE TFRECORD RAW DATASETS COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:51.86013Z","iopub.execute_input":"2021-06-05T03:14:51.860753Z","iopub.status.idle":"2021-06-05T03:14:51.879945Z","shell.execute_reply.started":"2021-06-05T03:14:51.86072Z","shell.execute_reply":"2021-06-05T03:14:51.879291Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... RAW TFRECORD INVESTIGATION TO DETERMINE FEATURE DESCRIPTIONS STARTED ...\\n\")\n\nprint(\"\\n... EXAMPLE OF TRUNCATED RAW TFRECORD/TFEXAMPLE FROM TRAINING DATASET TO SHOW HOW TO FIND FEATURE DESCRIPTIONS:\\n\")\n# See an example\nfor raw in raw_train_ds.take(1):\n    example=tf.train.Example()\n    example.ParseFromString(raw.numpy())\n    for i, (k,v) in enumerate(example.features.feature.items()):\n        print(f\"\\tFEATURE #{i+1}\")\n        print(f\"\\t\\t--> KEY={k}\")\n        if k!=\"image\":\n            try:\n                print(f\"\\t\\t\\t--> TRUNCATED-VALUE={v.int64_list.value[:15]} ...\\n\")\n            except:\n                print(f\"\\t\\t\\t--> TRUNCATED-VALUE={v.bytes_list.value[0][:25]} ...\\n\")\n        else:\n            print(f\"\\t\\t\\t--> TRUNCATED-VALUE={str(v.bytes_list.value[0][:25])} ...\\n\")         \n\nprint(\"\\n... RAW TFRECORD INVESTIGATION TO DETERMINE FEATURE DESCRIPTIONS COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:51.881009Z","iopub.execute_input":"2021-06-05T03:14:51.881481Z","iopub.status.idle":"2021-06-05T03:14:51.953168Z","shell.execute_reply.started":"2021-06-05T03:14:51.881428Z","shell.execute_reply":"2021-06-05T03:14:51.951977Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def decode(serialized_example, is_test=False, tokenized_inchi=True):\n    \"\"\" Parses a set of features and label from the given `serialized_example`.\n        \n        It is used as a map function for `dataset.map`\n\n    Args:\n        serialized_example (tf.Example): A serialized example containing the\n            following features:\n                – 'image'\n                – 'image_id'\n                – 'inchi'\n        is_test (bool, optional): Whether to allow for the InChI feature\n        drop_id (bool, optional): Whether or not to drop the ID feature\n        \n    Returns:\n        A decoded tf.data.Dataset object representing the tfrecord dataset\n    \"\"\"\n    feature_dict={\n        'image': tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value=''),\n    }\n    \n    if not is_test:\n        if tokenized_inchi:\n            feature_dict[\"inchi\"]=tf.io.FixedLenFeature(shape=[MAX_LEN], dtype=tf.int64, default_value=[0]*MAX_LEN)\n        else:\n            feature_dict[\"inchi\"]=tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value='')\n    else:\n        feature_dict[\"image_id\"]=tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value='')\n    \n    # Define a parser\n    features=tf.io.parse_single_example(serialized_example, features=feature_dict)\n    \n    # Decode the tf.string\n    image=decode_image(features['image'], resize_to=IMG_SHAPE)\n    \n    # Figure out the correct information to return\n    if is_test:\n        image_id=features[\"image_id\"] \n        return image, image_id\n    else:\n        if tokenized_inchi:\n            target=tf.cast(features[\"inchi\"], tf.uint8)\n        else:\n            target=features[\"inchi\"]\n        return image, target","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:51.954812Z","iopub.execute_input":"2021-06-05T03:14:51.955237Z","iopub.status.idle":"2021-06-05T03:14:51.964919Z","shell.execute_reply.started":"2021-06-05T03:14:51.955188Z","shell.execute_reply":"2021-06-05T03:14:51.964042Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... DECODING RAW TFRECORD DATASETS STARTING ...\\n\")\n\n# Decode the tfrecords completely –– decode is our `_parse_function` (from recipe above)\ntrain_ds=raw_train_ds.map(lambda x: decode(x, is_test=False))\nval_ds=raw_val_ds.map(lambda x: decode(x, is_test=False))\n#test_ds=raw_test_ds.map(lambda x: decode(x, is_test=True))\n\nprint(f\"\\n... THE DECODED TF.DATA.TFRECORDDATASET OBJECT:\" \\\n      f\"\\n\\t--> ((image), (image_id - optional), (inchi))\" \\\n      f\"\\n\\t--> {train_ds}\\n\")\n\nprint(\"\\n... 2 EXAMPLES OF IMAGES AND LABELS AFTER DECODING ...\")\nfor i, (img, inchi) in enumerate(train_ds.take(2)):\n    print(f\"\\nIMAGE SHAPE : {img.shape}\")\n    print(f\"IMAGE INCHI : {[int_2_tok[x] for x in inchi.numpy()]}\\n\")\n    plt.figure(figsize=(10,10))\n    plt.imshow(img.numpy().astype(np.int64), cmap=\"gray\")\n    plt.title(f\"{''.join([int_2_tok[x] for x in inchi.numpy() if x!=0][:50])} ... [truncated]\")\n    plt.show()\n\nprint(\"\\n... DECODING RAW TFRECORD DATASETS COMPLETED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:51.965924Z","iopub.execute_input":"2021-06-05T03:14:51.966336Z","iopub.status.idle":"2021-06-05T03:14:52.673042Z","shell.execute_reply.started":"2021-06-05T03:14:51.966305Z","shell.execute_reply":"2021-06-05T03:14:52.672291Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_dataset(filenames, is_test=False, ordered=False, tokenized_inchi=True):\n    \"\"\"Read from TFRecords.\n    \n    For optimal performance, reading from multiple files at once and disregarding data order (if `ordered=False`).\n        - If pulling InChI from TFRecords than order does not matter since we will \n          be shuffling the data anyway (for training dataset).\n          \n    Args:\n        filenames (list of strings): List of paths to that point to the respective TFRecord files\n        is_test (bool, optional): Whether or not to include the image ID or label in the returned dataset\n        ordered (bool, optional): Whether to ensured ordered results or maximize parallelization\n        tokenized_inchi (bool, optional): Whether our dataset includes the tokenized inchi or we will be \n            creating it from the caption numpy array\n        \n    Returns:\n        Decoded tf.data.Dataset object\n    \"\"\"\n\n    options=tf.data.Options()\n    if not ordered:\n        # disable order, increase speed\n        options.experimental_deterministic=False\n        N_PARALLEL=tf.data.AUTOTUNE\n    else:\n        N_PARALLEL=None\n        \n    # If not-ordered, this will read in by automatically interleaving multiple tfrecord files.\n    dataset=tf.data.TFRecordDataset(filenames, num_parallel_reads=N_PARALLEL)\n    \n    # If not-ordered, this will ensure that we use data as soon as it \n    # streams in, rather than in its original order.\n    dataset=dataset.with_options(options) \n    \n    # parse and return a dataset w/ the appropriate configuration\n    dataset=dataset.map(\n        lambda x: decode(x, is_test, tokenized_inchi),\n        num_parallel_calls=N_PARALLEL,\n    )\n    \n    return dataset\n\ndef get_dataset(filenames, batch_size, \n                is_test=False, \n                shuffle_buffer_size=1, \n                repeat_dataset=True, \n                preserve_file_order=False, \n                drop_remainder=True,\n                tokenized_inchi=True,\n                external_inchi_dataset=None,\n                test_padding=0):\n    \"\"\" Get a tf.data.Dataset w/ the appropriate configuration\n    \n    Args:\n        TBD\n        test_padding (int, optional): Amount required to pad dataset to have only full batches\n        \n    Returns:\n        TBD\n        \n    \"\"\"\n    # Load the dataset\n    dataset=load_dataset(filenames, is_test, preserve_file_order, tokenized_inchi)\n    \n    if test_padding!=0:\n        pad_dataset=tf.data.Dataset.from_tensor_slices((\n            tf.zeros((test_padding, *IMG_SHAPE), dtype=TARGET_DTYPE),       # Fake Images\n            tf.constant([\"000000000000\",]*test_padding, dtype=tf.string))   # Fake IDs\n        )\n        dataset=dataset.concatenate(pad_dataset)\n    \n    # If we are training than we will want to repeat the dataset. \n    # We will determine the number of steps (or updates) later for 1 training epoch.\n    if repeat_dataset:\n        dataset=dataset.repeat()\n    \n    # If we need to add on manually the inchi\n    if external_inchi_dataset is not None:\n        # Zip the datasets and tile the 1 channel image to 3 channels & drop the old inchi value\n        dataset=tf.data.Dataset.zip((dataset, external_inchi_dataset))\n        dataset=dataset.map(lambda x,y: (tf.tile(tf.expand_dims(x[0], -1), tf.constant([1,1,3], tf.int32)), y))\n                              \n    # Shuffling\n    if shuffle_buffer_size!=1:\n        dataset=dataset.shuffle(shuffle_buffer_size)\n    \n    # Batching\n    dataset=dataset.batch(batch_size, drop_remainder=drop_remainder)\n    \n    # prefetch next batch while training (autotune prefetch buffer size)\n    dataset=dataset.prefetch(tf.data.experimental.AUTOTUNE)\n    \n    return dataset","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:52.674233Z","iopub.execute_input":"2021-06-05T03:14:52.674507Z","iopub.status.idle":"2021-06-05T03:14:52.688256Z","shell.execute_reply.started":"2021-06-05T03:14:52.674481Z","shell.execute_reply":"2021-06-05T03:14:52.687414Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"####### ####### ####### ####### ####### ####### ####### #######\n\n# Template Configuration\nDS_TEMPLATE_CONFIG=dict(\n    filenames=[],\n    batch_size=1,\n    is_test=False, \n    shuffle_buffer_size=1, \n    repeat_dataset=True, \n    preserve_file_order=False, \n    drop_remainder=True,\n    tokenized_inchi=True,\n    external_inchi_dataset=None,\n    test_padding=0\n)\n\n####### ####### ####### ####### ####### ####### ####### #######\n\n# Individual Respective Configurations\nTRAIN_DS_CONFIG=DS_TEMPLATE_CONFIG.copy()\nTRAIN_DS_CONFIG.update(dict(\n    filenames=TRAIN_TFREC_PATHS,\n    batch_size=OVERALL_BATCH_SIZE,\n    shuffle_buffer_size=OVERALL_BATCH_SIZE*6,\n))\n\nVAL_DS_CONFIG=DS_TEMPLATE_CONFIG.copy()\nVAL_DS_CONFIG.update(dict(\n    filenames=VAL_TFREC_PATHS,\n    batch_size=OVERALL_BATCH_SIZE,\n))\n\n####### ####### ####### ####### ####### ####### ####### #######\n\ntrain_ds=get_dataset(**TRAIN_DS_CONFIG)\nval_ds=get_dataset(**VAL_DS_CONFIG)\n#test_ds=get_dataset(**TEST_DS_CONFIG)\n\nfor SPLIT, CONFIG in zip([\"TRAINING\", \"VALIDATION\"\n                          #, \"TESTING\"\n                         ], [TRAIN_DS_CONFIG, VAL_DS_CONFIG\n                           #  , TEST_DS_CONFIG\n                            ]): # , TEST_DS_CONFIG]\n    print(f\"\\n... {SPLIT} CONFIGURATION:\")\n    for k,v in CONFIG.items():\n        if k==\"filenames\":\n            print(f\"\\t--> {k:<23}: {[path.split('/', 4)[-1] for path in v[:2]]+['...']}\")\n        else:\n            print(f\"\\t--> {k:<23}: {v}\")\n\nprint(f\"\\n\\n... TRAINING DATASET   : {train_ds} ...\")\nprint(f\"... VALIDATION DATASET : {val_ds} ...\")\n#print(f\"... TESTING DATASET    : {test_ds}    ...\\n\")\n\nprint(\"\\n\\n ... SOME VALIDATION EXAMPLES ... \\n\\n\")\nfor x,y in val_ds.take(1):\n    for i in range(2):\n        plt.figure(figsize=(12,12))\n        plt.imshow(x[i].numpy().astype(np.int64))\n        plt.title(f\"IMAGE INCHI : {''.join([int_2_tok[z] for z in y[i].numpy() if z not in [0,1,2]])}\\n\")\n        plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:52.689659Z","iopub.execute_input":"2021-06-05T03:14:52.689946Z","iopub.status.idle":"2021-06-05T03:14:53.786575Z","shell.execute_reply.started":"2021-06-05T03:14:52.689909Z","shell.execute_reply":"2021-06-05T03:14:53.78557Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"# 先把 decoder的参数放在这\nnum_layer=4\nd_model=768\ndff= 512\nnum_heads=8\ntarget_vocab_size=VOCAB_LEN\ndropout_rate=0.1 \n\n# vision encoder 参数\nimage_size = (IMG_SHAPE[0], IMG_SHAPE[1])  # We'll resize input images to this size\npatch_size = 16  # Size of the patches to be extract from the input images\nnum_patches = (image_size[0] // patch_size) * (image_size[1] // patch_size) \nprojection_dim = d_model\nvit_heads = 12\ntransformer_units = [\n    projection_dim * 2,\n    projection_dim,\n]  # Size of the transformer layers\ntransformer_layers = 12","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.787727Z","iopub.execute_input":"2021-06-05T03:14:53.787987Z","iopub.status.idle":"2021-06-05T03:14:53.795094Z","shell.execute_reply.started":"2021-06-05T03:14:53.787962Z","shell.execute_reply":"2021-06-05T03:14:53.793346Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Attention","metadata":{}},{"cell_type":"code","source":"def scaled_dot_product_attention(q, k, v, mask):\n    \n    \"\"\"Calculate the attention weights.\n    q, k, v must have matching leading dimensions.\n    k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.\n    The mask has different shapes depending on its type(padding or look ahead) \n    but it must be broadcastable for addition.\n\n    Args:\n        q: query shape == (..., seq_len_q, depth)\n        k: key shape == (..., seq_len_k, depth)\n        v: value shape == (..., seq_len_v, depth_v)\n        mask: Float tensor with shape broadcastable \n            to (..., seq_len_q, seq_len_k). Defaults to None.\n\n    Returns:\n        output, attention_weights\n    \"\"\"\n    \n    # (..., seq_len_q, seq_len_k)\n    matmul_qk=tf.matmul(q, k, transpose_b=True)  \n    \n    # scale matmul_qk\n    dk=tf.cast(tf.shape(k)[-1], TARGET_DTYPE)\n    \n    # Calculate scaled attention logits\n    scaled_attention_logits=matmul_qk / tf.math.sqrt(dk)\n    \n    # add the mask to the scaled tensor.\n    if mask is not None:\n        scaled_attention_logits += (mask * -1e9) \n    \n    # softmax is normalized on the last axis (seq_len_k) \n    # so that the scores add up to 1.\n    #   - shape --> (..., seq_len_q, seq_len_k)\n    attention_weights=tf.nn.softmax(scaled_attention_logits, axis=-1) \n    \n     #   - shape --> (..., seq_len_q, depth_v)\n    output=tf.matmul(attention_weights, v)  \n\n    return output, attention_weights","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.796089Z","iopub.execute_input":"2021-06-05T03:14:53.796612Z","iopub.status.idle":"2021-06-05T03:14:53.808959Z","shell.execute_reply.started":"2021-06-05T03:14:53.796585Z","shell.execute_reply":"2021-06-05T03:14:53.808289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def point_wise_feed_forward_network(d_model, dff):\n    return tf.keras.Sequential([\n                # INNER LAYER\n                #   – (batch_size, seq_len, dff)\n                tf.keras.layers.Dense(dff, activation='relu'),\n                 # INNER LAYER\n                #   – (batch_size, seq_len, dff)\n                tf.keras.layers.Dense(d_model)])","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.810121Z","iopub.execute_input":"2021-06-05T03:14:53.810612Z","iopub.status.idle":"2021-06-05T03:14:53.819372Z","shell.execute_reply.started":"2021-06-05T03:14:53.81058Z","shell.execute_reply":"2021-06-05T03:14:53.81854Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class 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        \"\"\" Split the last dimension into (num_heads, depth).\n            Then we transpose the result such that the shape is \n                - (batch_size, num_heads, seq_len, depth)\n        \"\"\"\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, v, k, q, mask=None):\n        batch_size=tf.shape(q)[0]\n        q=self.wq(q)  # (batch_size, seq_len, d_model)\n        k=self.wk(k)  # (batch_size, seq_len, d_model)\n        v=self.wv(v)  # (batch_size, seq_len, d_model)\n\n        q=self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)\n        k=self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)\n        v=self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)\n    \n        # scaled_attention.shape – (batch_size, num_heads, seq_len_q, depth)\n        # attention_weights.shape – (batch_size, num_heads, seq_len_q, seq_len_k)\n        scaled_attention, attention_weights=scaled_dot_product_attention(q, k, v, mask)\n        \n        # (batch_size, seq_len_q, num_heads, depth)\n        scaled_attention=tf.transpose(scaled_attention, perm=[0, 2, 1, 3]) \n    \n        # (batch_size, seq_len_q, d_model)\n        concat_attention=tf.reshape(scaled_attention,\n                                 (batch_size, -1, self.d_model))  \n\n        # (batch_size, seq_len_q, d_model)\n        output=self.dense(concat_attention)  \n        \n        return output, attention_weights","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.820553Z","iopub.execute_input":"2021-06-05T03:14:53.820905Z","iopub.status.idle":"2021-06-05T03:14:53.834448Z","shell.execute_reply.started":"2021-06-05T03:14:53.820786Z","shell.execute_reply":"2021-06-05T03:14:53.83358Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Position embedding","metadata":{}},{"cell_type":"code","source":"def get_angles(pos, i, d_model):\n    \"\"\"\n        angle_rates=1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))\n        return pos * angle_rates\n    \"\"\"\n    # 浮点类型不匹配\n    angle_rates = tf.constant(1, TARGET_DTYPE) / tf.math.pow(tf.constant(10000, TARGET_DTYPE), (tf.constant(2, dtype=TARGET_DTYPE) * tf.cast((i//2), TARGET_DTYPE))/d_model)\n    return pos * angle_rates\n\n# https://kazemnejad.com/blog/transformer_architecture_positional_encoding/\n# https://stackoverflow.com/questions/46452020/sinusoidal-embedding-attention-is-all-you-need\n\n# 这个需要改成tnt 或者vit\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),row,axis=0).reshape(-1,1)\n\n    angle_rads_row=get_angles(row_pos,np.arange(d_model//2)[np.newaxis,:],d_model//2)\n    angle_rads_col=get_angles(col_pos,np.arange(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)[np.newaxis, ...]\n    #return tf.cast(pos_encoding, dtype=tf.float32)\n    return tf.cast(pos_encoding, dtype=TARGET_DTYPE)\n\"\"\"","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.83575Z","iopub.execute_input":"2021-06-05T03:14:53.836183Z","iopub.status.idle":"2021-06-05T03:14:53.850124Z","shell.execute_reply.started":"2021-06-05T03:14:53.836077Z","shell.execute_reply":"2021-06-05T03:14:53.849049Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def do_interleave(arr_a, arr_b):\n    a_arr_tf_column = tf.range(arr_a.shape[1])*2 # [0 2 4 ...]\n    b_arr_tf_column = tf.range(arr_b.shape[1])*2+1 # [1 3 5 ...]\n    column_indices = tf.argsort(tf.concat([a_arr_tf_column,b_arr_tf_column],axis=-1))\n    column, row = tf.meshgrid(column_indices,tf.range(arr_a.shape[0]))\n    combine_indices = tf.stack([row,column],axis=-1)\n    combine_value = tf.concat([arr_a,arr_b],axis=1)\n    return tf.gather_nd(combine_value,combine_indices)\n\ndef positional_encoding_1d(position, d_model):\n    angle_rads = get_angles(tf.cast(tf.range(position)[:, tf.newaxis], TARGET_DTYPE),\n                            tf.cast(tf.range(d_model)[tf.newaxis, :], TARGET_DTYPE),\n                            d_model)\n    \n    # apply sin to even indices in the array; 2i\n    sin_angle_rads = tf.math.sin(angle_rads[:, ::2])\n    cos_angle_rads = tf.math.cos(angle_rads[:, 1::2])\n    angle_rads = do_interleave(sin_angle_rads, cos_angle_rads)\n    pos_encoding = angle_rads[tf.newaxis, ...]\n    return pos_encoding","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.851118Z","iopub.execute_input":"2021-06-05T03:14:53.851395Z","iopub.status.idle":"2021-06-05T03:14:53.860276Z","shell.execute_reply.started":"2021-06-05T03:14:53.85137Z","shell.execute_reply":"2021-06-05T03:14:53.859475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Encoder","metadata":{}},{"cell_type":"code","source":"from tensorflow.keras import layers","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.861628Z","iopub.execute_input":"2021-06-05T03:14:53.862171Z","iopub.status.idle":"2021-06-05T03:14:53.873278Z","shell.execute_reply.started":"2021-06-05T03:14:53.862132Z","shell.execute_reply":"2021-06-05T03:14:53.87262Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class PatchEncoder(layers.Layer):\n    def __init__(self, patch_size, num_patches, projection_dim):\n        super(PatchEncoder, self).__init__()\n        self.num_patches = num_patches\n        self.patch_size = patch_size\n        self.projection_dim = projection_dim\n\n        self.conv =  layers.Conv2D(\n            filters=self.projection_dim,\n            kernel_size=self.patch_size,\n            strides=self.patch_size,\n            padding=\"valid\",\n            name=\"embedding\",)\n        \n        self.reshape = tf.keras.layers.Reshape((self.num_patches, self.projection_dim))\n        \n        self.position_embedding = layers.Embedding(\n            input_dim=self.num_patches, output_dim=self.projection_dim\n        )\n        self.add = layers.Add()\n\n    def call(self, images):\n        positions = tf.range(start=0, limit=self.num_patches, delta=1)\n\n        patches = self.conv(images)\n        proj = self.reshape(patches)\n\n        pos = self.position_embedding(positions)\n\n        encoded = proj + pos\n        return encoded","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.877132Z","iopub.execute_input":"2021-06-05T03:14:53.877395Z","iopub.status.idle":"2021-06-05T03:14:53.886005Z","shell.execute_reply.started":"2021-06-05T03:14:53.877372Z","shell.execute_reply":"2021-06-05T03:14:53.884997Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class VIT_ENCODER(layers.Layer):\n    def __init__(self, patch_size, num_patches, projection_dim, transformer_layers, transformer_units):\n        super(VIT_ENCODER, self).__init__()\n        self.patch_size = patch_size\n        self.num_patches = num_patches\n        self.projection_dim = projection_dim\n        self.transformer_layers = transformer_layers\n        self.transformer_units = transformer_units\n        self.dropout_rate = 0.1\n\n        # divide patch\n        #self.patches = Patches(self.patch_size)\n        # Encode patches.\n        self.PatchEncoder = PatchEncoder(self.patch_size, self.num_patches, self.projection_dim)\n\n        self.layerNorm1 = layers.LayerNormalization(epsilon=1e-6)\n        self.layerNorm2 = layers.LayerNormalization(epsilon=1e-6)\n        self.layerNorm3 = layers.LayerNormalization(epsilon=1e-6)\n\n        self.mha = layers.MultiHeadAttention(\n                num_heads=vit_heads, key_dim=projection_dim, dropout=0.1\n        )\n        \n        self.add1 = layers.Add()\n        self.add2 = layers.Add()\n        \n\n        #def mlp(x, hidden_units, dropout_rate):\n       \n        self.Denses = [layers.Dense(units, activation=tf.nn.gelu) for units in self.transformer_units]\n        self.Dropouts = [layers.Dropout(self.dropout_rate) for _ in range(len(self.transformer_units))]\n\n        #self.mlp = mlp\n\n    def call(self, x, training):\n        \n        # Create patches.\n        #patches = self.patches(x)\n        # Encode patches.\n        encoded_patches = self.PatchEncoder(x)\n\n        # Create multiple layers of the Transformer block.\n        for _ in range(self.transformer_layers):\n            # Layer normalization 1.\n            x1 = self.layerNorm1(encoded_patches)\n            # Create a multi-head attention layer.\n            attention_output = self.mha(x1, x1)\n            # Skip connection 1.\n            x2 = self.add1([attention_output, encoded_patches])\n            # Layer normalization 2.\n            x3 = self.layerNorm2(x2)\n            # MLP.\n            #x3 = self.mlp(x3, hidden_units=transformer_units, dropout_rate=0.1)\n            for i in range(len(self.transformer_units)):\n              x3 = self.Denses[i](x3)\n              x3 = self.Dropouts[i](x3, training=training)\n\n            # Skip connection 2.\n            encoded_patches = self.add2([x3, x2])\n\n        # Create a [batch_size, projection_dim] tensor.\n        x = self.layerNorm3(encoded_patches)\n\n        return x","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.887427Z","iopub.execute_input":"2021-06-05T03:14:53.887848Z","iopub.status.idle":"2021-06-05T03:14:53.901856Z","shell.execute_reply.started":"2021-06-05T03:14:53.887824Z","shell.execute_reply":"2021-06-05T03:14:53.900989Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Decoder","metadata":{}},{"cell_type":"code","source":"class DecoderLayer(tf.keras.layers.Layer):\n    def __init__(self, d_model, num_heads, dff, rate=0.1):\n        \"\"\" Decoder Layer Component Of Transformer Decoder Block \n        \n        Args:\n            d_model: the depth of the d-dimensional space used for positional encoding of image embedding\n            num_heads: The number of heads to use in the multi-head-attention block\n            dff (int): Number of units to use in the feed-forward neural network\n            dropout_rate (float): Percentage of nodes to drop in a given layer\n        \n        Returns:\n            none; This is an intiailization\n        \"\"\"\n        super(DecoderLayer, self).__init__()\n        self.mha1=MultiHeadAttention(d_model, num_heads)\n        self.mha2=MultiHeadAttention(d_model, num_heads)\n        \n        # keras 多头注意力实现\n        # # Multi Head Attention Layers\n        # self.mha1 = tf.keras.layers.MultiHeadAttention(num_heads, key_dim=d_model,)\n        # self.mha2 = tf.keras.layers.MultiHeadAttention(num_heads, key_dim=d_model,)\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, training,look_ahead_mask=None, padding_mask=None):\n        \"\"\" Call function for our encoder layer\n        \n        Args:\n            self (class attributes): For the given layer, this exposes the attributes\n            x (tf.float32 array): token embeddinng (batch_size, output_seq_len, embedding_dim)\n            enc_output (tf.float32 array): token embeddinng (batch_size, output_seq_len, embedding_dim)\n            training (bool): Whether or not to apply certain operations (i.e. disable/enable dropout)\n            mask (): tbd\n            \n        Returns:\n            The encoded input sequence\n               - shape --> (batch_size, input_seq_len, d_model)\n        \"\"\"\n        attn1, attn_weights_block1=self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)\n        attn1=self.dropout1(attn1, training=training)\n        \n        # Residual connection followed by layer normalization\n        #   – (batch_size, target_seq_len, d_model)\n        out1=self.layernorm1(attn1 + x)\n        \n        # Merging connection between encoder and decoder (MHA)\n        #   – (batch_size, target_seq_len, d_model)\n        attn2, attn_weights_block2=self.mha2(enc_output, enc_output, out1, padding_mask) \n        attn2=self.dropout2(attn2, training=training)\n        \n        # Residual connection followed by layer normalization\n        #   – (batch_size, target_seq_len, d_model)\n        out2=self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)\n        \n        # (batch_size, target_seq_len, d_model)\n        ffn_output=self.ffn(out2)  # (batch_size, target_seq_len, d_model)\n        ffn_output=self.dropout3(ffn_output, training=training)\n        \n        # Residual connection followed by layer normalization\n        #   – (batch_size, target_seq_len, d_model)\n        out3=self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)\n\n        return out3, attn_weights_block1, attn_weights_block2\n\nprint(\"\\n... ATTENTION MECHANISM LAYER CREATION FINISHED ...\\n\")\n\nprint(\"\\n... DECODER MODEL CREATION STARTING ...\\n\")\n\n\n\nclass Decoder(tf.keras.layers.Layer):\n    def __init__(self, num_layers,d_model,num_heads,dff, target_vocab_size, maximum_position_encoding,   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=positional_encoding_1d(maximum_position_encoding, d_model)\n\n        self.dec_layers=[DecoderLayer(d_model, num_heads, dff, rate)\n                         for _ in range(num_layers)]\n        self.dropout=tf.keras.layers.Dropout(rate)\n\n    def call(self, x, enc_output, training,look_ahead_mask=None, padding_mask=None):\n        seq_len=tf.shape(x)[1]\n        attention_weights={}\n\n        x=self.embedding(x)  # (batch_size, target_seq_len, d_model)\n        #x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n        x *= tf.math.sqrt(tf.cast(self.d_model, TARGET_DTYPE))\n        #print(self.pos_encoding.shape, self.pos_encoding[:, :seq_len, :].shape)\n        x += self.pos_encoding[:, :seq_len, :]\n        x=self.dropout(x, training=training)\n\n        for i in range(self.num_layers):\n            x, block1, block2=self.dec_layers[i](x, enc_output, training,\n                                            look_ahead_mask, padding_mask)\n\n            attention_weights['decoder_layer{}_block1'.format(i+1)]=block1\n            attention_weights['decoder_layer{}_block2'.format(i+1)]=block2\n\n        return x, attention_weights\n\n\n\nprint(\"\\n... DECODER MODEL CREATION FINISHED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.903249Z","iopub.execute_input":"2021-06-05T03:14:53.903657Z","iopub.status.idle":"2021-06-05T03:14:53.926677Z","shell.execute_reply.started":"2021-06-05T03:14:53.903617Z","shell.execute_reply":"2021-06-05T03:14:53.925582Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# learning Rate Schedule ","metadata":{}},{"cell_type":"code","source":"print(\"\\n... LEARNING RATE SCHEDULE CREATION STARTING ...\\n\")\n\n\nprint(\"\\n... LEARNING RATE SCHEDULE CREATION STARTING ...\\n\")\n\n# Part of the Training Configuration\nEPOCHS = 4\nTOTAL_STEPS = TRAIN_STEPS*EPOCHS\n\n# Learning Rate Scheduler Configuration\nWARM_STEPS = (TRAIN_STEPS-1)*4 # Suuuuuper long ramp-up\n\nclass CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n    def __init__(self, d_model, warmup_steps=4000):\n        super(CustomSchedule, self).__init__()\n\n        self.d_model = tf.cast(d_model, tf.float32)\n        self.warmup_steps = warmup_steps\n\n    def __call__(self, step):\n        arg1 = tf.math.rsqrt(step)\n        arg2 = step * (self.warmup_steps ** -1.5)\n\n        return tf.math.rsqrt(self.d_model*1.75) * tf.math.minimum(arg1, arg2)\n    \nprint(\"\\n... LEARNING RATE SCHEDULE CREATION FINISHED ...\\n\")\n","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.927921Z","iopub.execute_input":"2021-06-05T03:14:53.928215Z","iopub.status.idle":"2021-06-05T03:14:53.938758Z","shell.execute_reply.started":"2021-06-05T03:14:53.928187Z","shell.execute_reply":"2021-06-05T03:14:53.937784Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Config():\n    def __init__(self,):\n        self.encoder_config={}\n        self.decoder_config={}\n        self.lr_config={}\n    def initialize_lr_config(self, warm_steps, n_epochs):\n        self.lr_config = dict(\n            warm_steps=warm_steps, \n            n_epochs=n_epochs,\n        )\n        \ntraining_config=Config()\ntraining_config.initialize_lr_config(warm_steps=WARM_STEPS, n_epochs=EPOCHS,)\n\nprint(f\"TRAINING LEARNING RATE CONFIG:\\n\\t--> {training_config.lr_config}\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.940014Z","iopub.execute_input":"2021-06-05T03:14:53.94035Z","iopub.status.idle":"2021-06-05T03:14:53.950042Z","shell.execute_reply.started":"2021-06-05T03:14:53.940302Z","shell.execute_reply":"2021-06-05T03:14:53.948889Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"\\n... TRAINING PREPERATION STARTING ...\\n\")\n\ndef prepare_for_training(lr_config, encoder_wts=None, decoder_wts=None, verbose=0):\n    \n\n    # Everything must be declared within the scope when leveraging the TPU strategy\n    #     - This will still function properly if scope is set to another type of accelerator\n    with strategy.scope():\n        \n        print(\"\\t--> CREATING LOSS FUNCTION ...\")\n        # Declare the loss object\n        #     - Sparse categorical cross entropy loss is used as root loss\n        loss_object=tf.keras.losses.SparseCategoricalCrossentropy(\n            from_logits=True, reduction=tf.keras.losses.Reduction.NONE\n        )\n        \n        def loss_fn(real, pred):\n            # Convert to uint8\n            mask = tf.math.not_equal(real, 0)\n            loss_ = loss_object(real, pred)\n            loss_ *= tf.cast(mask, dtype=loss_.dtype)\n\n            # https://www.tensorflow.org/tutorials/distribute/custom_training#define_the_loss_function\n            loss_ = tf.nn.compute_average_loss(loss_, global_batch_size=REPLICA_BATCH_SIZE)\n            return loss_\n        \n        \n        # Declare the metrics\n        #    - Loss (train only) and sparse categorical accuracy will be used\n        print(\"\\t--> CREATING METRICS ...\")\n        metrics={\n            'batch_loss':tf.keras.metrics.Mean(),\n            'train_loss': tf.keras.metrics.Mean(),\n            'train_acc': tf.keras.metrics.SparseCategoricalAccuracy(),\n            'val_loss': tf.keras.metrics.Mean(),\n            'val_acc': tf.keras.metrics.SparseCategoricalAccuracy(),\n            'val_lsd': tf.keras.metrics.Mean(), \n        }\n        \n        print(\"\\t--> CREATING OPTIMIZER ...\")\n        \n        lr_scheduler = CustomSchedule(d_model, lr_config[\"warm_steps\"])\n        \n        # Instiate an optimizer\n        optimizer=tf.keras.optimizers.Adam(lr_scheduler)\n        #optimizer=tf.keras.optimizers.Adam()\n        \n        print(\"\\t--> CREATING LEARNING RATE SCHEDULER ...\")\n        # Declare the learning rate schedule (try this as actual lr schedule and list...)\n        \n\n\n    # Show the model architectures and plot the learning rate\n    if verbose:\n\n        print(\"\\n\\n... LR SCHEDULE PLOT...\\n\")\n        plot_lr_schedule(lr_schedule)\n  \n    return loss_fn, metrics, optimizer, lr_scheduler\n    \n    \nprint(\"\\n... GENERATING THE FOLLOWING:\")\n# Instantiate our required training components in the correct scope\nloss_fn, metrics, optimizer, lr_scheduler=\\\n    prepare_for_training(lr_config=training_config.lr_config,\n                         encoder_wts=(ENCODER_CKPT_PATH if ENCODER_CKPT_PATH!=\"\" else None),\n                         decoder_wts=(DECODER_CKPT_PATH if DECODER_CKPT_PATH!=\"\" else None),\n                         verbose=0,)\n\nprint(\"\\n... TRAINING PREPERATION FINISHED ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:53.951807Z","iopub.execute_input":"2021-06-05T03:14:53.952577Z","iopub.status.idle":"2021-06-05T03:14:54.310655Z","shell.execute_reply.started":"2021-06-05T03:14:53.952509Z","shell.execute_reply":"2021-06-05T03:14:54.309619Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Transformer","metadata":{}},{"cell_type":"code","source":"class Transformer(tf.keras.Model):\n    def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size, pe_input, pe_target, rate=0.1):\n        super(Transformer, self).__init__()\n        print(\"num_layers, d_model, num_heads, dff, target_vocab_size, pe_input, pe_target\", num_layers, d_model, num_heads, dff, target_vocab_size, pe_input, pe_target, rate)\n        self.encoder= VIT_ENCODER(patch_size, num_patches, projection_dim, transformer_layers, transformer_units)\n        self.decoder=Decoder(num_layers, d_model, num_heads, dff,\n                          target_vocab_size, pe_target, rate)\n        self.final_layer=tf.keras.layers.Dense(target_vocab_size)\n\n    def call(self, inp, tar, training,look_ahead_mask=None,dec_padding_mask=None,enc_padding_mask=None):\n        enc_output=self.encoder(inp, training)  # (batch_size, inp_seq_len, d_model      )\n        dec_output, attention_weights=self.decoder(tar, enc_output, training, look_ahead_mask, dec_padding_mask)\n        final_output=self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)\n        return final_output, attention_weights","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:54.311951Z","iopub.execute_input":"2021-06-05T03:14:54.312236Z","iopub.status.idle":"2021-06-05T03:14:54.320031Z","shell.execute_reply.started":"2021-06-05T03:14:54.312209Z","shell.execute_reply":"2021-06-05T03:14:54.319361Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"transformer = Transformer(num_layer, \n                          d_model,num_heads, \n                          dff, \n                          target_vocab_size, \n                          num_patches,\n                          MAX_LEN,\n                          dropout_rate)","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:54.320999Z","iopub.execute_input":"2021-06-05T03:14:54.321559Z","iopub.status.idle":"2021-06-05T03:14:54.476909Z","shell.execute_reply.started":"2021-06-05T03:14:54.321516Z","shell.execute_reply":"2021-06-05T03:14:54.475962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## LookForward mask","metadata":{}},{"cell_type":"code","source":"def create_padding_mask(seq):\n    # add extra dimensions to add the padding to the attention logits.\n    #    - (batch_size, 1, 1, seq_len)\n    seq=tf.cast(tf.math.equal(seq, 0), TARGET_DTYPE)\n    return seq[:, tf.newaxis, tf.newaxis, :] \n\ndef create_look_ahead_mask(size):\n    mask=1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n    #return mask  # (seq_len, seq_len)\n    return tf.cast(mask, TARGET_DTYPE)\ndef create_masks_decoder(tar):\n\n    look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])\n    dec_target_padding_mask = create_padding_mask(tar)\n\n    combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)\n\n    return combined_mask","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:54.478408Z","iopub.execute_input":"2021-06-05T03:14:54.478687Z","iopub.status.idle":"2021-06-05T03:14:54.485968Z","shell.execute_reply.started":"2021-06-05T03:14:54.478661Z","shell.execute_reply":"2021-06-05T03:14:54.484252Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"def train_step(_image_batch, _inchi_batch):\n\n    tar_inp = _inchi_batch[:, :-1]\n    tar_real = _inchi_batch[:, 1:]\n    \n    # TODO 此处可能不对\n    #enc_padding_mask = create_padding_mask(_inchi_batch)\n    # 其实是lookahead mask\n    dec_mask = create_masks_decoder(tar_inp)\n    \n    #dec_padding_mask = create_padding_mask(_inchi_batch)\n    \n    \n    with tf.GradientTape() as tape:\n            predictions, _ = transformer(_image_batch, tar_inp, True, look_ahead_mask=dec_mask)\n            batch_loss = loss_fn(tar_real, predictions)\n            # TODO 此处可能不对\n            # Update Accuracy Metric\n            metrics[\"train_acc\"].update_state(tar_real, predictions, \n                                              sample_weight=tf.where(tf.not_equal(tar_real, PAD_TOKEN), 1.0, 0.0))\n        \n    gradients = tape.gradient(batch_loss, transformer.trainable_variables)   \n    #gradients, _ = tf.clip_by_global_norm(gradients, 10.0)\n    optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))\n    #train_loss(loss)\n    #train_accuracy(tar_real, predictions)\n    metrics[\"batch_loss\"].update_state(batch_loss)\n    metrics[\"train_loss\"].update_state(batch_loss)\n    \n@tf.function\ndef dist_train_step(_image_batch, _inchi_batch):\n    strategy.run(train_step, args=(_image_batch, _inchi_batch))","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:54.487572Z","iopub.execute_input":"2021-06-05T03:14:54.487948Z","iopub.status.idle":"2021-06-05T03:14:54.958123Z","shell.execute_reply.started":"2021-06-05T03:14:54.48791Z","shell.execute_reply":"2021-06-05T03:14:54.957254Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def val_step(_image_batch, _inchi_batch):\n    \"\"\" Forward pass (calculate gradients)\n    \n    Args:\n        image_batch (): TBD\n        inchi_batch (): TBD\n    \n    Returns:\n        tbd\n    \"\"\"\n    \n    # Initialize batch_loss\n    batch_loss = tf.constant(0.0, TARGET_DTYPE)       \n    transformer_pred_batch = tf.ones((REPLICA_BATCH_SIZE, 1), dtype=tf.uint8)\n    #transformer_pred_batch = tf.ones((OVERALL_BATCH_SIZE, 1), dtype=tf.uint8)\n    \n    # Get image embedding (once)\n    #_image_embedding = encoder(_image_batch, training=False)\n    \n    # Teacher forcing - feeding the target as the next input\n    for c_idx in range(1, MAX_LEN):\n        gt_batch_id = _inchi_batch[:, c_idx]\n        combined_mask = create_masks_decoder(transformer_pred_batch)\n        \n        # predictions.shape == (batch_size, seq_len, vocab_size)\n        prediction_batch, attention_weights = transformer(_image_batch, transformer_pred_batch, training=False, look_ahead_mask=combined_mask)\n        predicted_batch_id = prediction_batch[:, -1:, :]\n\n        # Update Loss Accumulator\n        batch_loss += loss_fn(gt_batch_id, predicted_batch_id[:, -1])\n    \n        # Update Accuracy Metric\n        metrics[\"val_acc\"].update_state(gt_batch_id, predicted_batch_id[:, -1],\n                                        sample_weight=tf.where(tf.not_equal(gt_batch_id, PAD_TOKEN), 1.0, 0.0))\n\n        # no teacher forcing, predicted char is next transformer input\n        transformer_pred_batch = tf.concat([transformer_pred_batch, tf.cast(tf.argmax(predicted_batch_id, axis=-1), tf.uint8)], axis=-1)\n        \n    # Update Loss Metric\n    metrics[\"val_loss\"].update_state(batch_loss)\n    return transformer_pred_batch    \n\n    \n@tf.function\ndef dist_val_step(_val_dist_ds):\n    _val_image_batch, _val_inchi_batch = next(_val_dist_ds)\n    predictions_seq_batch_per_replica = strategy.run(val_step, args=(_val_image_batch, _val_inchi_batch))\n    predictions_seq_batch_accum = strategy.gather(predictions_seq_batch_per_replica, axis=0)\n    _val_inchi_batch_accum = strategy.gather(_val_inchi_batch, axis=0)\n    return predictions_seq_batch_accum, _val_inchi_batch_accum","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:54.96147Z","iopub.execute_input":"2021-06-05T03:14:54.961894Z","iopub.status.idle":"2021-06-05T03:14:54.972404Z","shell.execute_reply.started":"2021-06-05T03:14:54.961842Z","shell.execute_reply":"2021-06-05T03:14:54.971441Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class StatLogger():\n    def __init__(self, verbose_frequency=100, print_style=\"tight\"):\n        self.train_loss = []\n        self.train_acc = []\n        self.val_loss = []\n        self.val_acc = []\n        self.val_lsd = []\n        self.step = []\n        self.epoch = []\n        self.lr = []\n        \n        self.current_step = 0\n        self.epoch_start_time = 0\n        self.batch_start_time = 0\n        self.verbose_frequency = verbose_frequency\n        self.print_style = print_style\n        \n    def print_last_val(self, current_time):\n        if self.print_style==\"tight\":\n            print(f\"| VAL DATA |  STEP {VAL_STEPS:>4}/{VAL_STEPS} |  \" \\\n                  f\"ACC: {str(self.val_acc[-1]*100)[:5]:<5} – \" \\\n                  f\"LOSS: {str(self.val_loss[-1])[:5]:<5} – \" \\\n                  f\"LSD: {str(self.val_lsd[-1]):<3} |\")\n        else:\n            print(f'\\n\\n{\"-\"*100}\\n{\"-\"*100}\\n' \\\n                  f'{\"-\"*25:<25}{\"VALIDATION ACCURACY : \"+str(self.val_acc[-1]*100): ^50}{\"-\"*25:>25}\\n' \\\n                  f'{\"-\"*100}\\n' \\\n                  f'{\"-\"*25:<25}{\"VALIDATION LOSS     : \"+str(self.val_loss[-1]): ^50}{\"-\"*25:>25}\\n' \\\n                  f'{\"-\"*100}\\n' \\\n                  f'{\"-\"*25:<25}{\"VALIDATION LSD      : \"+str(self.val_lsd[-1]): ^50}{\"-\"*25:>25}\\n' \\\n                  f'{\"-\"*100}\\n{\"-\"*100}\\n\\n')\n    \n    def print_current_train(self, step, train_acc, train_loss, batch_loss, current_time, current_lr):\n        if self.print_style==\"tight\":\n            print(f\"| TRAIN DATA |  STEP {self.current_step:>4}/{TRAIN_STEPS} | \" \\\n                  f\"ACC: {str(train_acc*100)[:5]:<5} – \" \\\n                  f\"LOSS: {str(train_loss)[:5]:<5} – \" \\\n                  f\"LR: {current_lr:.2e} \" \\\n                  f\"|   | TIME |  EPOCH: {str(round((current_time-self.epoch_start_time)/3600,1))+'h':<5} – \" \\\n                  f\"SUBSET: {str(round((current_time-self.batch_start_time)*self.verbose_frequency,1))+'s':<6} – \" \\\n                  f\"BATCH: {str(round(current_time-self.batch_start_time,1))+'s':<5} |\")\n        else:\n            print(f'\\n\\n{\"-\"*100}\\n{\"-\"*100}\\n' \\\n                  f'{\"-\"*25:<25}{\"CURRENT STEP : \"+str(step)+\" OF \"+str(TRAIN_STEPS): ^50}{\"-\"*25:>25}\\n' \\\n                  f'{\"-\"*100}\\n' \\\n                  f'{\"-\"*25:<25}{\"CURRENT TRAIN ACCURACY : \"+str(train_acc*100): ^50}{\"-\"*25:>25}\\n' \\\n                  f'{\"-\"*100}\\n' \\\n                  f'{\"-\"*25:<25}{\"CURRENT TRAIN LOSS     : \"+str(train_loss): ^50}{\"-\"*25:>25}\\n' \\\n                  f'{\"-\"*100}\\n' \\\n                  f'{\"-\"*25:<25}{\"LAST BATCH LOSS        : \"+str(batch_loss): ^50}{\"-\"*25:>25}\\n' \\\n                  f'{\"-\"*100}\\n{\"-\"*100}\\n' \\\n                  f'{\"-\"*25:<25}{\"EPOCH ELAPSED TIME  : \"+str(round(current_time-self.epoch_start_time,1))+\" SECONDS\": ^50}{\"-\"*25:>25}\\n' \\\n                  f'{\"-\"*100}\\n' \\\n                  f'{\"-\"*25:<25}{\"LAST SET OF BATCHES TOOK  : ~\"+str(round((current_time-self.batch_start_time)*self.verbose_frequency,1))+\" SECONDS\": ^50}{\"-\"*25:>25}\\n' \\\n                  f'{\"-\"*100}\\n' \\\n                  f'{\"-\"*25:<25}{\"LAST SINGLE BATCH TOOK  : \"+str(round(current_time-self.batch_start_time,1))+\" SECONDS\": ^50}{\"-\"*25:>25}\\n' \\\n                  f'{\"-\"*100}\\n{\"-\"*100}\\n\\n')","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:54.974003Z","iopub.execute_input":"2021-06-05T03:14:54.97441Z","iopub.status.idle":"2021-06-05T03:14:54.989377Z","shell.execute_reply.started":"2021-06-05T03:14:54.97437Z","shell.execute_reply":"2021-06-05T03:14:54.988441Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Instantiate our tool for logging\nstat_logger = StatLogger()\n    \nfor epoch in range(1,EPOCHS+1):\n    print(f'\\n\\n{\"=\"*100}\\n{\"=\"*25:<25}{\"EPOCH #\"+str(epoch): ^50}{\"=\"*25:>25}\\n{\"=\"*100}\\n')\n    \n    stat_logger.current_step=0\n    stat_logger.epoch_start_time = time.time() # to compute epoch duration\n    \n    # create distributed versions of dataset to run on TPU with 8 computation units\n    train_dist_ds = strategy.experimental_distribute_dataset(train_ds)\n    val_dist_ds = iter(strategy.experimental_distribute_dataset(val_ds))\n\n    for image_batch, inchi_batch in train_dist_ds:\n                \n        # Update current step\n        stat_logger.batch_start_time = time.time()\n        \n        # Update the current step\n        stat_logger.current_step += 1\n        \n        # Calculate training step\n        dist_train_step(image_batch, inchi_batch)\n        \n        # end of epoch validation step\n        if stat_logger.current_step == TRAIN_STEPS and epoch%2==0:\n            print(\"\\n... VALIDATION DATASET STATISTICS ... \\n\")\n            \"\"\"\n            for _ in range(VAL_STEPS):\n                preds, lbls = dist_val_step(val_dist_ds)\n                if epoch > EPOCHS - 2:\n                    metrics[\"val_lsd\"].update_state(get_levenshtein_distance(preds, lbls))\n            \"\"\"\n                \n            # Record this epochs statistics\n            stat_logger.train_loss.append(metrics[\"train_loss\"].result().numpy())\n            stat_logger.train_acc.append(metrics[\"train_acc\"].result().numpy())\n            stat_logger.val_loss.append(metrics[\"val_loss\"].result().numpy())\n            stat_logger.val_acc.append(metrics[\"val_acc\"].result().numpy())\n            stat_logger.val_lsd.append(metrics[\"val_lsd\"].result().numpy())\n            stat_logger.step.append(stat_logger.current_step)\n            stat_logger.epoch.append(epoch)\n            #stat_logger.lr.append(lr_scheduler.lr)\n            stat_logger.lr.append(lr_scheduler(tf.cast(stat_logger.current_step+TRAIN_STEPS*(epoch-1), tf.float32)))\n            \n            # Reset the validation metrics as one epoch should not effect the next\n            metrics[\"val_lsd\"].reset_states()\n            metrics[\"val_acc\"].reset_states()\n            metrics[\"val_loss\"].reset_states()\n            metrics[\"train_acc\"].reset_states()\n            metrics[\"train_loss\"].reset_states()\n            metrics[\"batch_loss\"].reset_states()\n            \n            # Print validation scores\n            stat_logger.print_last_val(current_time=time.time())\n        \n        # verbose logging step\n        if stat_logger.current_step % stat_logger.verbose_frequency == 0:    \n            stat_logger.print_current_train(\n                stat_logger.current_step,\n                metrics[\"train_acc\"].result().numpy(), \n                metrics[\"train_loss\"].result().numpy(), \n                metrics[\"batch_loss\"].result().numpy(), \n                current_time=time.time(),\n                #current_lr=lr_scheduler.lr\n                current_lr=lr_scheduler(tf.cast(stat_logger.current_step+TRAIN_STEPS*(epoch-1), tf.float32))\n            )\n            metrics[\"train_acc\"].reset_states()\n            metrics[\"train_loss\"].reset_states()\n            metrics[\"batch_loss\"].reset_states()\n\n        # stop training when NaN loss is detected\n        if stat_logger.current_step == TRAIN_STEPS:\n            break\n            \n        # update learning rate\n        #lr_scheduler.step(stat_logger.current_step+((epoch-1)*TRAIN_STEPS))\n        \n    # Save every other epoch (starting with first epoch)\n    # Save after last epoch too...\n    # if epoch%2==1 or epoch==EPOCHS:\n    # save weights\n    print(\"\\n...SAVING MODELS TO DISK ... \\n\")\n    #encoder.save_weights(f'./encoder_epoch_{epoch}.h5')\n    #decoder.save_weights(f'./decoder_epoch_{epoch}.h5’）\n    transformer.save_weights(f'./transformer_epoch_{epoch}.h5')","metadata":{"execution":{"iopub.status.busy":"2021-06-05T03:14:54.990896Z","iopub.execute_input":"2021-06-05T03:14:54.991305Z"},"trusted":true},"execution_count":null,"outputs":[]}]}