{"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":"
\n\n

Tensorflow/Keras Code Implementation
CLIP (Contrastive Learning Image Pretraining

\n
CREATED BY: DARIEN SCHETTLER
\n\n
\n\n---\n\n
\n\n**LINKS TO THE OTHER [MORE UNINTELLIGIBLE] PARTS:**\n* PART 1 - TFRecord Creation From COCO Dataset\n* PART 2 - Training (This Notebook)\n* PART 3 - Inference

\n\n---\n\n
\n\n
\n πŸ‘   THIS NOTEBOOK IS HEAVILY INSPIRED BY THIS KERAS IMPLEMENTATION   πŸ‘
πŸ‘   AUTHORED BY KHALID SALAMA   πŸ‘\n
\n\n
\n πŸ‘   IF YOU FORK THIS OR FIND THIS HELPFUL   πŸ‘

PLEASE UPVOTE!

This was a lot of work for me and while it may seem silly, it makes me feel appreciated when others like my work. πŸ˜…\n
\n\n\n","metadata":{}},{"cell_type":"markdown","source":"

\n\n

\n\n

TABLE OF CONTENTS

\n\n---\n\n

0    IMPORTS

\n\n---\n\n

1    BACKGROUND INFORMATION

\n\n---\n\n

2    SETUP

\n\n---\n\n

3    HELPER FUNCTIONS

\n\n---\n\n

4    DATASET CREATION

\n\n---\n\n

5    CLIP IMPLEMENTATION

\n\n---\n\n

6    MODEL TRAINING

\n\n---","metadata":{}},{"cell_type":"markdown","source":"
\n\n\n\n

0  IMPORTS    ⤒

","metadata":{}},{"cell_type":"code","source":"!pip install \"tensorflow-text==2.6.*\"\n\n# Machine Learning and Data Science Imports\nimport tensorflow as tf; print(f\"\\t\\t– TENSORFLOW VERSION: {tf.__version__}\"); tf.get_logger().setLevel('INFO')\nimport tensorflow_text; print(f\"\\t\\t– TENSORFLOW TEXT VERSION: {tensorflow_text.__version__}\");\nimport tensorflow_hub as tfhub; print(f\"\\t\\t– TENSORFLOW HUB VERSION: {tfhub.__version__}\");\nimport tensorflow_addons as tfa; print(f\"\\t\\t– TENSORFLOW ADDONS VERSION: {tfa.__version__}\");\nimport pandas as pd; pd.options.mode.chained_assignment = None;\nimport numpy as np; print(f\"\\t\\t– NUMPY VERSION: {np.__version__}\");\nimport sklearn; print(f\"\\t\\t– SKLEARN VERSION: {sklearn.__version__}\");\nfrom sklearn.preprocessing import RobustScaler, PolynomialFeatures\nfrom sklearn.model_selection import GroupKFold, StratifiedKFold\n\n# Built In Imports\nfrom kaggle_datasets import KaggleDatasets\nfrom collections import Counter\nfrom datetime import datetime\nfrom zipfile import ZipFile\nfrom glob import glob\nimport warnings\nimport requests\nimport hashlib\nimport imageio\nimport IPython\nimport sklearn\nimport urllib\nimport zipfile\nimport pickle\nimport random\nimport shutil\nimport string\nimport json\nimport math\nimport time\nimport gzip\nimport ast\nimport sys\nimport io\nimport os\nimport gc\nimport re\n\n# Visualization Imports\nfrom matplotlib.colors import ListedColormap\nfrom matplotlib.patches import Rectangle\nimport matplotlib.patches as patches\nimport plotly.graph_objects as go\nimport matplotlib.pyplot as plt\nfrom tqdm.notebook import tqdm; tqdm.pandas();\nimport plotly.express as px\nimport tifffile as tif\nimport seaborn as sns\nfrom PIL import Image, ImageEnhance; Image.MAX_IMAGE_PIXELS = 5_000_000_000;\nimport matplotlib; print(f\"\\t\\t– MATPLOTLIB VERSION: {matplotlib.__version__}\");\nfrom matplotlib import animation, rc; rc('animation', html='jshtml')\nimport plotly\nimport PIL\nimport cv2\n\nimport plotly.io as pio\nprint(pio.renderers)\n\n# Basic helpers and seeding\ndef seed_it_all(seed=7):\n \"\"\" Attempt to be Reproducible \"\"\"\n os.environ['PYTHONHASHSEED'] = str(seed)\n random.seed(seed)\n np.random.seed(seed)\n tf.random.set_seed(seed)\nseed_it_all()\n\ndef 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\ndef load_json_to_dict(json_path):\n \"\"\" tbd \"\"\"\n with open(json_path) as json_file:\n data = json.load(json_file)\n return data\n \nprint(\"\\n\\n... IMPORTS COMPLETE ...\\n\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-08-01T16:45:17.731480Z","iopub.execute_input":"2022-08-01T16:45:17.731794Z","iopub.status.idle":"2022-08-01T16:45:57.259185Z","shell.execute_reply.started":"2022-08-01T16:45:17.731727Z","shell.execute_reply":"2022-08-01T16:45:57.257962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n\n\n\n

1  BACKGROUND INFORMATION    ⤒

\n\n---\n","metadata":{}},{"cell_type":"markdown","source":"

1.1 GENERAL OVERVIEW

\n\n---\n\n**Abstract Excerpt:**\n> CLIP (Contrastive Language-Image Pre-Training) is a neural network trained on a variety of (image, text) pairs. It can be instructed in natural language to predict the most relevant text snippet, given an image, without directly optimizing for the task, similarly to the zero-shot capabilities of GPT-2 and 3. We found CLIP matches the performance of the original ResNet50 on ImageNet β€œzero-shot” without using any of the original 1.28M labeled examples, overcoming several major challenges in computer vision.\n\n
\n\n**Here are some relevant links:**\n* [[Blog]](https://openai.com/blog/clip/) \n* [[Paper]](https://arxiv.org/abs/2103.00020) \n* [[Model Card]](model-card.md) \n* [[Colab]](https://colab.research.google.com/github/openai/clip/blob/master/notebooks/Interacting_with_CLIP.ipynb)","metadata":{}},{"cell_type":"markdown","source":"

1.2 VISUAL OVERVIEW

\n\n---\n\n![CLIP](https://github.com/openai/CLIP/raw/main/CLIP.png)","metadata":{}},{"cell_type":"markdown","source":"
\n\n\n\n

2  SETUP    ⤒

\n\n---\n","metadata":{}},{"cell_type":"markdown","source":"

2.1 ACCELERATOR DETECTION

\n\n---\n\nIn order to use **`TPU`**, we use **`TPUClusterResolver`** for the initialization which is necessary to connect to the remote cluster and initialize cloud TPUs. Let's go over two important points\n\n1. When using TPU on Kaggle, you don't need to specify arguments for **`TPUClusterResolver`**\n2. However, on **G**oogle **C**ompute **E**ngine (**GCE**), you will need to do the following:\n\n
\n\n```python\n# The name you gave to the TPU to use\nTPU_WORKER = 'my-tpu-name'\n\n# or you can also specify the grpc path directly\n# TPU_WORKER = 'grpc://xxx.xxx.xxx.xxx:8470'\n\n# The zone you chose when you created the TPU to use on GCP.\nZONE = 'us-east1-b'\n\n# The name of the GCP project where you created the TPU to use on GCP.\nPROJECT = 'my-tpu-project'\n\ntpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=TPU_WORKER, zone=ZONE, project=PROJECT)\n```\n\n
\n πŸ›‘   WARNING:

- Although the Tensorflow documentation says it is the project name that should be provided for the argument `project`, it is actually the Project ID, that you should provide. This can be found on the GCP project dashboard page.
\n
\n\n
\n πŸ“–   REFERENCES:

\n - Guide - Use TPUs
\n - Doc - TPUClusterResolver
\n\n
","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 physical_devices = tf.config.list_physical_devices('GPU')\n try: tf.config.experimental.set_memory_growth(physical_devices[0], True)\n except: pass # Invalid device or cannot modify virtual devices once initialized.\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":"2022-08-01T16:45:57.261115Z","iopub.execute_input":"2022-08-01T16:45:57.261389Z","iopub.status.idle":"2022-08-01T16:45:57.349170Z","shell.execute_reply.started":"2022-08-01T16:45:57.261334Z","shell.execute_reply":"2022-08-01T16:45:57.348019Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

2.2 DATASET ACCESS

\n\n---\n\nTPUs read data must be read directly from **G**oogle **C**loud **S**torage **(GCS)**. Kaggle provides a utility library – **`KaggleDatasets`** – which has a utility function **`.get_gcs_path`** that will allow us to access the location of our input datasets within **GCS**.

\n\n
\n πŸ“Œ   TIPS:

- If you have multiple datasets attached to the notebook, you should pass the name of a specific dataset to the `get_gcs_path()` function. In our case, the name of the dataset is the name of the directory the dataset is mounted within.

\n
","metadata":{}},{"cell_type":"code","source":"print(\"\\n... DATA ACCESS SETUP STARTED ...\\n\")\n\nif TPU:\n # Google Cloud Dataset path to training and validation images\n DATA_DIR = KaggleDatasets().get_gcs_path('coco-2017-tfrecords')\n save_locally = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')\n load_locally = tf.saved_model.LoadOptions(experimental_io_device='/job:localhost')\nelse:\n # Local path to training and validation images\n DATA_DIR = \"/kaggle/input/coco-2017-tfrecords\"\n save_locally = None\n load_locally = None\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":"2022-08-01T16:45:57.350598Z","iopub.execute_input":"2022-08-01T16:45:57.351061Z","iopub.status.idle":"2022-08-01T16:45:57.366105Z","shell.execute_reply.started":"2022-08-01T16:45:57.351021Z","shell.execute_reply":"2022-08-01T16:45:57.365420Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

2.3 LEVERAGING XLA OPTIMIZATIONS

\n\n---\n\n\n**XLA** (Accelerated Linear Algebra) is a domain-specific compiler for linear algebra that can accelerate TensorFlow models with potentially no source code changes. **The results are improvements in speed and memory usage**.\n\n
\n\nWhen a TensorFlow program is run, all of the operations are executed individually by the TensorFlow executor. Each TensorFlow operation has a precompiled GPU/TPU kernel implementation that the executor dispatches to.\n\nXLA provides us with an alternative mode of running models: it compiles the TensorFlow graph into a sequence of computation kernels generated specifically for the given model. Because these kernels are unique to the model, they can exploit model-specific information for optimization.

\n\n
\n πŸ›‘   WARNING:

- XLA can not currently compile functions where dimensions are not inferrable: that is, if it's not possible to infer the dimensions of all tensors without running the entire computation
\n
\n\n
\n πŸ“Œ   NOTE:

- XLA compilation is only applied to code that is compiled into a graph (in TF2 that's only a code inside tf.function).
- The jit_compile API has must-compile semantics, i.e. either the entire function is compiled with XLA, or an errors.InvalidArgumentError exception is thrown)\n
\n\n
\n πŸ“–   REFERENCE:

- XLA: Optimizing Compiler for Machine Learning
\n
","metadata":{}},{"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":"2022-08-01T16:45:57.368862Z","iopub.execute_input":"2022-08-01T16:45:57.369509Z","iopub.status.idle":"2022-08-01T16:45:57.377811Z","shell.execute_reply.started":"2022-08-01T16:45:57.369474Z","shell.execute_reply":"2022-08-01T16:45:57.376988Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

2.4 BASIC DATA DEFINITIONS & INITIALIZATIONS

\n\n---\n","metadata":{}},{"cell_type":"code","source":"# Dataset image shape\nIMG_DIM = int([x for x in os.listdir(DATA_DIR) if not x.startswith(\"_\")][0].split(\"x\")[0])\nIMG_SHAPE = (IMG_DIM, IMG_DIM, 3)\n\n# Batching and shuffling\nREPLICA_BATCH_SIZE = 32\nBATCH_SIZE = N_REPLICAS*REPLICA_BATCH_SIZE\nSHUFFLE_BUFFER = 10*BATCH_SIZE\n\n# Collection of tfrecords\nTRAIN_TFRECORDS = tf.io.gfile.glob(os.path.join(DATA_DIR, f\"{IMG_DIM}x{IMG_DIM}\", \"train\", \"*.tfrec\"))\nVAL_TFRECORDS = tf.io.gfile.glob(os.path.join(DATA_DIR, f\"{IMG_DIM}x{IMG_DIM}\", \"val\", \"*.tfrec\"))\n\nVAL_TFRECORDS","metadata":{"execution":{"iopub.status.busy":"2022-08-01T16:45:57.378969Z","iopub.execute_input":"2022-08-01T16:45:57.379330Z","iopub.status.idle":"2022-08-01T16:45:57.427014Z","shell.execute_reply.started":"2022-08-01T16:45:57.379232Z","shell.execute_reply":"2022-08-01T16:45:57.426390Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n\n\n\n\n\n

\n 3  HELPER FUNCTION & CLASSES    ⤒\n

\n\n---","metadata":{}},{"cell_type":"code","source":"def flatten_l_o_l(nested_list):\n \"\"\" Flatten a list of lists \"\"\"\n nested_list = [x if type(x) is list else [x,] for x in nested_list]\n return [item for sublist in nested_list for item in sublist]\n\n\ndef decode_image(image_data, resize_to=(224,224,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 to_norm_val (int, optional): The value that will allow us to normalize the tensor to 0-1\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, dtype=tf.uint8)\n return tf.reshape(image, resize_to)\n\n\ndef decode_example(serialized_example, resize_to=(224,224,3)):\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 – 'label'\n is_test (bool, optional): Whether to allow for the label 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 'caption': tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value=''),\n }\n \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)\n caption = features[\"caption\"]\n return (image, caption)","metadata":{"execution":{"iopub.status.busy":"2022-08-01T16:45:57.428112Z","iopub.execute_input":"2022-08-01T16:45:57.428349Z","iopub.status.idle":"2022-08-01T16:45:57.440597Z","shell.execute_reply.started":"2022-08-01T16:45:57.428317Z","shell.execute_reply":"2022-08-01T16:45:57.439864Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n\n\n\n\n\n

\n 4  DATASET CREATION    ⤒\n

\n\n---","metadata":{}},{"cell_type":"code","source":"train_ds = tf.data.TFRecordDataset(TRAIN_TFRECORDS, num_parallel_reads=tf.data.AUTOTUNE)\\\n .map(lambda x: decode_example(x, IMG_SHAPE), num_parallel_calls=tf.data.AUTOTUNE)\\\n .shuffle(SHUFFLE_BUFFER)\\\n .batch(BATCH_SIZE, drop_remainder=True, num_parallel_calls=tf.data.AUTOTUNE)\\\n .prefetch(tf.data.AUTOTUNE)\n\nval_ds = tf.data.TFRecordDataset(VAL_TFRECORDS, num_parallel_reads=tf.data.AUTOTUNE)\\\n .map(lambda x: decode_example(x, IMG_SHAPE), num_parallel_calls=tf.data.AUTOTUNE)\\\n .shuffle(BATCH_SIZE*2)\\\n .batch(BATCH_SIZE, drop_remainder=True, num_parallel_calls=tf.data.AUTOTUNE)\\\n .prefetch(tf.data.AUTOTUNE)\n\nprint(\"TRAIN DATASET : \", train_ds)\nprint(\"VAL DATASET : \", val_ds)","metadata":{"execution":{"iopub.status.busy":"2022-08-01T16:45:57.441809Z","iopub.execute_input":"2022-08-01T16:45:57.442104Z","iopub.status.idle":"2022-08-01T16:46:00.085558Z","shell.execute_reply.started":"2022-08-01T16:45:57.442069Z","shell.execute_reply":"2022-08-01T16:46:00.084339Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n\n\n\n\n\n

\n 5  CLIP IMPLEMENTATION    ⤒\n

\n\n---","metadata":{}},{"cell_type":"markdown","source":"

5.1 THE PROJECTION HEAD

\n\n---\n\nThe projection head is used to transform the image and the text embeddings to the same embedding space with the same dimensionality.","metadata":{}},{"cell_type":"code","source":"def project_embeddings(embeddings, num_projection_layers, projection_dims, dropout_rate):\n \"\"\"\n Transform a given embedding (vision or text) and project \n into a shared embedding dimensional space\n \n Args:\n embeddings (tf.constant): TBD\n num_projection_layers (int): TBD\n projection_dims (int): TBD\n dropout_rate (float): TBD\n \n Returns:\n Embedding projected into shared dimensional space\n \n \"\"\"\n \n projected_embeddings = tf.keras.layers.Dense(units=projection_dims)(embeddings)\n for _ in range(num_projection_layers):\n x = tf.nn.gelu(projected_embeddings)\n x = tf.keras.layers.Dense(projection_dims)(x)\n x = tf.keras.layers.Dropout(dropout_rate)(x)\n x = tf.keras.layers.Add()([projected_embeddings, x])\n projected_embeddings = tf.keras.layers.LayerNormalization()(x)\n return projected_embeddings","metadata":{"execution":{"iopub.status.busy":"2022-08-01T16:46:00.086874Z","iopub.execute_input":"2022-08-01T16:46:00.087117Z","iopub.status.idle":"2022-08-01T16:46:00.094228Z","shell.execute_reply.started":"2022-08-01T16:46:00.087082Z","shell.execute_reply":"2022-08-01T16:46:00.093410Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

5.2 THE VISION ENCODER

\n\n---\n\nYou can use any vision encoder here. The point is for this to be completely modular. \n* We will use an EfficientNet model\n* We will freeze the majority of the model to accelerate training and prevent information loss","metadata":{}},{"cell_type":"code","source":"def create_vision_encoder(backbone_model, pp_fn=None, \n num_projection_layers=1, \n projection_dims=64, \n post_backbone_dropout=0.4,\n dropout_rate=0.2):\n\n # Receive the images as inputs.\n _inputs = tf.keras.layers.Input(shape=IMG_SHAPE, name=\"image_input\")\n \n # Preprocess the input image if required\n if pp_fn is not None: x = pp_fn(_inputs)\n else: x = _inputs\n \n # Generate the embeddings for the images using the xception model.\n embeddings = backbone_model(x)\n \n if post_backbone_dropout is not None and type(post_backbone_dropout) is float:\n embeddings = tf.keras.layers.Dropout(post_backbone_dropout)(embeddings)\n \n # Project the embeddings produced by the model.\n _outputs = project_embeddings(embeddings, num_projection_layers, projection_dims, dropout_rate)\n \n # Create the vision encoder model.\n return tf.keras.Model(inputs=_inputs, outputs=_outputs, name=\"vision_encoder\")\n\nwith strategy.scope():\n # EfficientNetB5 Pretrained on ImageNet1K\n bb_model = tf.keras.applications.EfficientNetB5(include_top=False, weights=\"imagenet\", pooling=\"avg\")\n bb_preprocessing_fn = tf.keras.applications.efficientnet.preprocess_input\n\n # Freeze all but the \"top\" layers\n for _l in bb_model.layers:\n if not _l.name.startswith(\"top\"):\n _l.trainable=False\n\n # Get our vision encoder\n vision_encoder = create_vision_encoder(bb_model, pp_fn=bb_preprocessing_fn,\n num_projection_layers=1, projection_dims=64, \n dropout_rate=0.2, post_backbone_dropout=0.4)\n\nvision_encoder.summary()","metadata":{"execution":{"iopub.status.busy":"2022-08-01T16:46:00.095873Z","iopub.execute_input":"2022-08-01T16:46:00.096128Z","iopub.status.idle":"2022-08-01T16:46:06.558140Z","shell.execute_reply.started":"2022-08-01T16:46:00.096093Z","shell.execute_reply":"2022-08-01T16:46:06.557436Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

5.3 THE TEXT ENCODER

\n\n---\n\nYou can use any language/text encoder here. The point is for this to be completely modular. \n* We will use a sentence encoder\n* We will freeze the entire model","metadata":{}},{"cell_type":"code","source":"def create_text_encoder(backbone_model, pp_fn=None,\n num_projection_layers=1, \n projection_dims=64, \n dropout_rate=0.2):\n \n # Receive the text as inputs.\n \n \n # Preprocess the text if required\n _inputs = tf.keras.layers.Input(shape=(), dtype=tf.string, name=\"text_input\")\n \n if pp_fn is not None: x = pp_fn(_inputs)\n else: x = _inputs\n \n embeddings = backbone_model(x)\n \n # Project the embeddings produced by the model.\n _outputs = project_embeddings(embeddings, num_projection_layers, projection_dims, dropout_rate)\n \n # Create the text encoder model.\n return tf.keras.Model(inputs=_inputs, outputs=_outputs, name=\"text_encoder\")\n\nwith strategy.scope():\n # Load the model and freeze weights\n bb_model = tfhub.KerasLayer(tfhub.load(\"https://tfhub.dev/google/universal-sentence-encoder/4\"))\n bb_model.trainable=False\n \n # No preprocessing required\n bb_preprocessing_fn = None\n \n text_encoder = create_text_encoder(bb_model, pp_fn=bb_preprocessing_fn,\n num_projection_layers=1, \n projection_dims=64, \n dropout_rate=0.2)\n \ntext_encoder.summary()","metadata":{"execution":{"iopub.status.busy":"2022-08-01T16:46:06.563037Z","iopub.execute_input":"2022-08-01T16:46:06.564981Z","iopub.status.idle":"2022-08-01T16:46:26.500884Z","shell.execute_reply.started":"2022-08-01T16:46:06.564942Z","shell.execute_reply":"2022-08-01T16:46:26.500098Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"

5.4 THE DUAL ENCODER

\n\n---\n\nWe need to combine our two modalities into a shared space\n\n> To calculate the loss, we compute the pairwise dot-product similarity between each $caption_i$ and $images_j$ in the batch as the predictions. The target similarity between $caption_i$ and $image_j$ is computed as the average of the (dot-product similarity between $caption_i$ and $caption_j$) and (the dot-product similarity between $image_i$ and $image_j$). Then, we use crossentropy to compute the loss between the targets and the predictions.","metadata":{}},{"cell_type":"code","source":"class DualEncoder(tf.keras.Model):\n def __init__(self, text_encoder, vision_encoder, temperature=0.075, **kwargs):\n super(DualEncoder, self).__init__(**kwargs)\n self.text_encoder = text_encoder\n self.vision_encoder = vision_encoder\n self.temperature = temperature\n self.loss_tracker = tf.keras.metrics.Mean(name=\"loss\")\n\n @property\n def metrics(self):\n return [self.loss_tracker]\n\n # Ensure second feature is text and first is image\n def call(self, features, training=False):\n \n # Get the embeddings for the images.\n image_embeddings = vision_encoder(features[0], training=training)\n \n # Get the embeddings for the captions.\n caption_embeddings = text_encoder(features[1], training=training)\n\n return caption_embeddings, image_embeddings\n\n def compute_loss(self, caption_embeddings, image_embeddings):\n # logits[i][j] is the dot_similarity(caption_i, image_j).\n logits = (tf.matmul(caption_embeddings, image_embeddings, transpose_b=True) / self.temperature)\n \n # images_similarity[i][j] is the dot_similarity(image_i, image_j).\n images_similarity = tf.matmul(image_embeddings, image_embeddings, transpose_b=True)\n \n # captions_similarity[i][j] is the dot_similarity(caption_i, caption_j).\n captions_similarity = tf.matmul(caption_embeddings, caption_embeddings, transpose_b=True)\n \n # targets[i][j] = avarage dot_similarity(caption_i, caption_j) and dot_similarity(image_i, image_j).\n targets = tf.keras.activations.softmax((captions_similarity + images_similarity) / (2 * self.temperature))\n \n # Compute the loss for the captions using crossentropy\n captions_loss = tf.keras.losses.categorical_crossentropy(y_true=targets, y_pred=logits, from_logits=True)\n \n # Compute the loss for the images using crossentropy\n images_loss = tf.keras.losses.categorical_crossentropy(y_true=tf.transpose(targets), y_pred=tf.transpose(logits), from_logits=True)\n \n # Return the mean of the loss over the batch.\n return (captions_loss + images_loss) / 2\n\n def train_step(self, features):\n with tf.GradientTape() as tape:\n # Forward pass\n caption_embeddings, image_embeddings = self(features, training=True)\n loss = self.compute_loss(caption_embeddings, image_embeddings)\n \n # Backward pass\n gradients = tape.gradient(loss, self.trainable_variables)\n self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))\n \n # Monitor loss\n self.loss_tracker.update_state(loss)\n return {\"loss\": self.loss_tracker.result()}\n\n def test_step(self, features):\n caption_embeddings, image_embeddings = self(features, training=False)\n loss = self.compute_loss(caption_embeddings, image_embeddings)\n self.loss_tracker.update_state(loss)\n return {\"loss\": self.loss_tracker.result()}\n\n# I really like 75\nwith strategy.scope():\n dual_encoder = DualEncoder(text_encoder, vision_encoder, temperature=0.075)\n dual_encoder.compile(optimizer=tfa.optimizers.AdamW(learning_rate=0.0005, weight_decay=0.00075))","metadata":{"execution":{"iopub.status.busy":"2022-08-01T18:15:57.517205Z","iopub.execute_input":"2022-08-01T18:15:57.517827Z","iopub.status.idle":"2022-08-01T18:15:57.560037Z","shell.execute_reply.started":"2022-08-01T18:15:57.517769Z","shell.execute_reply":"2022-08-01T18:15:57.558617Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"
\n\n\n\n\n\n

\n 6  MODEL TRAINING    ⤒\n

\n\n---\n\nTrain the dual encoder mode","metadata":{}},{"cell_type":"code","source":"# Train for up to N epochs\nN_EPOCHS=25\n\n# Cleanup callback (stupid TF...)\nclass GCCallback(tf.keras.callbacks.Callback):\n def on_epoch_end(self, epoch, logs=None):\n gc.collect(); tf.keras.backend.clear_session(); gc.collect()\n\n# Create a learning rate scheduler callback.\nlr_cb = tf.keras.callbacks.ReduceLROnPlateau(monitor=\"val_loss\", factor=0.75, patience=2)\n\n# Create an early stopping callback.\nes_cb = tf.keras.callbacks.EarlyStopping(monitor=\"val_loss\", patience=4, restore_best_weights=True)\n\n# Train the model\nprint(\"\\n\\n ... TRAINING STARTED ...\\n\")\nhistory = dual_encoder.fit(train_ds, validation_data=val_ds, epochs=N_EPOCHS, callbacks=[lr_cb, es_cb, GCCallback()])\nprint(\"\\n... TRAINING COMPLETE ...\\n\")\n\n# Save the models\nprint(\"\\n\\n ... SAVING STARTED ...\\n\")\nprint(f\"\\t--> Saving vision encoder --> /kaggle/working/vision_encoder_{IMG_DIM}\")\nvision_encoder.save(f\"vision_encoder_{IMG_DIM}\")\nprint(f\"\\t--> Saving text encoder --> /kaggle/working/text_encoder_{IMG_DIM}\")\ntext_encoder.save(f\"text_encoder_{IMG_DIM}\")\nprint(\"\\n ... SAVING COMPLETE ...\\n\")","metadata":{"execution":{"iopub.status.busy":"2022-08-01T18:16:11.680403Z","iopub.execute_input":"2022-08-01T18:16:11.680986Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}