Image Captioning Transformers.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "image_captioning",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.0"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9CpLds8XuJIn"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tx9QfKTBuJIn"
      },
      "source": [
        "import os\n",
        "import re\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n",
        "from tensorflow.keras.applications import efficientnet\n",
        "from tensorflow.keras.layers import TextVectorization\n",
        "\n",
        "\n",
        "seed = 111\n",
        "np.random.seed(seed)\n",
        "tf.random.set_seed(seed)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s7DKOIHquJIp"
      },
      "source": [
        "!wget -q https://github.com/jbrownlee/Datasets/releases/download/Flickr8k/Flickr8k_Dataset.zip\n",
        "!wget -q https://github.com/jbrownlee/Datasets/releases/download/Flickr8k/Flickr8k_text.zip\n",
        "!unzip -qq Flickr8k_Dataset.zip\n",
        "!unzip -qq Flickr8k_text.zip\n",
        "!rm Flickr8k_Dataset.zip Flickr8k_text.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fnMDYpfeuJIp"
      },
      "source": [
        "\n",
        "# Path to the images\n",
        "IMAGES_PATH = \"Flicker8k_Dataset\"\n",
        "\n",
        "# Desired image dimensions\n",
        "IMAGE_SIZE = (299, 299)\n",
        "\n",
        "# Vocabulary size\n",
        "VOCAB_SIZE = 10000\n",
        "\n",
        "# Fixed length allowed for any sequence\n",
        "SEQ_LENGTH = 25\n",
        "\n",
        "# Dimension for the image embeddings and token embeddings\n",
        "EMBED_DIM = 512\n",
        "\n",
        "# Per-layer units in the feed-forward network\n",
        "FF_DIM = 512\n",
        "\n",
        "# Other training parameters\n",
        "BATCH_SIZE = 64\n",
        "EPOCHS = 30\n",
        "AUTOTUNE = tf.data.AUTOTUNE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JjSBa1VpuJIq"
      },
      "source": [
        "## Preparing the dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y7oSmsU1uJIq"
      },
      "source": [
        "\n",
        "def load_captions_data(filename):\n",
        "    with open(filename) as caption_file:\n",
        "        caption_data = caption_file.readlines()\n",
        "        caption_mapping = {}\n",
        "        text_data = []\n",
        "        images_to_skip = set()\n",
        "\n",
        "        for line in caption_data:\n",
        "            line = line.rstrip(\"\\n\")\n",
        "            # Image name and captions are separated using a tab\n",
        "            img_name, caption = line.split(\"\\t\")\n",
        "\n",
        "            # Each image is repeated five times for the five different captions.\n",
        "            # Each image name has a suffix `#(caption_number)`\n",
        "            img_name = img_name.split(\"#\")[0]\n",
        "            img_name = os.path.join(IMAGES_PATH, img_name.strip())\n",
        "\n",
        "            # We will remove caption that are either too short to too long\n",
        "            tokens = caption.strip().split()\n",
        "\n",
        "            if len(tokens) < 5 or len(tokens) > SEQ_LENGTH:\n",
        "                images_to_skip.add(img_name)\n",
        "                continue\n",
        "\n",
        "            if img_name.endswith(\"jpg\") and img_name not in images_to_skip:\n",
        "                # We will add a start and an end token to each caption\n",
        "                caption = \"<start> \" + caption.strip() + \" <end>\"\n",
        "                text_data.append(caption)\n",
        "\n",
        "                if img_name in caption_mapping:\n",
        "                    caption_mapping[img_name].append(caption)\n",
        "                else:\n",
        "                    caption_mapping[img_name] = [caption]\n",
        "\n",
        "        for img_name in images_to_skip:\n",
        "            if img_name in caption_mapping:\n",
        "                del caption_mapping[img_name]\n",
        "\n",
        "        return caption_mapping, text_data\n",
        "\n",
        "\n",
        "def train_val_split(caption_data, train_size=0.8, shuffle=True):\n",
        "    # 1. Get the list of all image names\n",
        "    all_images = list(caption_data.keys())\n",
        "\n",
        "    # 2. Shuffle if necessary\n",
        "    if shuffle:\n",
        "        np.random.shuffle(all_images)\n",
        "\n",
        "    # 3. Split into training and validation sets\n",
        "    train_size = int(len(caption_data) * train_size)\n",
        "\n",
        "    training_data = {\n",
        "        img_name: caption_data[img_name] for img_name in all_images[:train_size]\n",
        "    }\n",
        "    validation_data = {\n",
        "        img_name: caption_data[img_name] for img_name in all_images[train_size:]\n",
        "    }\n",
        "\n",
        "    # 4. Return the splits\n",
        "    return training_data, validation_data\n",
        "\n",
        "\n",
        "# Load the dataset\n",
        "captions_mapping, text_data = load_captions_data(\"Flickr8k.token.txt\")\n",
        "\n",
        "# Split the dataset into training and validation sets\n",
        "train_data, valid_data = train_val_split(captions_mapping)\n",
        "print(\"Number of training samples: \", len(train_data))\n",
        "print(\"Number of validation samples: \", len(valid_data))\n",
        "\n",
        "\n",
        "# Load the dataset\n",
        "captions_mapping, text_data = load_captions_data(\"Flickr8k.token.txt\")\n",
        "\n",
        "# Split the dataset into training and validation sets\n",
        "train_data, valid_data = train_val_split(captions_mapping)\n",
        "print(\"Number of training samples: \", len(train_data))\n",
        "print(\"Number of validation samples: \", len(valid_data))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GitzBvd2uJIs"
      },
      "source": [
        "## Vectorizing the text data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mcuXvJVxuJIs"
      },
      "source": [
        "\n",
        "def custom_standardization(input_string):\n",
        "    lowercase = tf.strings.lower(input_string)\n",
        "    return tf.strings.regex_replace(lowercase, \"[%s]\" % re.escape(strip_chars), \"\")\n",
        "\n",
        "\n",
        "strip_chars = \"!\\\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~\"\n",
        "strip_chars = strip_chars.replace(\"<\", \"\")\n",
        "strip_chars = strip_chars.replace(\">\", \"\")\n",
        "\n",
        "vectorization = TextVectorization(\n",
        "    max_tokens=VOCAB_SIZE,\n",
        "    output_mode=\"int\",\n",
        "    output_sequence_length=SEQ_LENGTH,\n",
        "    standardize=custom_standardization,\n",
        ")\n",
        "vectorization.adapt(text_data)\n",
        "\n",
        "# Data augmentation for image data\n",
        "image_augmentation = keras.Sequential(\n",
        "    [\n",
        "        layers.RandomFlip(\"horizontal\"),\n",
        "        layers.RandomRotation(0.2),\n",
        "        layers.RandomContrast(0.3),\n",
        "    ]\n",
        ")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pLHl5z0zuJIt"
      },
      "source": [
        "## Building a pipeline for training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XV8n5evluJIt"
      },
      "source": [
        "\n",
        "def decode_and_resize(img_path, size=IMAGE_SIZE):\n",
        "    img = tf.io.read_file(img_path)\n",
        "    img = tf.image.decode_jpeg(img, channels=3)\n",
        "    img = tf.image.resize(img, IMAGE_SIZE)\n",
        "    return img\n",
        "\n",
        "\n",
        "def read_train_image(img_path, size=IMAGE_SIZE):\n",
        "    img = decode_and_resize(img_path)\n",
        "    img = image_augmentation(tf.expand_dims(img, 0))[0]\n",
        "    img = tf.image.convert_image_dtype(img, tf.float32)\n",
        "    return img\n",
        "\n",
        "\n",
        "def read_valid_image(img_path, size=IMAGE_SIZE):\n",
        "    img = decode_and_resize(img_path)\n",
        "    img = tf.image.convert_image_dtype(img, tf.float32)\n",
        "    return img\n",
        "\n",
        "\n",
        "def make_dataset(images, captions, split=\"train\"):\n",
        "    if split == \"train\":\n",
        "        img_dataset = tf.data.Dataset.from_tensor_slices(images).map(\n",
        "            read_train_image, num_parallel_calls=AUTOTUNE\n",
        "        )\n",
        "    else:\n",
        "        img_dataset = tf.data.Dataset.from_tensor_slices(images).map(\n",
        "            read_valid_image, num_parallel_calls=AUTOTUNE\n",
        "        )\n",
        "\n",
        "    cap_dataset = tf.data.Dataset.from_tensor_slices(captions).map(\n",
        "        vectorization, num_parallel_calls=AUTOTUNE\n",
        "    )\n",
        "\n",
        "    dataset = tf.data.Dataset.zip((img_dataset, cap_dataset))\n",
        "    dataset = dataset.batch(BATCH_SIZE).shuffle(256).prefetch(AUTOTUNE)\n",
        "    return dataset\n",
        "\n",
        "\n",
        "# Pass the list of images and the list of corresponding captions\n",
        "train_dataset = make_dataset(\n",
        "    list(train_data.keys()), list(train_data.values()), split=\"train\"\n",
        ")\n",
        "\n",
        "valid_dataset = make_dataset(\n",
        "    list(valid_data.keys()), list(valid_data.values()), split=\"valid\"\n",
        ")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7JG1zr9CuJIu"
      },
      "source": [
        "## Building the model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xgRXqjUsuJIu"
      },
      "source": [
        "\n",
        "def get_cnn_model():\n",
        "    base_model = efficientnet.EfficientNetB0(\n",
        "        input_shape=(*IMAGE_SIZE, 3), include_top=False, weights=\"imagenet\",\n",
        "    )\n",
        "    # We freeze our feature extractor\n",
        "    base_model.trainable = False\n",
        "    base_model_out = base_model.output\n",
        "    base_model_out = layers.Reshape((-1, base_model_out.shape[-1]))(base_model_out)\n",
        "    cnn_model = keras.models.Model(base_model.input, base_model_out)\n",
        "    return cnn_model\n",
        "\n",
        "\n",
        "class TransformerEncoderBlock(layers.Layer):\n",
        "    def __init__(self, embed_dim, dense_dim, num_heads, **kwargs):\n",
        "        super().__init__(**kwargs)\n",
        "        self.embed_dim = embed_dim\n",
        "        self.dense_dim = dense_dim\n",
        "        self.num_heads = num_heads\n",
        "        self.attention_1 = layers.MultiHeadAttention(\n",
        "            num_heads=num_heads, key_dim=embed_dim, dropout=0.0\n",
        "        )\n",
        "        self.layernorm_1 = layers.LayerNormalization()\n",
        "        self.layernorm_2 = layers.LayerNormalization()\n",
        "        self.dense_1 = layers.Dense(embed_dim, activation=\"relu\")\n",
        "\n",
        "    def call(self, inputs, training, mask=None):\n",
        "        inputs = self.layernorm_1(inputs)\n",
        "        inputs = self.dense_1(inputs)\n",
        "\n",
        "        attention_output_1 = self.attention_1(\n",
        "            query=inputs,\n",
        "            value=inputs,\n",
        "            key=inputs,\n",
        "            attention_mask=None,\n",
        "            training=training,\n",
        "        )\n",
        "        out_1 = self.layernorm_2(inputs + attention_output_1)\n",
        "        return out_1\n",
        "\n",
        "\n",
        "class PositionalEmbedding(layers.Layer):\n",
        "    def __init__(self, sequence_length, vocab_size, embed_dim, **kwargs):\n",
        "        super().__init__(**kwargs)\n",
        "        self.token_embeddings = layers.Embedding(\n",
        "            input_dim=vocab_size, output_dim=embed_dim\n",
        "        )\n",
        "        self.position_embeddings = layers.Embedding(\n",
        "            input_dim=sequence_length, output_dim=embed_dim\n",
        "        )\n",
        "        self.sequence_length = sequence_length\n",
        "        self.vocab_size = vocab_size\n",
        "        self.embed_dim = embed_dim\n",
        "        self.embed_scale = tf.math.sqrt(tf.cast(embed_dim, tf.float32))\n",
        "\n",
        "    def call(self, inputs):\n",
        "        length = tf.shape(inputs)[-1]\n",
        "        positions = tf.range(start=0, limit=length, delta=1)\n",
        "        embedded_tokens = self.token_embeddings(inputs)\n",
        "        embedded_tokens = embedded_tokens * self.embed_scale\n",
        "        embedded_positions = self.position_embeddings(positions)\n",
        "        return embedded_tokens + embedded_positions\n",
        "\n",
        "    def compute_mask(self, inputs, mask=None):\n",
        "        return tf.math.not_equal(inputs, 0)\n",
        "\n",
        "\n",
        "class TransformerDecoderBlock(layers.Layer):\n",
        "    def __init__(self, embed_dim, ff_dim, num_heads, **kwargs):\n",
        "        super().__init__(**kwargs)\n",
        "        self.embed_dim = embed_dim\n",
        "        self.ff_dim = ff_dim\n",
        "        self.num_heads = num_heads\n",
        "        self.attention_1 = layers.MultiHeadAttention(\n",
        "            num_heads=num_heads, key_dim=embed_dim, dropout=0.1\n",
        "        )\n",
        "        self.attention_2 = layers.MultiHeadAttention(\n",
        "            num_heads=num_heads, key_dim=embed_dim, dropout=0.1\n",
        "        )\n",
        "        self.ffn_layer_1 = layers.Dense(ff_dim, activation=\"relu\")\n",
        "        self.ffn_layer_2 = layers.Dense(embed_dim)\n",
        "\n",
        "        self.layernorm_1 = layers.LayerNormalization()\n",
        "        self.layernorm_2 = layers.LayerNormalization()\n",
        "        self.layernorm_3 = layers.LayerNormalization()\n",
        "\n",
        "        self.embedding = PositionalEmbedding(\n",
        "            embed_dim=EMBED_DIM, sequence_length=SEQ_LENGTH, vocab_size=VOCAB_SIZE\n",
        "        )\n",
        "        self.out = layers.Dense(VOCAB_SIZE, activation=\"softmax\")\n",
        "\n",
        "        self.dropout_1 = layers.Dropout(0.3)\n",
        "        self.dropout_2 = layers.Dropout(0.5)\n",
        "        self.supports_masking = True\n",
        "\n",
        "    def call(self, inputs, encoder_outputs, training, mask=None):\n",
        "        inputs = self.embedding(inputs)\n",
        "        causal_mask = self.get_causal_attention_mask(inputs)\n",
        "\n",
        "        if mask is not None:\n",
        "            padding_mask = tf.cast(mask[:, :, tf.newaxis], dtype=tf.int32)\n",
        "            combined_mask = tf.cast(mask[:, tf.newaxis, :], dtype=tf.int32)\n",
        "            combined_mask = tf.minimum(combined_mask, causal_mask)\n",
        "\n",
        "        attention_output_1 = self.attention_1(\n",
        "            query=inputs,\n",
        "            value=inputs,\n",
        "            key=inputs,\n",
        "            attention_mask=combined_mask,\n",
        "            training=training,\n",
        "        )\n",
        "        out_1 = self.layernorm_1(inputs + attention_output_1)\n",
        "\n",
        "        attention_output_2 = self.attention_2(\n",
        "            query=out_1,\n",
        "            value=encoder_outputs,\n",
        "            key=encoder_outputs,\n",
        "            attention_mask=padding_mask,\n",
        "            training=training,\n",
        "        )\n",
        "        out_2 = self.layernorm_2(out_1 + attention_output_2)\n",
        "\n",
        "        ffn_out = self.ffn_layer_1(out_2)\n",
        "        ffn_out = self.dropout_1(ffn_out, training=training)\n",
        "        ffn_out = self.ffn_layer_2(ffn_out)\n",
        "\n",
        "        ffn_out = self.layernorm_3(ffn_out + out_2, training=training)\n",
        "        ffn_out = self.dropout_2(ffn_out, training=training)\n",
        "        preds = self.out(ffn_out)\n",
        "        return preds\n",
        "\n",
        "    def get_causal_attention_mask(self, inputs):\n",
        "        input_shape = tf.shape(inputs)\n",
        "        batch_size, sequence_length = input_shape[0], input_shape[1]\n",
        "        i = tf.range(sequence_length)[:, tf.newaxis]\n",
        "        j = tf.range(sequence_length)\n",
        "        mask = tf.cast(i >= j, dtype=\"int32\")\n",
        "        mask = tf.reshape(mask, (1, input_shape[1], input_shape[1]))\n",
        "        mult = tf.concat(\n",
        "            [tf.expand_dims(batch_size, -1), tf.constant([1, 1], dtype=tf.int32)],\n",
        "            axis=0,\n",
        "        )\n",
        "        return tf.tile(mask, mult)\n",
        "\n",
        "\n",
        "class ImageCaptioningModel(keras.Model):\n",
        "    def __init__(\n",
        "        self, cnn_model, encoder, decoder, num_captions_per_image=5,\n",
        "    ):\n",
        "        super().__init__()\n",
        "        self.cnn_model = cnn_model\n",
        "        self.encoder = encoder\n",
        "        self.decoder = decoder\n",
        "        self.loss_tracker = keras.metrics.Mean(name=\"loss\")\n",
        "        self.acc_tracker = keras.metrics.Mean(name=\"accuracy\")\n",
        "        self.num_captions_per_image = num_captions_per_image\n",
        "\n",
        "    def calculate_loss(self, y_true, y_pred, mask):\n",
        "        loss = self.loss(y_true, y_pred)\n",
        "        mask = tf.cast(mask, dtype=loss.dtype)\n",
        "        loss *= mask\n",
        "        return tf.reduce_sum(loss) / tf.reduce_sum(mask)\n",
        "\n",
        "    def calculate_accuracy(self, y_true, y_pred, mask):\n",
        "        accuracy = tf.equal(y_true, tf.argmax(y_pred, axis=2))\n",
        "        accuracy = tf.math.logical_and(mask, accuracy)\n",
        "        accuracy = tf.cast(accuracy, dtype=tf.float32)\n",
        "        mask = tf.cast(mask, dtype=tf.float32)\n",
        "        return tf.reduce_sum(accuracy) / tf.reduce_sum(mask)\n",
        "\n",
        "    def _compute_caption_loss_and_acc(self, img_embed, batch_seq, training=True):\n",
        "        encoder_out = self.encoder(img_embed, training=training)\n",
        "        batch_seq_inp = batch_seq[:, :-1]\n",
        "        batch_seq_true = batch_seq[:, 1:]\n",
        "        mask = tf.math.not_equal(batch_seq_true, 0)\n",
        "        batch_seq_pred = self.decoder(\n",
        "            batch_seq_inp, encoder_out, training=training, mask=mask\n",
        "        )\n",
        "        loss = self.calculate_loss(batch_seq_true, batch_seq_pred, mask)\n",
        "        acc = self.calculate_accuracy(batch_seq_true, batch_seq_pred, mask)\n",
        "        return loss, acc\n",
        "\n",
        "    def train_step(self, batch_data):\n",
        "        batch_img, batch_seq = batch_data\n",
        "        batch_loss = 0\n",
        "        batch_acc = 0\n",
        "\n",
        "        # 1. Get image embeddings\n",
        "        img_embed = self.cnn_model(batch_img)\n",
        "\n",
        "        # 2. Pass each of the five captions one by one to the decoder\n",
        "        # along with the encoder outputs and compute the loss as well as accuracy\n",
        "        # for each caption.\n",
        "        for i in range(self.num_captions_per_image):\n",
        "            with tf.GradientTape() as tape:\n",
        "                loss, acc = self._compute_caption_loss_and_acc(\n",
        "                    img_embed, batch_seq[:, i, :], training=True\n",
        "                )\n",
        "\n",
        "                # 3. Update loss and accuracy\n",
        "                batch_loss += loss\n",
        "                batch_acc += acc\n",
        "\n",
        "            # 4. Get the list of all the trainable weights\n",
        "            train_vars = (\n",
        "                self.encoder.trainable_variables + self.decoder.trainable_variables\n",
        "            )\n",
        "\n",
        "            # 5. Get the gradients\n",
        "            grads = tape.gradient(loss, train_vars)\n",
        "\n",
        "            # 6. Update the trainable weights\n",
        "            self.optimizer.apply_gradients(zip(grads, train_vars))\n",
        "\n",
        "        # 7. Update the trackers\n",
        "        batch_acc /= float(self.num_captions_per_image)\n",
        "        self.loss_tracker.update_state(batch_loss)\n",
        "        self.acc_tracker.update_state(batch_acc)\n",
        "\n",
        "        # 8. Return the loss and accuracy values\n",
        "        return {\"loss\": self.loss_tracker.result(), \"acc\": self.acc_tracker.result()}\n",
        "\n",
        "    def test_step(self, batch_data):\n",
        "        batch_img, batch_seq = batch_data\n",
        "        batch_loss = 0\n",
        "        batch_acc = 0\n",
        "\n",
        "        # 1. Get image embeddings\n",
        "        img_embed = self.cnn_model(batch_img)\n",
        "\n",
        "        # 2. Pass each of the five captions one by one to the decoder\n",
        "        # along with the encoder outputs and compute the loss as well as accuracy\n",
        "        # for each caption.\n",
        "        for i in range(self.num_captions_per_image):\n",
        "            loss, acc = self._compute_caption_loss_and_acc(\n",
        "                img_embed, batch_seq[:, i, :], training=False\n",
        "            )\n",
        "\n",
        "            # 3. Update batch loss and batch accuracy\n",
        "            batch_loss += loss\n",
        "            batch_acc += acc\n",
        "\n",
        "        batch_acc /= float(self.num_captions_per_image)\n",
        "\n",
        "        # 4. Update the trackers\n",
        "        self.loss_tracker.update_state(batch_loss)\n",
        "        self.acc_tracker.update_state(batch_acc)\n",
        "\n",
        "        # 5. Return the loss and accuracy values\n",
        "        return {\"loss\": self.loss_tracker.result(), \"acc\": self.acc_tracker.result()}\n",
        "\n",
        "    @property\n",
        "    def metrics(self):\n",
        "        # We need to list our metrics here so the `reset_states()` can be\n",
        "        # called automatically.\n",
        "        return [self.loss_tracker, self.acc_tracker]\n",
        "\n",
        "\n",
        "cnn_model = get_cnn_model()\n",
        "encoder = TransformerEncoderBlock(embed_dim=EMBED_DIM, dense_dim=FF_DIM, num_heads=1)\n",
        "decoder = TransformerDecoderBlock(embed_dim=EMBED_DIM, ff_dim=FF_DIM, num_heads=2)\n",
        "caption_model = ImageCaptioningModel(\n",
        "    cnn_model=cnn_model, encoder=encoder, decoder=decoder\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iUdMnYHGuJIv"
      },
      "source": [
        "## Model training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fRVnsolHuJIv"
      },
      "source": [
        "\n",
        "# Define the loss function\n",
        "cross_entropy = keras.losses.SparseCategoricalCrossentropy(\n",
        "    from_logits=False, reduction=\"none\"\n",
        ")\n",
        "\n",
        "# EarlyStopping criteria\n",
        "early_stopping = keras.callbacks.EarlyStopping(patience=3, restore_best_weights=True)\n",
        "\n",
        "\n",
        "# Learning Rate Scheduler for the optimizer\n",
        "class LRSchedule(keras.optimizers.schedules.LearningRateSchedule):\n",
        "    def __init__(self, post_warmup_learning_rate, warmup_steps):\n",
        "        super().__init__()\n",
        "        self.post_warmup_learning_rate = post_warmup_learning_rate\n",
        "        self.warmup_steps = warmup_steps\n",
        "\n",
        "    def __call__(self, step):\n",
        "        global_step = tf.cast(step, tf.float32)\n",
        "        warmup_steps = tf.cast(self.warmup_steps, tf.float32)\n",
        "        warmup_progress = global_step / warmup_steps\n",
        "        warmup_learning_rate = self.post_warmup_learning_rate * warmup_progress\n",
        "        return tf.cond(\n",
        "            global_step < warmup_steps,\n",
        "            lambda: warmup_learning_rate,\n",
        "            lambda: self.post_warmup_learning_rate,\n",
        "        )\n",
        "\n",
        "\n",
        "# Create a learning rate schedule\n",
        "num_train_steps = len(train_dataset) * EPOCHS\n",
        "num_warmup_steps = num_train_steps // 15\n",
        "lr_schedule = LRSchedule(post_warmup_learning_rate=1e-4, warmup_steps=num_warmup_steps)\n",
        "\n",
        "# Compile the model\n",
        "caption_model.compile(optimizer=keras.optimizers.Adam(lr_schedule), loss=cross_entropy)\n",
        "\n",
        "# Fit the model\n",
        "caption_model.fit(\n",
        "    train_dataset,\n",
        "    epochs=EPOCHS,\n",
        "    validation_data=valid_dataset,\n",
        "    callbacks=[early_stopping],\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PHrBOeQUuJIw"
      },
      "source": [
        "## Check sample predictions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VmJRuL1QuJIw"
      },
      "source": [
        "vocab = vectorization.get_vocabulary()\n",
        "index_lookup = dict(zip(range(len(vocab)), vocab))\n",
        "max_decoded_sentence_length = SEQ_LENGTH - 1\n",
        "valid_images = list(valid_data.keys())\n",
        "\n",
        "\n",
        "def generate_caption():\n",
        "    # Select a random image from the validation dataset\n",
        "    sample_img = np.random.choice(valid_images)\n",
        "\n",
        "    # Read the image from the disk\n",
        "    sample_img = read_valid_image(sample_img)\n",
        "    img = sample_img.numpy().clip(0, 255).astype(np.uint8)\n",
        "    plt.imshow(img)\n",
        "    plt.show()\n",
        "\n",
        "    # Pass the image to the CNN\n",
        "    img = tf.expand_dims(sample_img, 0)\n",
        "    img = caption_model.cnn_model(img)\n",
        "\n",
        "    # Pass the image features to the Transformer encoder\n",
        "    encoded_img = caption_model.encoder(img, training=False)\n",
        "\n",
        "    # Generate the caption using the Transformer decoder\n",
        "    decoded_caption = \"<start> \"\n",
        "    for i in range(max_decoded_sentence_length):\n",
        "        tokenized_caption = vectorization([decoded_caption])[:, :-1]\n",
        "        mask = tf.math.not_equal(tokenized_caption, 0)\n",
        "        predictions = caption_model.decoder(\n",
        "            tokenized_caption, encoded_img, training=False, mask=mask\n",
        "        )\n",
        "        sampled_token_index = np.argmax(predictions[0, i, :])\n",
        "        sampled_token = index_lookup[sampled_token_index]\n",
        "        if sampled_token == \" <end>\":\n",
        "            break\n",
        "        decoded_caption += \" \" + sampled_token\n",
        "\n",
        "    decoded_caption = decoded_caption.replace(\"<start> \", \"\")\n",
        "    decoded_caption = decoded_caption.replace(\" <end>\", \"\").strip()\n",
        "    print(\"Predicted Caption: \", decoded_caption)\n",
        "\n",
        "\n",
        "# Check predictions for a few samples\n",
        "generate_caption()\n",
        "generate_caption()\n",
        "generate_caption()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}