neural_structured_learning/examples/notebooks/graph_keras_cnn_flowers.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "24gYiJcWNlpA"
      },
      "source": [
        "##### Copyright 2021 Google LLC"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ioaprt5q5US7"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ItXfxkxvosLH"
      },
      "source": [
        "# Graph regularization for image classification using synthesized graphs\n",
        "\n",
        "*By [Sayak Paul](https://sayak.dev)*\n",
        "\n",
        "\u003cbr\u003e\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/neural-structured-learning/blob/master/neural_structured_learning/examples/notebooks/graph_keras_cnn_flowers.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/neural-structured-learning/blob/master/neural_structured_learning/examples/notebooks/graph_keras_cnn_flowers.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z3otbdCMmJiJ"
      },
      "source": [
        "## Overview"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Eg62Pmz3o83v"
      },
      "source": [
        "This notebook is the counterpart of [Graph regularization for sentiment \n",
        "classification using synthesized graphs](https://www.tensorflow.org/neural_structured_learning/tutorials/graph_keras_lstm_imdb) for image classification. In this notebook, we will build a flower classification model that can categorize images of flowers into discrete classes such as sunflowers, roses,\n",
        "tulips, etc. \n",
        "\n",
        "We will demonstrate the use of graph regularization in this notebook by building a graph from the given input. The general recipe for building a graph-regularized model using the Neural Structured Learning (NSL) framework when the input does not contain an explicit graph is as follows:\n",
        "\n",
        "1.  Create embeddings for each image sample in the input. This can be done using\n",
        "    pre-trained models such as [EfficientNet](https://arxiv.org/abs/1905.11946),\n",
        "    [Inception](https://arxiv.org/abs/1512.00567),\n",
        "    [BiT](https://arxiv.org/abs/1912.11370) etc.\n",
        "2.  Build a graph based on these embeddings by using a similarity metric such as\n",
        "    the 'L2' distance, 'cosine' distance, etc. Nodes in the graph correspond to\n",
        "    samples and edges in the graph correspond to similarity between pairs of\n",
        "    samples.\n",
        "3.  Generate training data from the above synthesized graph and sample features.\n",
        "    The resulting training data will contain neighbor features in addition to\n",
        "    the original node features.\n",
        "4.  Create a neural network as a base model using the Keras sequential,\n",
        "    functional, or subclass API.\n",
        "5.  Wrap the base model with the GraphRegularization wrapper class, which is\n",
        "    provided by the NSL framework, to create a new graph Keras model. This new\n",
        "    model will include a graph regularization loss as the regularization term in\n",
        "    its training objective.\n",
        "6.  Train and evaluate the graph Keras model.\n",
        "\n",
        "**Note**: We expect that it would take readers about 1 hour to go through this\n",
        "tutorial."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nDOFbB34KY1R"
      },
      "source": [
        "## Requirements\n",
        "\n",
        "1.  Install the Neural Structured Learning package.\n",
        "2.  Install tensorflow-hub."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uVnjPmOaQlnH"
      },
      "outputs": [],
      "source": [
        "!pip install --quiet neural-structured-learning\n",
        "!pip install --quiet tensorflow-hub"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ngUVB9hDu8CK"
      },
      "source": [
        "## Dependencies and imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OpLo-U08i3PX"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "import neural_structured_learning as nsl\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "# Resets notebook state\n",
        "tf.keras.backend.clear_session()\n",
        "\n",
        "tfds.disable_progress_bar()\n",
        "\n",
        "print(\"Version: \", tf.__version__)\n",
        "print(\"Eager mode: \", tf.executing_eagerly())\n",
        "print(\n",
        "  \"GPU is\",\n",
        "  \"available\" if tf.config.list_physical_devices(\"GPU\") else \"NOT AVAILABLE\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mlWeG07DvAEa"
      },
      "source": [
        "## Flowers dataset\n",
        "\n",
        "The flowers dataset contains a total of 3670 images of flowers categorized into 5 classes - \n",
        "\n",
        "* daisy\n",
        "* dandelion\n",
        "* roses\n",
        "* sunflowers\n",
        "* tulips\n",
        "\n",
        "The dataset is balanced i.e., each class contains roughly the same number of examples. In the cells below, we first load the dataset using [TensorFlow Datasets](https://www.tensorflow.org/datasets) and then visualize a couple of examples from the dataset. \n",
        "\n",
        "The dataset does not come pre-split into training, validation, and test splits. However, when downloading the dataset, we can specify a splitting ratio. Here we'll be using a split of 85:15 for the training and validation sets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ntzN4Am1i5-U"
      },
      "outputs": [],
      "source": [
        "train_ds, validation_ds = tfds.load(\n",
        "  \"tf_flowers\",\n",
        "  split=[\"train[:85%]\", \"train[85%:]\"],\n",
        "  as_supervised=True\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mz9S3Cc3Odc7"
      },
      "source": [
        "After downloading the dataset and splitting it, we can visualize a few samples from it. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K7ZHPYrBjjYl"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(10, 10))\n",
        "for i, (image, label) in enumerate(train_ds.take(9)):\n",
        "  ax = plt.subplot(3, 3, i + 1)\n",
        "  plt.imshow(image)\n",
        "  plt.title(int(label))\n",
        "  plt.axis(\"off\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XhliNWZjOl0o"
      },
      "source": [
        "In the next cell, we determine the number of samples present in the splits."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WsPVr-3WjrUi"
      },
      "outputs": [],
      "source": [
        "num_train_examples = tf.data.experimental.cardinality(train_ds)\n",
        "num_val_examples = tf.data.experimental.cardinality(validation_ds)\n",
        "print(f\"Total training examples: {num_train_examples}\")\n",
        "print(f\"Total validation examples: {num_val_examples}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rVmqL-zcWm8v"
      },
      "source": [
        "## Graph construction\n",
        "\n",
        "Graph construction involves creating embeddings for the images and then using\n",
        "a similarity function to compare the embeddings."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uUyHEa-3TB2X"
      },
      "source": [
        "### Create sample embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qCe9vOy7-Br9"
      },
      "source": [
        "We will use a pre-trained [DenseNet121](https://arxiv.org/abs/1608.06993) model to create embeddings in the\n",
        "`tf.train.Example` format for each sample in the input. We will store the\n",
        "resulting embeddings in the `TFRecord` format along with an additional feature\n",
        "that represents the ID of each sample. This is important and will allow us match\n",
        "sample embeddings with corresponding nodes in the graph later. We'll start this section with a utility function to build a feature extraction model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NKnvW9p6aDcD"
      },
      "source": [
        "One important detail to note here is that we are using [random projections](https://en.wikipedia.org/wiki/Random_projection) in order to reduce the dimensionality of the final vector coming out of the pre-trained model. The final embedding vector is 1024-dimensional. When the size of the dataset is large, such high-dimensional vectors can consume a lot of memory. Hence depending on your use-case, it might be a good idea to further reduce the dimensionality. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jicfZ0poW88t"
      },
      "outputs": [],
      "source": [
        "IMG_SIZE = 224 #@param [\"128\", \"224\"] {type:\"raw\"}\n",
        "PROJECTED_DIM = 128 #@param {type:\"slider\", min:128, max:1024, step:128}\n",
        "#@markdown `IMG_SIZE` of 224 denotes the 224 $\\times$ 224 resolution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dSeESYsuj97o"
      },
      "outputs": [],
      "source": [
        "def create_feature_extractor_model():\n",
        "  \"\"\"Creates a feature extractor model with DenseNet121.\"\"\"\n",
        "  inputs = tf.keras.layers.Input((IMG_SIZE, IMG_SIZE, 3))\n",
        "  \n",
        "  densenet_model = tf.keras.applications.DenseNet121(weights=\"imagenet\", \n",
        "    input_shape=(IMG_SIZE, IMG_SIZE, 3),\n",
        "    pooling=\"avg\", include_top=False\n",
        "  )\n",
        "  densenet_model.trainable = False\n",
        "  x = tf.keras.applications.densenet.preprocess_input(inputs)\n",
        "  outputs = densenet_model(x, training=False)\n",
        "\n",
        "  return tf.keras.Model(inputs, outputs, name=\"densenet_feature_extractor\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f_-XRns3maF5"
      },
      "outputs": [],
      "source": [
        "feature_extractor = create_feature_extractor_model()\n",
        "feature_extractor.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wNzWo6E-YKR4"
      },
      "source": [
        "We encourage you to try out other pre-trained models available in the [`tf.keras.applications`](https://www.tensorflow.org/api_docs/python/tf/keras/applications) module and also on [TensorFlow Hub](https://tfhub.dev/s?fine-tunable=yes\u0026module-type=image-feature-vector\u0026publisher=google,tensorflow\u0026tf-version=tf2). We'll now write a couple of utility functions to create the sample embeddings for graph construction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-RO8vAGQm8xc"
      },
      "outputs": [],
      "source": [
        "def resize(image, label):\n",
        "  \"\"\"Resizes the images to (IMG_SIZE x IMG_SIZE) size.\"\"\"\n",
        "  image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE))\n",
        "  return image, label\n",
        "\n",
        "# Resize all the images to uniform shape so that they can\n",
        "# be batched.\n",
        "train_ds = train_ds.map(resize)\n",
        "validation_ds = validation_ds.map(resize)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FAIhVW6InktB"
      },
      "outputs": [],
      "source": [
        "def _int64_feature(value):\n",
        "  \"\"\"Returns int64 tf.train.Feature.\"\"\"\n",
        "  return tf.train.Feature(int64_list=tf.train.Int64List(value=value.tolist()))\n",
        "\n",
        "\n",
        "def _bytes_feature(value):\n",
        "  \"\"\"Returns bytes tf.train.Feature.\"\"\"\n",
        "  return tf.train.Feature(\n",
        "    bytes_list=tf.train.BytesList(value=[value.encode('utf-8')]))\n",
        "\n",
        "\n",
        "def _float_feature(value):\n",
        "  \"\"\"Returns float tf.train.Feature.\"\"\"\n",
        "  return tf.train.Feature(float_list=tf.train.FloatList(value=value.tolist()))\n",
        "\n",
        "\n",
        "def create_embedding_example(feature_extractor, image,\n",
        "                             projection_matrix, record_id):\n",
        "  \"\"\"Create tf.Example containing the sample's embedding and its ID.\"\"\"\n",
        "\n",
        "  image_features = feature_extractor(image[None, ...])\n",
        "  image_features_numpy = image_features.numpy().squeeze()\n",
        "  compressed_image_features = image_features_numpy.dot(projection_matrix)\n",
        "\n",
        "  features = {\n",
        "    \"id\": _bytes_feature(str(record_id)),\n",
        "    \"embedding\": _float_feature(compressed_image_features)\n",
        "  }\n",
        "  return tf.train.Example(features=tf.train.Features(feature=features))\n",
        "\n",
        "\n",
        "def generate_random_projection_weights(original_dim=1024,\n",
        "                                       projected_dim=PROJECTED_DIM):\n",
        "  \"\"\"Generates a random projection matrix.\"\"\"\n",
        "  random_projection_matrix = np.random.randn(\n",
        "    projected_dim, original_dim).T\n",
        "  return random_projection_matrix\n",
        "\n",
        "\n",
        "def create_embeddings(feature_extractor, dataset, output_path,\n",
        "                      starting_record_id):\n",
        "  \"\"\"Creates TFRecords with embeddings of the images.\"\"\"\n",
        "  projection_matrix = generate_random_projection_weights()\n",
        "  record_id = int(starting_record_id)\n",
        "  with tf.io.TFRecordWriter(output_path) as writer:\n",
        "    for image, _ in dataset:\n",
        "      example = create_embedding_example(feature_extractor,\n",
        "                                         image,\n",
        "                                         projection_matrix,\n",
        "                                         record_id)\n",
        "      record_id = record_id + 1\n",
        "      writer.write(example.SerializeToString())\n",
        "  return record_id\n",
        "\n",
        "\n",
        "# Persist TF.Example features containing embeddings for training data in\n",
        "# TFRecord format.\n",
        "create_embeddings(feature_extractor, train_ds, \"flowers_embeddings.tfr\", 0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C8s06RuI_vKs"
      },
      "source": [
        "### Graph building\n",
        "\n",
        "Now that we have the sample embeddings, we will use them to build a similarity\n",
        "graph, i.e, nodes in this graph will correspond to samples and edges in this\n",
        "graph will correspond to similarity between pairs of nodes.\n",
        "\n",
        "Neural Structured Learning provides a graph building library to build a graph\n",
        "based on sample embeddings. It uses\n",
        "[**cosine similarity**](https://en.wikipedia.org/wiki/Cosine_similarity) as the\n",
        "similarity measure to compare embeddings and build edges between them. It also\n",
        "allows us to specify a similarity threshold, which can be used to discard\n",
        "dissimilar edges from the final graph. In this example, using 0.7 as the\n",
        "similarity threshold and 12345 as the random seed, we end up with a graph that\n",
        "has 987,246 bi-directional edges. Here we're using the graph builder's support\n",
        "for [locality-sensitive hashing](https://en.wikipedia.org/wiki/Locality-sensitive_hashing)\n",
        "(LSH) to speed up graph building. For details on using the graph builder's LSH\n",
        "support, see the\n",
        "[`build_graph_from_config`](https://www.tensorflow.org/neural_structured_learning/api_docs/python/nsl/tools/build_graph_from_config)\n",
        "API documentation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2Vv5doOUuZFz"
      },
      "outputs": [],
      "source": [
        "similarity_threshold = 0.7 \n",
        "graph_builder_config = nsl.configs.GraphBuilderConfig(\n",
        "    similarity_threshold=similarity_threshold, \n",
        "    lsh_splits=10, lsh_rounds=15, random_seed=12345)\n",
        "nsl.tools.build_graph_from_config([\"flowers_embeddings.tfr\"],\n",
        "                                  \"flowers_graph_70.tsv\",\n",
        "                                  graph_builder_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4dk9xfQcK553"
      },
      "source": [
        "Each bi-directional edge is represented by two directed edges in the output TSV\n",
        "file, so that file contains 493,623 * 2 = 98,246 total lines:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ob9GwDWOy4vt"
      },
      "outputs": [],
      "source": [
        "!wc -l flowers_graph_70.tsv"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_USkfut69gNW"
      },
      "source": [
        "## Sample features\n",
        "\n",
        "We create sample features for our problem using the `tf.train.Example` format\n",
        "and persist them in the `TFRecord` format. Each sample will include the\n",
        "following three features:\n",
        "\n",
        "1.  **id**: The node ID of the sample.\n",
        "2.  **image**: A byte list containing raw image vectors.\n",
        "3.  **label**: A singleton int64 identifying the target class of the image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f-RTbN6y04qv"
      },
      "outputs": [],
      "source": [
        "def _bytes_feature_image(value):\n",
        "    \"\"\"Returns bytes tf.train.Feature.\"\"\"\n",
        "    return tf.train.Feature(\n",
        "        bytes_list=tf.train.BytesList(value=[value]))\n",
        "    \n",
        "\n",
        "def create_example(image, label, record_id):\n",
        "    \"\"\"Create tf.Example containing the image, label, and ID.\"\"\"\n",
        "    features = {\n",
        "        \"id\": _bytes_feature(str(record_id)),\n",
        "        \"image\": _bytes_feature_image(image.numpy()),\n",
        "        \"label\": _int64_feature(np.asarray([label])),\n",
        "    }\n",
        "    return tf.train.Example(features=tf.train.Features(feature=features))\n",
        "\n",
        "\n",
        "def create_records(dataset, record_path, starting_record_id):\n",
        "  \"\"\"Generates TFRecords from a tf.data.Dataset object.\"\"\"\n",
        "  record_id = int(starting_record_id)\n",
        "  with tf.io.TFRecordWriter(record_path) as writer:\n",
        "    for image, label in dataset:\n",
        "      image = tf.cast(image, tf.uint8)\n",
        "      image = tf.image.encode_jpeg(image, optimize_size=True,\n",
        "                                    chroma_downsampling=False)\n",
        "      \n",
        "      example = create_example(image, label, record_id)\n",
        "      record_id = record_id + 1\n",
        "      writer.write(example.SerializeToString())\n",
        "  return record_id\n",
        "\n",
        "# Persist TF.Example features (images and labels) for training and validation\n",
        "# data in TFRecord format.\n",
        "next_record_id = create_records(train_ds,\n",
        "                                \"train_data.tfr\", 0)\n",
        "create_records(validation_ds, \"validation_data.tfr\",\n",
        "               next_record_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T94mB0beWKoa"
      },
      "source": [
        "**A note on `create_records()`**:\n",
        "\n",
        "Images are serialized as byte-strings in TFRecords. `tf.image.encode_jpeg()` function allows us to do that in an optimal way (when `optimize_size` is set to `True`). Inputs to that function should be integers. This is why, we first cast the image to `tf.uint8` and then pass it to `tf.image.encode_jpeg()`. To know more, refer to [this tutorial](https://www.tensorflow.org/tutorials/load_data/tfrecord#walkthrough_reading_and_writing_image_data)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0DOaTUQVvNQ5"
      },
      "source": [
        "## Augment training data with graph neighbors\n",
        "\n",
        "Since we have the sample features and the synthesized graph, we can generate the\n",
        "augmented training data for Neural Structured Learning. The NSL framework\n",
        "provides a library to combine the graph and the sample features to produce\n",
        "the final training data for graph regularization. The resulting training data\n",
        "will include original sample features as well as features of their corresponding\n",
        "neighbors.\n",
        "\n",
        "In this tutorial, we consider undirected edges and use a maximum of 3 neighbors\n",
        "per sample to augment training data with graph neighbors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mK5Bc5QS6A-Z"
      },
      "outputs": [],
      "source": [
        "nsl.tools.pack_nbrs(\n",
        "    \"train_data.tfr\",\n",
        "    \"\",\n",
        "    \"flowers_graph_70.tsv\",\n",
        "    \"nsl_train_data.tfr\",\n",
        "    add_undirected_edges=True,\n",
        "    max_nbrs=3\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dZeyYg1Rvg8W"
      },
      "outputs": [],
      "source": [
        "!ls -lh *_data.tfr"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fO36xEgHvRzk"
      },
      "source": [
        "## Base model\n",
        "\n",
        "We'll first build a model without graph regularization. We'll use a simple convolutional neural network (CNN) for the purpose of this tutorial. However, this can be easily replaced with more sophisticated network architectures. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kLSbRFguBUNl"
      },
      "source": [
        "### Global variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h6PEUJ256XLG"
      },
      "outputs": [],
      "source": [
        "NBR_FEATURE_PREFIX = \"NL_nbr_\"\n",
        "NBR_WEIGHT_SUFFIX = \"_weight\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s8gMVBw6t6CI"
      },
      "source": [
        "### Hyperparameters\n",
        "\n",
        "We will use an instance of `HParams` to inclue various hyperparameters and\n",
        "constants used for training and evaluation. We briefly describe each of them\n",
        "below:\n",
        "\n",
        "-   **num_classes**: There are 5 classes.\n",
        "\n",
        "-   **num_train_examples** The number of training examples. \n",
        "\n",
        "-   **num_val_examples**:  The number of validation examples.\n",
        "\n",
        "-   **distance_type**: This is the distance metric used to regularize the sample\n",
        "    with its neighbors.\n",
        "\n",
        "-   **graph_regularization_multiplier**: This controls the relative weight of\n",
        "    the graph regularization term in the overall loss function.\n",
        "\n",
        "-   **num_neighbors**: The number of neighbors used for graph regularization.\n",
        "    This value has to be less than or equal to the `max_nbrs` argument used\n",
        "    above when invoking `nsl.tools.pack_nbrs`.\n",
        "\n",
        "-   **num_channels**: The number of channels to be used in the convolutional layers.\n",
        "\n",
        "-   **kernel_size**: Kernel size to be used in the convolutional layers.  \n",
        "\n",
        "-   **train_epochs**: The number of training epochs.\n",
        "\n",
        "-   **batch_size**: Batch size used for training and evaluation.\n",
        "\n",
        "-   **eval_steps**: The number of batches to process before deeming evaluation\n",
        "    is complete. If set to `None`, all instances in the test set are evaluated."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j40LH7b26cQE"
      },
      "outputs": [],
      "source": [
        "class HParams(object):\n",
        "  \"\"\"Hyperparameters used for training.\"\"\"\n",
        "  def __init__(self):\n",
        "    ### dataset parameters\n",
        "    self.num_classes = 5\n",
        "    self.num_train_examples = num_train_examples\n",
        "    self.num_val_examples = num_val_examples\n",
        "    \n",
        "    ### neural graph learning parameters\n",
        "    self.distance_type = nsl.configs.DistanceType.L2\n",
        "    self.graph_regularization_multiplier = 0.3\n",
        "    self.num_neighbors = 2\n",
        "\n",
        "    ### network architecture parameters\n",
        "    self.num_channels = 32\n",
        "    self.kernel_size = 3\n",
        "    \n",
        "    ### training parameters\n",
        "    self.train_epochs = 30 \n",
        "    self.batch_size = 64 \n",
        "    \n",
        "    ### eval parameters\n",
        "    self.eval_steps = None  # All instances in the test set are evaluated.\n",
        "\n",
        "HPARAMS = HParams()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PCypw0H1fr7n"
      },
      "source": [
        "### Prepare the data\n",
        "\n",
        "Now, we will prepare our dataset with TensorFlow's `data` module ([`tf.data`](https://www.tensorflow.org/api_docs/python/tf/data)). This will include parsing the TFRecords, structuring them, batching them, and shuffling them if necessary.\n",
        "\n",
        "In the next two cells, we first define a default value for the image examples which will come in handy when parsing the neighbor examples. Then we write a utility function to create the `tf.data.Dataset` objects which will be fed to our model in a moment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gY2n7_4zx9by"
      },
      "outputs": [],
      "source": [
        "default_jpeg_value = tf.ones((IMG_SIZE, IMG_SIZE, 3), dtype=tf.uint8)\n",
        "default_jpeg_value *= 255\n",
        "default_jpeg_value = tf.image.encode_jpeg(default_jpeg_value, optimize_size=True,\n",
        "                                         chroma_downsampling=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w5u6_zuU6wLr"
      },
      "outputs": [],
      "source": [
        "def make_dataset(file_path, training=False):\n",
        "  \"\"\"Creates a `tf.data.TFRecordDataset`.\n",
        "\n",
        "  Args:\n",
        "    file_path: Name of the file in the `.tfrecord` format containing\n",
        "      `tf.train.Example` objects.\n",
        "    training: Boolean indicating if we are in training mode.\n",
        "\n",
        "  Returns:\n",
        "    An instance of `tf.data.TFRecordDataset` containing the `tf.train.Example`\n",
        "    objects.\n",
        "  \"\"\"\n",
        "\n",
        "  def parse_example(example_proto):\n",
        "    \"\"\"Extracts relevant fields from the `example_proto`.\n",
        "\n",
        "    Args:\n",
        "      example_proto: An instance of `tf.train.Example`.\n",
        "\n",
        "    Returns:\n",
        "      A pair whose first value is a dictionary containing relevant features\n",
        "      and whose second value contains the ground truth labels.\n",
        "    \"\"\"\n",
        "    feature_spec = {\n",
        "      'image': tf.io.FixedLenFeature([], tf.string, \n",
        "                                      default_value=default_jpeg_value),\n",
        "      'label': tf.io.FixedLenFeature((), tf.int64, default_value=-1),\n",
        "    }\n",
        "    # We also extract corresponding neighbor features in a similar manner to\n",
        "    # the features above during training.\n",
        "    if training:\n",
        "      for i in range(HPARAMS.num_neighbors):\n",
        "        nbr_feature_key = '{}{}_{}'.format(NBR_FEATURE_PREFIX, i, 'image')\n",
        "        nbr_weight_key = '{}{}{}'.format(NBR_FEATURE_PREFIX, i,\n",
        "                                         NBR_WEIGHT_SUFFIX)\n",
        "        feature_spec[nbr_feature_key] = tf.io.FixedLenFeature([], tf.string,\n",
        "                                            default_value=default_jpeg_value)\n",
        "\n",
        "        # We assign a default value of 0.0 for the neighbor weight so that\n",
        "        # graph regularization is done on samples based on their exact number\n",
        "        # of neighbors. In other words, non-existent neighbors are discounted.\n",
        "        feature_spec[nbr_weight_key] = tf.io.FixedLenFeature(\n",
        "            [1], tf.float32, default_value=tf.constant([0.0]))\n",
        "    \n",
        "    features = tf.io.parse_single_example(example_proto, feature_spec)\n",
        "    labels = features.pop('label')\n",
        "    \n",
        "    # We need to convert the byte-strings back to images.\n",
        "    features['image'] = tf.image.decode_jpeg(features['image'], channels=3)\n",
        "    if training:\n",
        "      for i in range(HPARAMS.num_neighbors):\n",
        "        nbr_feature_key = '{}{}_{}'.format(NBR_FEATURE_PREFIX, i, 'image')\n",
        "        features[nbr_feature_key] = tf.image.decode_jpeg(features[nbr_feature_key],\n",
        "                                                         channels=3)\n",
        "\n",
        "    return features, labels\n",
        "\n",
        "  dataset = tf.data.TFRecordDataset([file_path])\n",
        "  if training:\n",
        "    dataset = dataset.shuffle(HPARAMS.batch_size * 10)\n",
        "  dataset = dataset.map(parse_example)\n",
        "  dataset = dataset.batch(HPARAMS.batch_size)\n",
        "  return dataset\n",
        "\n",
        "\n",
        "train_dataset = make_dataset('nsl_train_data.tfr', True)\n",
        "validation_dataset = make_dataset('validation_data.tfr')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BBB-TY-F80NJ"
      },
      "source": [
        "### Visualization\n",
        "\n",
        "In this section, we visualize the neighbor images as computed by Neural Structured Learning during building the graph. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2QeTI1q98YOR"
      },
      "outputs": [],
      "source": [
        "sample = next(iter(train_dataset))\n",
        "sample[0].keys()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k2M-OUoH83KQ"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(20, 20))\n",
        "\n",
        "plt.subplot(1, 3, 1)\n",
        "plt.imshow(sample[0][\"NL_nbr_0_image\"][0])\n",
        "neighbor_one_weight = float(sample[0][\"NL_nbr_0_weight\"][0].numpy())\n",
        "plt.title(f\"Neighbor 1 with weight: {neighbor_one_weight:.3f}\", fontsize=14)\n",
        "plt.axis(\"off\")\n",
        "\n",
        "plt.subplot(1, 3, 2)\n",
        "plt.imshow(sample[0][\"NL_nbr_1_image\"][0])\n",
        "neighbor_two_weight = float(sample[0][\"NL_nbr_1_weight\"][0].numpy())\n",
        "plt.title(f\"Neighbor 2 with weight: {neighbor_two_weight:.3f}\", fontsize=14)\n",
        "plt.axis(\"off\")\n",
        "\n",
        "plt.subplot(1, 3, 3)\n",
        "plt.imshow(sample[0][\"image\"][0])\n",
        "plt.title(f\"Original image with label: {int(sample[1][0])}\", fontsize=14)\n",
        "plt.axis(\"off\")\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cWhY84H2lfh_"
      },
      "source": [
        "In the figure above, `weight` denotes the similarity strength of the examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QDr206Vb8xFt"
      },
      "source": [
        "### Model training (no graph regularization)\n",
        "\n",
        "With our datasets prepared, we are now ready to train our shallow CNN model *without* graph regularization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z4mGjVLw3jNc"
      },
      "outputs": [],
      "source": [
        "def make_cnn_model():\n",
        "  \"\"\"Creates a simple CNN model.\"\"\"\n",
        "  model = tf.keras.Sequential([\n",
        "      tf.keras.layers.InputLayer(\n",
        "          input_shape=(IMG_SIZE, IMG_SIZE, 3), name='image'),\n",
        "      tf.keras.layers.experimental.preprocessing.Rescaling(scale=1. / 255),\n",
        "      tf.keras.layers.Conv2D(\n",
        "          HPARAMS.num_channels, HPARAMS.kernel_size, activation='relu'),\n",
        "      tf.keras.layers.MaxPooling2D(),\n",
        "      tf.keras.layers.Conv2D(\n",
        "          HPARAMS.num_channels, HPARAMS.kernel_size, activation='relu'),\n",
        "      tf.keras.layers.MaxPooling2D(),\n",
        "      tf.keras.layers.GlobalAvgPool2D(),\n",
        "      tf.keras.layers.Dense(HPARAMS.num_classes)\n",
        "  ])\n",
        "  return model\n",
        "\n",
        "\n",
        "model = make_cnn_model()\n",
        "model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvkUsggciw1I"
      },
      "source": [
        "After building and initializing the model in Keras, we can compile it and finally train it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WblCOGZC5kGX"
      },
      "outputs": [],
      "source": [
        "model.compile(\n",
        "    optimizer='adam',\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    metrics=['accuracy'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gSUYa1UL5rA4"
      },
      "outputs": [],
      "source": [
        "history = model.fit(\n",
        "  train_dataset,\n",
        "  validation_data=validation_dataset,\n",
        "  epochs=HPARAMS.train_epochs,\n",
        "  callbacks=[tf.keras.callbacks.EarlyStopping(patience=5)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1zjhxKn4iU-M"
      },
      "source": [
        "### Plot training metrics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E9CXstqo5zjW"
      },
      "outputs": [],
      "source": [
        "history_dict = history.history\n",
        "\n",
        "acc = history_dict['accuracy']\n",
        "val_acc = history_dict['val_accuracy']\n",
        "loss = history_dict['loss']\n",
        "val_loss = history_dict['val_loss']\n",
        "\n",
        "epochs = range(1, len(acc) + 1)\n",
        "\n",
        "# \"-r^\" is for solid red line with triangle markers.\n",
        "plt.plot(epochs, loss, '-r^', label='Training loss')\n",
        "# \"-b0\" is for solid blue line with circle markers.\n",
        "plt.plot(epochs, val_loss, '-bo', label='Validation loss')\n",
        "plt.title('Training and validation loss')\n",
        "plt.xlabel('Epochs')\n",
        "plt.ylabel('Loss')\n",
        "plt.legend(loc='best')\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8zl_bCYy68lT"
      },
      "outputs": [],
      "source": [
        "plt.clf()   # clear figure\n",
        "\n",
        "plt.plot(epochs, acc, '-r^', label='Training acc')\n",
        "plt.plot(epochs, val_acc, '-bo', label='Validation acc')\n",
        "plt.title('Training and validation accuracy')\n",
        "plt.xlabel('Epochs')\n",
        "plt.ylabel('Accuracy')\n",
        "plt.legend(loc='best')\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SymtYWWiMUum"
      },
      "source": [
        "## Graph regularization\n",
        "\n",
        "We are now ready to try graph regularization using the base model that we built\n",
        "above. We will use the `GraphRegularization` wrapper class provided by the\n",
        "Neural Structured Learning framework to wrap the base (CNN) model to include\n",
        "graph regularization. The rest of the steps for training and evaluating the\n",
        "graph-regularized model are similar to that of the base model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pCIkVe_QFX38"
      },
      "source": [
        "### Create graph-regularized model\n",
        "\n",
        "To assess the incremental benefit of graph regularization, we will create a new\n",
        "base model instance. This is because `model` has already been trained for a few\n",
        "iterations, and reusing this trained model to create a graph-regularized model\n",
        "will not be a fair comparison for `model`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uZImBkhJ7AFP"
      },
      "outputs": [],
      "source": [
        "# Build a new base CNN model.\n",
        "base_reg_model = make_cnn_model()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hdeDq6H77FWI"
      },
      "outputs": [],
      "source": [
        "# Wrap the base model with graph regularization.\n",
        "graph_reg_config = nsl.configs.make_graph_reg_config(\n",
        "  max_neighbors=HPARAMS.num_neighbors,\n",
        "  multiplier=HPARAMS.graph_regularization_multiplier,\n",
        "  distance_type=HPARAMS.distance_type,\n",
        "  sum_over_axis=-1)\n",
        "graph_reg_model = nsl.keras.GraphRegularization(base_reg_model,\n",
        "                                                graph_reg_config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RKQ3RXZPiq0m"
      },
      "outputs": [],
      "source": [
        "graph_reg_model.compile(\n",
        "    optimizer='adam',\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    metrics=['accuracy'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RLRxxV4G7McQ"
      },
      "outputs": [],
      "source": [
        "graph_reg_history = graph_reg_model.fit(\n",
        "  train_dataset,\n",
        "  validation_data=validation_dataset,\n",
        "  epochs=HPARAMS.train_epochs,\n",
        "  callbacks=[tf.keras.callbacks.EarlyStopping(patience=5)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QyLIa0qEi74A"
      },
      "source": [
        "### Plot training metrics of the graph-regularized model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "84ffIFmd7PMz"
      },
      "outputs": [],
      "source": [
        "graph_reg_history_dict = graph_reg_history.history\n",
        "\n",
        "acc = graph_reg_history_dict['accuracy']\n",
        "val_acc = graph_reg_history_dict['val_accuracy']\n",
        "loss = graph_reg_history_dict['loss']\n",
        "graph_loss = graph_reg_history_dict['scaled_graph_loss']\n",
        "val_loss = graph_reg_history_dict['val_loss']\n",
        "\n",
        "epochs = range(1, len(acc) + 1)\n",
        "\n",
        "plt.clf()   # clear figure\n",
        "\n",
        "# \"-r^\" is for solid red line with triangle markers.\n",
        "plt.plot(epochs, loss, '-r^', label='Training loss')\n",
        "# \"-gD\" is for solid green line with diamond markers.\n",
        "plt.plot(epochs, graph_loss, '-gD', label='Training graph loss')\n",
        "# \"-b0\" is for solid blue line with circle markers.\n",
        "plt.plot(epochs, val_loss, '-bo', label='Validation loss')\n",
        "plt.title('Training and validation loss')\n",
        "plt.xlabel('Epochs')\n",
        "plt.ylabel('Loss')\n",
        "plt.legend(loc='best')\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5EJUJJja7lEB"
      },
      "outputs": [],
      "source": [
        "plt.clf()   # clear figure\n",
        "\n",
        "plt.plot(epochs, acc, '-r^', label='Training acc')\n",
        "plt.plot(epochs, val_acc, '-bo', label='Validation acc')\n",
        "plt.title('Training and validation accuracy')\n",
        "plt.xlabel('Epochs')\n",
        "plt.ylabel('Accuracy')\n",
        "plt.legend(loc='best')\n",
        "\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cShQNXGNjIOs"
      },
      "source": [
        "Based on the performances of the two models (with and without graph regularization) we observe that the model performs better with graph regularization. Here is ablation study with respect to various values of embedding dimensionality:\n",
        "\n",
        "| Embedding \u003cbr\u003eDim \t| With \u003cbr\u003eGraph-reg \t| Without\u003cbr\u003eGraph-reg \t|\n",
        "|:---:\t|:---:\t|:---:\t|\n",
        "| 128 \t| 60.55% \t| 56.00% \t|\n",
        "| 256 \t| 57.82% \t| 55.45% \t|\n",
        "| 384 \t| 58% \t| 60.18% \t|\n",
        "| 512 \t| 57.45% \t| 57.64% \t|\n",
        "\n",
        "The table above reports the validation top-1 accuracies with and without graph regularization under different reduced embedding dimensionalities. We encourage you to further experiment with hyperparameters like `similarity_threshold` and `graph_regularization_multiplier` to study how these hyperparameters affect the final performance. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8X4zCEyPhIp-"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "We have demonstrated the use of graph regularization using the Neural Structured\n",
        "Learning (NSL) framework even when the input does not contain an explicit graph.\n",
        "We considered the task of flower species' classification of flower images for\n",
        "which we synthesized a similarity graph based on review embeddings. We encourage\n",
        "users to experiment further by varying hyperparameters, the amount of\n",
        "supervision, and by using different model architectures."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AV6RysJlm-D0"
      },
      "source": [
        "## Acknowledgements\n",
        "\n",
        "Thanks to [Arjun Gopalan](https://research.google/people/ArjunGopalan/) from Google. Without his guidance, this tutorial would not have been possible."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "name": "graph_keras_cnn_flowers",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}